mengw15 commented on code in PR #5260: URL: https://github.com/apache/texera/pull/5260#discussion_r3454681721
########## frontend/src/app/workspace/service/notebook-migration/migration-llm.ts: ########## @@ -0,0 +1,303 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +import { Injectable } from "@angular/core"; +import { GuiConfigService } from "../../../common/service/gui-config.service"; +import { createOpenAI } from "@ai-sdk/openai"; +import { generateText, type ModelMessage } from "ai"; +import { AppSettings } from "../../../common/app-setting"; +import { v4 as uuidv4 } from "uuid"; +import { WorkflowUtilService } from "../workflow-graph/util/workflow-util.service"; +import { OperatorPredicate } from "../../types/workflow-common.interface"; +import { + TEXERA_OVERVIEW, + TUPLE_DOCUMENTATION, + TABLE_DOCUMENTATION, + OPERATOR_DOCUMENTATION, + UDF_INPUT_PORT_DOCUMENTATION, + EXAMPLE_OF_GOOD_CONVERSION, + VISUALIZER_DOCUMENTATION, + EXAMPLE_OF_MULTIPLE_UDF_CONVERSION, + WORKFLOW_PROMPT, + MAPPING_PROMPT, +} from "./migration-prompts"; + +interface Cell { + cell_type: string; + metadata: { [key: string]: any }; + source: string; +} + +export interface Notebook { + cells: Cell[]; +} + +interface WorkflowJSON { + operators: OperatorPredicate[]; + operatorPositions: Record<string, { x: number; y: number }>; + links: any[]; + commentBoxes: any[]; + settings: { + dataTransferBatchSize: number; + }; +} + +interface CombinedMapping { + operator_to_cell: Record<string, string[]>; + cell_to_operator: Record<string, string[]>; +} + +@Injectable() +export class NotebookMigrationLLM { + private model: any; + private messages: ModelMessage[] = []; + private initialized = false; + + private static readonly DOCUMENTATION: string[] = [ + TEXERA_OVERVIEW, + TUPLE_DOCUMENTATION, + TABLE_DOCUMENTATION, + OPERATOR_DOCUMENTATION, + EXAMPLE_OF_GOOD_CONVERSION, + VISUALIZER_DOCUMENTATION, + UDF_INPUT_PORT_DOCUMENTATION, + EXAMPLE_OF_MULTIPLE_UDF_CONVERSION, + ]; + + constructor( + private config: GuiConfigService, + private workflowUtilService: WorkflowUtilService + ) {} + + private get enabled(): boolean { + return this.config.env.pythonNotebookMigrationEnabled; + } + + private assertEnabled(): void { + if (!this.enabled) { + throw new Error("Notebook migration feature is disabled"); + } + } + + private parseJsonResponse(raw: string, context: string): any { + // Trim first, then strip optional markdown code fences (```json ... ``` or ``` ... ```) + const cleaned = raw + .trim() + .replace(/^```[a-zA-Z]*\s*/, "") + .replace(/\s*```$/, "") + .trim(); + try { + return JSON.parse(cleaned); + } catch (err) { + throw new Error(`Failed to parse LLM ${context} response as JSON: ${(err as Error).message}`); + } + } + + /** + * Initialize a new LLM session with Texera documentation + */ + public initialize(modelType: string = "gpt-5-mini", apiKey: string = "dummy"): void { + this.assertEnabled(); + this.model = createOpenAI({ + baseURL: new URL(`${AppSettings.getApiEndpoint()}`, document.baseURI).toString(), + // apiKey is required by the library for creating the OpenAI compatible client; + // For security reason, we store the apiKey at the backend, thus the value is dummy here. + apiKey: apiKey, + }).chat(modelType); + + this.messages = [ + ...NotebookMigrationLLM.DOCUMENTATION.map( + (doc): ModelMessage => ({ + role: "system", + content: doc, + }) + ), + ]; + + this.initialized = true; + } + + /** + * Verify the connection to the LLM using the given API key + */ + public async verifyConnection(): Promise<boolean> { + if (!this.enabled) return false; + if (!this.initialized) { + throw new Error("LLM session not initialized"); + } + + try { + await generateText({ + model: this.model, + messages: [ + { + role: "user", + content: "ping", + }, + ], + maxOutputTokens: 10, + }); + + return true; + } catch (err) { + console.error("API key verification failed:", err); + return false; + } + } + + /** + * Send a prompt and receive a response. + * All prior documentation and conversation is preserved. + */ + private async sendPrompt(prompt: string): Promise<string> { + if (!this.initialized) { + throw new Error("LLM session not initialized"); + } + + this.messages.push({ + role: "user", + content: prompt, + }); + + const result = await generateText({ + model: this.model, + messages: this.messages, + }); + + this.messages.push({ + role: "assistant", + content: result.text, + }); + + return result.text; + } + + /** + * Send a Jupyter Notebook to be converted into a workflow and mapping. + */ + public async convertNotebookToWorkflow(notebook: Notebook): Promise<string> { + this.assertEnabled(); + if (!this.initialized) { + throw new Error("LLM session not initialized"); + } + + const codeCells = notebook.cells.filter(cell => cell.cell_type === "code"); + const notebookString = codeCells + .map(cell => { + const uuid = String(cell.metadata.uuid); + return `# START ${uuid}\n${cell.source}\n# END ${uuid}`; + }) + .join("\n\n"); + + const workflow = await this.sendPrompt(`${WORKFLOW_PROMPT}\n${notebookString}`); + const mapping = await this.sendPrompt(MAPPING_PROMPT); + + // Remove ```json blocks and parse + const udfLLMResponse = this.parseJsonResponse(workflow, "workflow"); + + const workflowJSON: WorkflowJSON = { + operators: [], + operatorPositions: {}, + links: [], + commentBoxes: [], + settings: { + dataTransferBatchSize: 400, + }, + }; + + const udfMappingToUUID: Record<string, string> = {}; + + Object.entries(udfLLMResponse.code).forEach(([udfId, udfCode], i) => { + let udfOutputColumns: { attributeName: string; attributeType: string }[] = []; + if (udfLLMResponse.outputs && udfLLMResponse.outputs[udfId]) { + udfOutputColumns = udfLLMResponse.outputs[udfId].map((attr: string) => ({ + attributeName: attr, + attributeType: "binary", + })); + } Review Comment: Follow-up to C2 (binary `attributeType`): thanks for explaining that binary makes sense for UDF-to-UDF pickled-object data flow. That covers internal flow, but I'd want to also flag the user-facing implication: - **Result panel viewing:** the final UDF in any LLM-generated chain emits binary columns, so Texera's result panel renders them as opaque blobs (`<binary data>` / base64) rather than the dataframe / values the user expects to see after running their migrated notebook. The user has to manually edit every output column's `attributeType` before they can view results. - **Non-UDF downstream operators:** if the user (post-generation) adds a Visualize / Filter / Aggregate / SQL operator after the chain, those operators can't interpret binary columns either. For an MVP this is a reasonable simplification (LLM prompt stays simple), but the friction shows up every time the user runs a generated workflow. Two cheap mitigations short of "make the LLM infer types" (which I agree is bigger work): 1. **Terminal-UDF default:** treat the operator(s) with no outgoing edge as the "result-facing" one and default its outputs to `string` (or some non-binary type) instead of `binary` — at least the final result renders. 2. **Documented limitation in the class doc-comment + a frontend toast/modal after generation** noting "Output columns default to binary; edit per-operator to view typed results." Either way is fine; even just (2) would close the "user confused why results don't render" UX gap. Up to you whether to scope this here or treat it as a follow-up issue. ########## frontend/src/app/workspace/service/notebook-migration/migration-llm.ts: ########## @@ -0,0 +1,303 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +import { Injectable } from "@angular/core"; +import { GuiConfigService } from "../../../common/service/gui-config.service"; +import { createOpenAI } from "@ai-sdk/openai"; +import { generateText, type ModelMessage } from "ai"; +import { AppSettings } from "../../../common/app-setting"; +import { v4 as uuidv4 } from "uuid"; +import { WorkflowUtilService } from "../workflow-graph/util/workflow-util.service"; +import { OperatorPredicate } from "../../types/workflow-common.interface"; +import { + TEXERA_OVERVIEW, + TUPLE_DOCUMENTATION, + TABLE_DOCUMENTATION, + OPERATOR_DOCUMENTATION, + UDF_INPUT_PORT_DOCUMENTATION, + EXAMPLE_OF_GOOD_CONVERSION, + VISUALIZER_DOCUMENTATION, + EXAMPLE_OF_MULTIPLE_UDF_CONVERSION, + WORKFLOW_PROMPT, + MAPPING_PROMPT, +} from "./migration-prompts"; + +interface Cell { + cell_type: string; + metadata: { [key: string]: any }; + source: string; +} + +export interface Notebook { + cells: Cell[]; +} + +interface WorkflowJSON { + operators: OperatorPredicate[]; + operatorPositions: Record<string, { x: number; y: number }>; + links: any[]; + commentBoxes: any[]; + settings: { + dataTransferBatchSize: number; + }; +} + +interface CombinedMapping { + operator_to_cell: Record<string, string[]>; + cell_to_operator: Record<string, string[]>; +} + +@Injectable() +export class NotebookMigrationLLM { + private model: any; + private messages: ModelMessage[] = []; + private initialized = false; + + private static readonly DOCUMENTATION: string[] = [ + TEXERA_OVERVIEW, + TUPLE_DOCUMENTATION, + TABLE_DOCUMENTATION, + OPERATOR_DOCUMENTATION, + EXAMPLE_OF_GOOD_CONVERSION, + VISUALIZER_DOCUMENTATION, + UDF_INPUT_PORT_DOCUMENTATION, + EXAMPLE_OF_MULTIPLE_UDF_CONVERSION, + ]; + + constructor( + private config: GuiConfigService, + private workflowUtilService: WorkflowUtilService + ) {} + + private get enabled(): boolean { + return this.config.env.pythonNotebookMigrationEnabled; + } + + private assertEnabled(): void { + if (!this.enabled) { + throw new Error("Notebook migration feature is disabled"); + } + } + + private parseJsonResponse(raw: string, context: string): any { + // Trim first, then strip optional markdown code fences (```json ... ``` or ``` ... ```) + const cleaned = raw + .trim() + .replace(/^```[a-zA-Z]*\s*/, "") + .replace(/\s*```$/, "") + .trim(); + try { + return JSON.parse(cleaned); + } catch (err) { + throw new Error(`Failed to parse LLM ${context} response as JSON: ${(err as Error).message}`); + } + } + + /** + * Initialize a new LLM session with Texera documentation + */ + public initialize(modelType: string = "gpt-5-mini", apiKey: string = "dummy"): void { + this.assertEnabled(); + this.model = createOpenAI({ + baseURL: new URL(`${AppSettings.getApiEndpoint()}`, document.baseURI).toString(), + // apiKey is required by the library for creating the OpenAI compatible client; + // For security reason, we store the apiKey at the backend, thus the value is dummy here. + apiKey: apiKey, + }).chat(modelType); + + this.messages = [ + ...NotebookMigrationLLM.DOCUMENTATION.map( + (doc): ModelMessage => ({ + role: "system", + content: doc, + }) + ), + ]; + + this.initialized = true; + } + + /** + * Verify the connection to the LLM using the given API key + */ + public async verifyConnection(): Promise<boolean> { + if (!this.enabled) return false; + if (!this.initialized) { + throw new Error("LLM session not initialized"); + } + + try { + await generateText({ + model: this.model, + messages: [ + { + role: "user", + content: "ping", + }, + ], + maxOutputTokens: 10, + }); + + return true; + } catch (err) { + console.error("API key verification failed:", err); + return false; + } + } + + /** + * Send a prompt and receive a response. + * All prior documentation and conversation is preserved. + */ + private async sendPrompt(prompt: string): Promise<string> { + if (!this.initialized) { + throw new Error("LLM session not initialized"); + } + + this.messages.push({ + role: "user", + content: prompt, + }); + + const result = await generateText({ + model: this.model, + messages: this.messages, + }); + + this.messages.push({ + role: "assistant", + content: result.text, + }); + + return result.text; + } + + /** + * Send a Jupyter Notebook to be converted into a workflow and mapping. + */ + public async convertNotebookToWorkflow(notebook: Notebook): Promise<string> { + this.assertEnabled(); + if (!this.initialized) { + throw new Error("LLM session not initialized"); + } + + const codeCells = notebook.cells.filter(cell => cell.cell_type === "code"); + const notebookString = codeCells + .map(cell => { + const uuid = String(cell.metadata.uuid); + return `# START ${uuid}\n${cell.source}\n# END ${uuid}`; + }) + .join("\n\n"); + + const workflow = await this.sendPrompt(`${WORKFLOW_PROMPT}\n${notebookString}`); + const mapping = await this.sendPrompt(MAPPING_PROMPT); Review Comment: `this.messages` accumulates user/assistant pairs across every `sendPrompt` call. A single `convertNotebookToWorkflow` adds 4 messages (workflow prompt + response + mapping prompt + response) on top of the 8 doc system messages from `initialize`. If a consumer calls `convertNotebookToWorkflow` twice without re-`initialize`-ing or calling `close()` between (e.g., user uploads notebook A, then B, in the same session), the second call sends the entire first conversation back to the LLM — tokens balloon, context window can blow, and the prior workflow/mapping output can pollute the second response. The class is built for a one-shot lifecycle (`initialize` → `verifyConnection` → `convertNotebookToWorkflow` → `close`) but doesn't enforce it. Three options: 1. **Reset at the top:** at the start of `convertNotebookToWorkflow`, reset `this.messages` to just the documentation prelude (same shape as what `initialize` builds). This makes each conversion stateless w.r.t. prior conversions. 2. **Throw on re-use:** track whether this session has already converted once; on second call, throw an error directing callers to `close()` + `initialize()` first. 3. **Document the one-shot lifecycle** clearly in the class doc-comment so the next consumer (and the spec) doesn't trip on it. Option 1 is probably most forgiving for callers; option 3 is the cheapest. The current spec uses a fresh `makeLLM()` per test so it doesn't catch the multi-call case, which would be a useful test to add either way. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
