zyratlo commented on code in PR #5260: URL: https://github.com/apache/texera/pull/5260#discussion_r3462903417
########## 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, + }, + }; Review Comment: Refactored in [95a610b](https://github.com/apache/texera/pull/5260/commits/95a610b41784eed6737bcbb88ccb2c6128963d0c) -- 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]
