From 95ceb37b4e98ecdf4d603997e39614792e6306c2 Mon Sep 17 00:00:00 2001 From: Ryan Zhang Date: Wed, 27 May 2026 21:33:50 -0700 Subject: [PATCH 01/20] added NotebookMigrationLLM class to handle communication with LiteLLM and prompt file --- .../notebook-migration/migration-llm.ts | 299 +++++++++++++ .../notebook-migration/migration-prompts.ts | 410 ++++++++++++++++++ 2 files changed, 709 insertions(+) create mode 100644 frontend/src/app/workspace/service/notebook-migration/migration-llm.ts create mode 100644 frontend/src/app/workspace/service/notebook-migration/migration-prompts.ts diff --git a/frontend/src/app/workspace/service/notebook-migration/migration-llm.ts b/frontend/src/app/workspace/service/notebook-migration/migration-llm.ts new file mode 100644 index 00000000000..0ba919b2a75 --- /dev/null +++ b/frontend/src/app/workspace/service/notebook-migration/migration-llm.ts @@ -0,0 +1,299 @@ +/** + * 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 { firstValueFrom, from, Observable, of } from "rxjs"; +import { map } from "rxjs/operators"; +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 { + 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: any[]; + operatorPositions: Record; + links: any[]; + commentBoxes: any[]; + settings: { + dataTransferBatchSize: number; + }; +} + +interface CombinedMapping { + operator_to_cell: Record; + cell_to_operator: Record; +} + +@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, + ]; + + /** + * Initialize a new LLM session with Texera documentation + */ + public initialize(modelType: string = "gpt-5-mini", apiKey: string = "dummy"): void { + 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 { + 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 sendPrompt(prompt: string): Observable { + if (!this.initialized) { + throw new Error("LLM session not initialized"); + } + + this.messages.push({ + role: "user", + content: prompt, + }); + + return from( + generateText({ + model: this.model, + messages: this.messages, + }) + ).pipe( + map(result => { + 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> { + 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 firstValueFrom(this.sendPrompt(`${WORKFLOW_PROMPT}\n${notebookString}`)); + const mapping = await firstValueFrom(this.sendPrompt(MAPPING_PROMPT)); + + // Remove ```json blocks and parse + const udfLLMResponse = JSON.parse(workflow.replace(/^```json\s*|```$/g, "").trim()); + + const workflowJSON: WorkflowJSON = { + operators: [], + operatorPositions: {}, + links: [], + commentBoxes: [], + settings: { + dataTransferBatchSize: 400, + }, + }; + + const udfMappingToUUID: Record = {}; + + Object.entries(udfLLMResponse.code).forEach(([udfId, udfCode], i) => { + const udfUUID = `PythonUDFV2-operator-${uuidv4()}`; + udfMappingToUUID[udfId] = udfUUID; + + let udfOutputColumns: { attributeName: string; attributeType: string }[] = []; + if (udfLLMResponse.outputs && udfLLMResponse.outputs[udfId]) { + udfOutputColumns = udfLLMResponse.outputs[udfId].map((attr: string) => ({ + attributeName: attr, + attributeType: "binary", + })); + } + + // Add UDF to operators + workflowJSON.operators.push({ + operatorID: udfUUID, + operatorType: "PythonUDFV2", + operatorVersion: "3d69fdcedbb409b47162c4b55406c77e54abe416", + operatorProperties: { + code: udfCode, + workers: 1, + retainInputColumns: false, + outputColumns: udfOutputColumns, + }, + inputPorts: [ + { + portID: "input-0", + displayName: "", + allowMultiInputs: true, + isDynamicPort: false, + dependencies: [], + }, + ], + outputPorts: [ + { + portID: "output-0", + displayName: "", + allowMultiInputs: false, + isDynamicPort: false, + }, + ], + showAdvanced: false, + isDisabled: false, + customDisplayName: udfId, + dynamicInputPorts: true, + dynamicOutputPorts: true, + }); + + // Add UDF to operatorPositions + workflowJSON.operatorPositions[udfUUID] = { x: 140 * (i + 1), y: 0 }; + }); + + // Add links/edges + (udfLLMResponse.edges || []).forEach(([source, target]: [string, string]) => { + workflowJSON.links.push({ + linkID: `link-${uuidv4()}`, + source: { + operatorID: udfMappingToUUID[source], + portID: "output-0", + }, + target: { + operatorID: udfMappingToUUID[target], + portID: "input-0", + }, + }); + }); + + // Parse mapping + const parsedMapping: Record = JSON.parse(mapping.replace(/^```json\s*|```$/g, "").trim()); + + const udfToCell: Record = {}; + const cellToUdf: Record = {}; + + Object.entries(parsedMapping).forEach(([udf, cells]) => { + const udfUUID = udfMappingToUUID[udf]; + udfToCell[udfUUID] = cells; + cells.forEach(cell => { + if (!cellToUdf[cell]) { + cellToUdf[cell] = [udfUUID]; + } else { + cellToUdf[cell].push(udfUUID); + } + }); + }); + + const workflowNotebookMapping: CombinedMapping = { + operator_to_cell: udfToCell, + cell_to_operator: cellToUdf, + }; + + const result = JSON.stringify({ workflowJSON, workflowNotebookMapping }); + return of(result); + } + + /** + * Closes the session. + * Clears all context and releases references. + */ + public close(): void { + this.messages = []; + this.model = null; + this.initialized = false; + } +} diff --git a/frontend/src/app/workspace/service/notebook-migration/migration-prompts.ts b/frontend/src/app/workspace/service/notebook-migration/migration-prompts.ts new file mode 100644 index 00000000000..0a752112349 --- /dev/null +++ b/frontend/src/app/workspace/service/notebook-migration/migration-prompts.ts @@ -0,0 +1,410 @@ +/** + * 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. + */ + +// TEXERA DOCUMENTATION + +// https://github.com/Texera/texera/wiki/Guide-to-Use-a-Python-UDF +export const TEXERA_OVERVIEW = ` +You are a robust compiler that takes python code and translates it to our personal workflow environment Texera that uses python. + + Texera is a data analytics tool that uses workflows to do machine learning and data analytics computation. User's are able to drag and drop operators and connect their inputs and outputs in a workflow graphical user interface, which the code we are going to create. + +Texera is able to use Python user defined functions. Documentation of a Python UDF in Texera follows: + Process Data APIs + +There are three APIs to process the data in different units. + + Tuple API. + + class ProcessTupleOperator(UDFOperatorV2): + +def process_tuple(self, tuple_: Tuple, port: int) -> Iterator[Optional[TupleLike]]: +yield tuple_ + +Tuple API takes one input tuple from a port at a time. It returns an iterator of optional TupleLike instances. A TupleLike is any data structure that supports key-value pairs, such as pytexera.Tuple, dict, defaultdict, NamedTuple, etc. + + Tuple API is useful for implementing functional operations which are applied to tuples one by one, such as map, reduce, and filter. + + Table API. + + class ProcessTableOperator(UDFTableOperator): + +def process_table(self, table: Table, port: int) -> Iterator[Optional[TableLike]]: +yield table + +Table API consumes a Table at a time, which consists of all the tuples from a port. It returns an iterator of optional TableLike instances. A TableLike is a collection of TupleLike, and currently, we support pytexera.Table and pandas.DataFrame as a TableLike instance. More flexible types will be supported down the road. + + Table API is useful for implementing blocking operations that will consume all the data from one port, such as join, sort, and machine learning training. + + Batch API. + + class ProcessBatchOperator(UDFBatchOperator): + +BATCH_SIZE = 10 + +def process_batch(self, batch: Batch, port: int) -> Iterator[Optional[BatchLike]]: +yield batch + +Batch API consumes a batch of tuples at a time. Similar to Table, a Batch is also a collection of Tuples; however, its size is defined by the BATCH_SIZE, and one port can have multiple batches. It returns an iterator of optional BatchLike instances. A BatchLike is a collection of TupleLike, and currently, we support pytexera.Batch and pandas.DataFrame as a BatchLike instance. More flexible types will be supported down the road. + + The Batch API serves as a hybrid API combining the features of both the Tuple and Table APIs. It is particularly valuable for striking a balance between time and space considerations, offering a trade-off that optimizes efficiency. + + All three APIs can return an empty iterator by yield None. + + The template code for a Python UDF follows: MAKE SURE TO USE THE CLASS NAMES AND FUNCTIONS DEFINED, THIS IS A MUST FOR THE PROGRAM TO WORK. SELECT 1 OUT OF THE 3 PROCESSING OPERATOR FUNCTIONS TO BUILD DEPENDINGO ON THE CONTEXT OF CODE TRANSLATION. +# Choose from the following templates: + # +# from pytexera import * +# +# class ProcessTupleOperator(UDFOperatorV2): +# +# @overrides +# def process_tuple(self, tuple_: Tuple, port: int) -> Iterator[Optional[TupleLike]]: +# yield tuple_ +# +# class ProcessBatchOperator(UDFBatchOperator): +# BATCH_SIZE = 10 # must be a positive integer +# +# @overrides +# def process_batch(self, batch: Batch, port: int) -> Iterator[Optional[BatchLike]]: +# yield batch +# +# class ProcessTableOperator(UDFTableOperator): +# +# @overrides +# def process_table(self, table: Table, port: int) -> Iterator[Optional[TableLike]]: +# yield table +`; + +// https://github.com/Texera/texera/blob/1fa249a9d55d4dcad36d93e093c2faed5c4434f0/core/amber/src/main/python/core/models/tuple.py +export const TUPLE_DOCUMENTATION = ` +### **Tuple Class Overview** + +The \`Tuple\` class is a **lazy-evaluated** data structure designed for efficient field storage and access. It provides: + + 1. **Support for Multiple Data Sources**: +* Can be initialized from a \`TupleLike\` object, such as a \`pandas.Series\`, \`OrderedDict\`, or another \`Tuple\` instance. +* Works with \`ArrowTableTupleProvider\` to access \`pyarrow.Table\` data. +2. **Lazy Field Evaluation**: +* Field values can be either **directly stored values** or **lazy accessors** (\`field_accessor\`). +* If a field is accessed and is an accessor, it is evaluated and cached. +3. **Schema (Schema) Enforcement**: + * A \`Tuple\` can be created without a schema but can be **finalized** with one using \`finalize(schema)\`, which: +* **Casts field values** (e.g., \`NaN → None\`, \`Object → Bytes\`). +* **Validates field completeness**, ensuring all fields match the \`Schema\`. +4. **Pythonic Access Patterns**: +* **Index-based access**: \`tuple["field_name"]\` or \`tuple[index]\` retrieves field values. +* **Dictionary-like operations**: \`tuple.as_dict()\` returns an \`OrderedDict\`, and \`tuple.as_series()\` converts to a \`pandas.Series\`. +* **Iterable support**: \`for field in tuple\` iterates over field values. +5. **Hashing and Comparisons**: +* Implements \`__hash__\` using a Java-like hashing algorithm, allowing usage as dictionary keys. +* Implements \`__eq__\`, supporting equality checks based on field contents. +6. **Partial Data Extraction**: +* \`tuple.get_partial_tuple(attribute_names)\` returns a new \`Tuple\` instance containing only the specified fields. +`; + +// https://github.com/Texera/texera/blob/1fa249a9d55d4dcad36d93e093c2faed5c4434f0/core/amber/src/main/python/core/models/table.py +export const TABLE_DOCUMENTATION = `### **Table Class Overview** + +The \`Table\` class extends \`pandas.DataFrame\`, providing **structured Tuple-based data management**. It is designed to integrate seamlessly with \`Tuple\` objects. + +#### **Key Features:** + +1. **Flexible Construction:** +* Can be initialized from various sources: +* Another \`Table\` (\`from_table(table)\`) +* A \`pandas.DataFrame\` (\`from_data_frame(df)\`) +* A list/iterator of \`TupleLike\` objects (\`from_tuple_likes(tuple_likes)\`) +* Ensures all \`Tuple\` objects in a \`Table\` have **consistent field names**. +2. **Tuple Conversion:** +* \`as_tuples()\`: Converts the table rows into an **iterator of Tuple instances**, preserving the row order. +3. **Equality Comparison (__eq__):** +* Supports **row-wise equality checks** by comparing the underlying \`Tuple\` objects. +4. **Universal Tuple Output (all_output_to_tuple):** +* A helper function to convert **various data types** into \`Tuple\` iterators, supporting: +* \`None\` → \`[None]\` +* \`Table\` → \`as_tuples()\` +* \`pandas.DataFrame\` → Converted into a \`Table\`, then to Tuples +* \`List[TupleLike]\` → Converted to \`Tuple\` instances +* A single \`TupleLike\` or \`Tuple\` → Wrapped in an iterator + +#### **Relation to Tuple:** + +* \`Table\` **stores multiple Tuple objects** and ensures schema consistency across rows. +* Provides an **efficient bridge** between \`Tuple\`-based data and \`pandas.DataFrame\`, enabling compatibility with Python's data analysis tools. +`; + +// https://github.com/Texera/texera/blob/42d803310c180978a9f02992f0e05556796b293c/core/amber/src/main/python/core/models/operator.py +export const OPERATOR_DOCUMENTATION = `### **Operator Class Overview** + +The \`Operator\` class is an **abstract base class (ABC)** for all operators, defining the fundamental structure for processing \`Tuple\`, \`Batch\`, and \`Table\` data in a workflow. + +#### **Key Features & Hierarchy** + +1. **Base Operator Class**: +* Defines lifecycle methods: \`open()\` and \`close()\`. +* Supports a **source flag (is_source)** to distinguish source operators from others. +2. **Tuple-Based Processing (TupleOperatorV2)**: +* Processes individual \`Tuple\` objects through \`process_tuple(tuple_, port)\`. +* Calls \`on_finish(port)\` when an input port is exhausted. +3. **Types of Operators**: +* **SourceOperator**: +* Produces data via \`produce()\`, yielding \`TupleLike\` or \`TableLike\` objects. +* Overrides \`on_finish(port)\` to output produced data. +* **BatchOperator**: +* Collects tuples into batches (\`BATCH_SIZE\`) before processing via \`process_batch(batch, port)\`. +* Converts processed batches (typically \`pandas.DataFrame\`) into \`Tuple\` output. +* **TableOperator**: +* Collects tuples into a \`Table\` before processing via \`process_table(table, port)\`. +* Converts processed \`Table\` output back into tuples. +4. **Data Flow & Processing**: +* Operators receive data **tuple-by-tuple**, **batch-by-batch**, or **table-by-table** depending on the type. +* Results are **iterators** of transformed data (\`TupleLike\`, \`BatchLike\`, or \`TableLike\`). +5. **Deprecated TupleOperator**: +* The older version of \`TupleOperator\` is deprecated in favor of \`TupleOperatorV2\`. + +#### Relation to Tuple and Table + +* Operators **consume and transform** \`Tuple\` and \`Table\` data within a workflow. +* **Tuple-based operators** process row-wise, while **Table operators** handle structured table transformations. +* **Source operators** initiate the data flow by generating tuples or tables.`; + +export const UDF_INPUT_PORT_DOCUMENTATION = ` +Python UDF operators support multiple input and output ports, allowing a single operator to receive different types of data from various upstream operators. In the process_tuple(self, tuple_: Tuple, port: int) function in ProcessTupleOperator and the process_table(self, table: Table, port: int) function in ProcessTableOperator, the port parameter indicates the input port. The port numbers are assigned in order, starting from 0 to N, from top to bottom. When input data have different schemas, it is necessary to assign them to different input ports. However, if all input data share the same schema, additional ports are not required. In both ProcessTupleOperator and ProcessTableOperator, there is an on_finish(self, port: int) function that is executed only after all the tuples from the specified port are processed. + +Using this knowledge, for situations where multiple upstream UDFs act as input to a single UDF, we can introduce an intermediary UDF that collects all of the input data and reformats it into a single table, which is then passed as input to the original next downstream UDF. When it is necessary for this to occur in your translation from notebook to UDFs, include the intermediary UDF and make sure that it and the next operator that uses its output is formatted correctly and handles the data transfer properly. +`; + +export const EXAMPLE_OF_GOOD_CONVERSION = ` +Here is an example of python code translated into a compatible Texera UDF that gives output that abides the output schema compatible with the Texera workflow operators for tuples. Other operators do not always follow this strict format, but the yielding output structure is important. + +Python Code (high level idea): We have a python code that given some data, we limit the number of data. + +Texera Operator code: +from pytexera import * + +class ProcessTupleOperator(UDFOperatorV2): +def __init__(self): +self.limit = 10 +self.count = 0 +@overrides +def process_tuple(self, tuple_: Tuple, port: int) -> Iterator[Optional[TupleLike]]: +if(self.count < self.limit): +self.count += 1 +yield tuple_ + +`; + +export const VISUALIZER_DOCUMENTATION = ` +Texera requires a unique way of generating visualizations from ML libraries: +1. Ensures one yield per operator (per Texera’s UDF constraints). +2. Uses Plotly for visualization and outputs results as embeddable HTML. +3. Error handling is built-in to notify users when data is missing. +`; + +export const EXAMPLE_OF_MULTIPLE_UDF_CONVERSION = ` +Here is an example of breaking up python code into multiple Texera UDFs. Format your response structure exactly like the given example. The "code" key contains a dictionary of the UDF ID's with their respective code. The "edges" key contains a list of pairs that contains the connections between UDFs. The "outputs" key contains a dictionary of the UDF ID's with a list of variable names that they yield in the UDF code. The UDFs can branch and merge, it does not have to be a linear chain depending on your implementation. + +Original Code: +\`\`\`python +# START CELL1 +import pandas as pd +from sklearn.model_selection import train_test_split +from sklearn.ensemble import RandomForestClassifier +from sklearn.svm import SVC +from sklearn.tree import DecisionTreeClassifier +from sklearn.linear_model import LogisticRegression +from sklearn.metrics import accuracy_score +from sklearn.preprocessing import StandardScaler +import matplotlib.pyplot as plt +# END CELL1 + +# START CELL2 +# Load the dataset +file_path = 'diabetes.csv' +data = pd.read_csv(file_path) +# END CELL2 + +# START CELL3 +# Remove duplicate rows +data = data.drop_duplicates() + +# Remove rows with null values +data = data.dropna() +# END CELL3 + +# START CELL4 +# Print the minimum, maximum, and mean for all fields +print("Minimum values:\n", data.min()) +print("\nMaximum values:\n", data.max()) +print("\nMean values:\n", data.mean()) +# END CELL 4 + +# START CELL5 +# Create a boxplot for the 'Pregnancies' field +plt.figure(figsize=(8, 6)) +plt.boxplot(data['Pregnancies'], vert=False, patch_artist=True) +plt.title('Boxplot of Pregnancies') +plt.xlabel('Number of Pregnancies') +plt.show() +# END CELL5 + +# START CELL6 +# Separate features and target variable +X = data.drop('Outcome', axis=1) +y = data['Outcome'] +# END CELL6 + +# START CELL7 +# Split data into training and testing sets (80% train, 20% test) +X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42) + +scaler = StandardScaler() +X_train = scaler.fit_transform(X_train) +X_test = scaler.transform(X_test) +# END CELL7 + +# START CELL8 +# Train Random Forest model +rf_model = RandomForestClassifier(random_state=42) +rf_model.fit(X_train, y_train) +rf_pred = rf_model.predict(X_test) +rf_accuracy = accuracy_score(y_test, rf_pred) +print(f"Random Forest Accuracy: {rf_accuracy:.2%}") +# END CELL8 + +# START CELL9 +# Train SVM model +svm_model = SVC(random_state=42) +svm_model.fit(X_train, y_train) +svm_pred = svm_model.predict(X_test) +svm_accuracy = accuracy_score(y_test, svm_pred) +print(f"SVM Accuracy: {svm_accuracy:.2%}") +# END CELL9 +\`\`\` + +Texera UDF conversion: +\`\`\`json +{ + "code": { + "UDF1": "# UDF1\nfrom pytexera import *\nimport pandas as pd\nfrom typing import Iterator, Optional\n\nclass ProcessTableOperator(UDFTableOperator):\n\n @overrides\n def process_table(self, table: Table, port: int) -> Iterator[Optional[TableLike]]:\n # Remove duplicate rows\n data = table.drop_duplicates()\n\n # Remove rows with null values\n data = data.dropna()\n\n # Calculate statistics\n min_values = data.min()\n max_values = data.max()\n mean_values = data.mean()\n\n # Create a DataFrame to yield\n result_table = pd.DataFrame({\n 'min_values': [min_values],\n 'max_values': [max_values],\n 'mean_values': [mean_values],\n 'data': [data]\n })\n\n yield Table(result_table)", + "UDF2": "# UDF2\nfrom pytexera import *\nimport pandas as pd\nimport plotly.express as px\nimport plotly.io\nfrom typing import Iterator, Optional\n\nclass ProcessTableOperator(UDFTableOperator):\n def render_error(self, error_msg):\n return '''

Boxplot is not available.

\n

Reason is: {}

\n '''.format(error_msg)\n\n @overrides\n def process_table(self, table: Table, port: int) -> Iterator[Optional[TableLike]]:\n data = table['data'].iloc[0]\n\n if data.empty:\n yield {'html-content': self.render_error('input table is empty.')}\n return\n\n # Create a boxplot for the 'Pregnancies' field\n fig = px.box(data, x='Pregnancies')\n fig.update_layout(margin=dict(l=0, r=0, t=0, b=0))\n\n # Convert fig to HTML content\n html = plotly.io.to_html(fig, include_plotlyjs='cdn', auto_play=False)\n yield {'html-content': html}", + "UDF3": "# UDF3\nfrom pytexera import *\nimport pandas as pd\nfrom sklearn.model_selection import train_test_split\nfrom sklearn.preprocessing import StandardScaler\nfrom typing import Iterator, Optional\n\nclass ProcessTableOperator(UDFTableOperator):\n\n @overrides\n def process_table(self, table: Table, port: int) -> Iterator[Optional[TableLike]]:\n data = table['data'].iloc[0]\n\n # Separate features and target variable\n X = data.drop('Outcome', axis=1)\n y = data['Outcome']\n\n # Split data into training and testing sets (80% train, 20% test)\n X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)\n\n scaler = StandardScaler()\n X_train = scaler.fit_transform(X_train)\n X_test = scaler.transform(X_test)\n\n # Create a DataFrame to yield\n result_table = pd.DataFrame({\n 'X_train': [X_train], 'X_test': [X_test],\n 'y_train': [y_train], 'y_test': [y_test]\n })\n\n yield Table(result_table)", + "UDF4": "# UDF4\nfrom pytexera import *\nimport pandas as pd\nfrom sklearn.ensemble import RandomForestClassifier\nfrom sklearn.metrics import accuracy_score\nfrom typing import Iterator, Optional\n\nclass ProcessTableOperator(UDFTableOperator):\n\n @overrides\n def process_table(self, table: Table, port: int) -> Iterator[Optional[TableLike]]:\n X_train = table['X_train'].iloc[0]\n y_train = table['y_train'].iloc[0]\n X_test = table['X_test'].iloc[0]\n y_test = table['y_test'].iloc[0]\n\n # Train Random Forest model\n rf_model = RandomForestClassifier(random_state=42)\n rf_model.fit(X_train, y_train)\n rf_pred = rf_model.predict(X_test)\n rf_accuracy = accuracy_score(y_test, rf_pred)\n\n # Create a DataFrame to yield\n result_table = pd.DataFrame({\n 'rf_model': [rf_model],\n 'rf_accuracy': [rf_accuracy],\n 'X_test': [X_test],\n 'y_test': [y_test]\n })\n\n yield Table(result_table)", + "UDF5": "# UDF5\nfrom pytexera import *\nimport pandas as pd\nfrom sklearn.svm import SVC\nfrom sklearn.metrics import accuracy_score\nfrom typing import Iterator, Optional\n\nclass ProcessTableOperator(UDFTableOperator):\n\n @overrides\n def process_table(self, table: Table, port: int) -> Iterator[Optional[TableLike]]:\n X_train = table['X_train'].iloc[0]\n y_train = table['y_train'].iloc[0]\n X_test = table['X_test'].iloc[0]\n y_test = table['y_test'].iloc[0]\n\n # Train SVM model\n svm_model = SVC(random_state=42)\n svm_model.fit(X_train, y_train)\n svm_pred = svm_model.predict(X_test)\n svm_accuracy = accuracy_score(y_test, svm_pred)\n\n # Create a DataFrame to yield\n result_table = pd.DataFrame({\n 'svm_model': [svm_model],\n 'svm_accuracy': [svm_accuracy],\n 'X_test': [X_test],\n 'y_test': [y_test]\n })\n\n yield Table(result_table)" + }, + "edges": [ + ["UDF1", "UDF2"], + ["UDF1", "UDF3"], + ["UDF3", "UDF4"], + ["UDF3", "UDF5"] + ], + "outputs": { + "UDF1": ["min_values", "max_values", "mean_values", "data"], + "UDF2": ["html-content"], + "UDF3": ["X_train", "X_test", "y_train", "y_test"], + "UDF4": ["rf_model", "rf_accuracy", "X_test", "y_test"], + "UDF5": ["svm_model", "svm_accuracy", "X_test", "y_test"] + } +} +\`\`\` +`; + +export const WORKFLOW_PROMPT = `You are an expert in Python coding and workflow systems. +Many users of Texera system are non-technical, but the notebooks they provide are written by technical people. +They want to convert their notebooks to Texera workflows. +Your goal is to help convert these notebooks into a Texera workflow that non-technical users can use directly. +So do not remove or modify any classes or functions, preserve their names and structure as they are. +Ensure that all essential logic remains intact. +Create multiple Texera UDF codes using the provided Python code. +Number each UDF, starting at 1 and incrementing, by starting with a comment that states that UDF number. + +Use the class and function names as shown in ProcessTupleOperator, ProcessTableOperator, and ProcessBatchOperator. +Do not change the class names, function names, or input parameters. +Use the ones that make sense and split the code meaningfully as instructed. + +Use the starter code provided for Python UDFs. + +Use the documentation of Table, Tuple, or Batch to work with parameters within Texera UDF. +Do not import other libraries to define these types. + +There is no need for an __init__ function. Assume all inputs are valid pandas DataFrames, +so do not use .to_pandas(), .to_dataframe(), etc. Do not load data from a file in the first UDF, assume +that the data is already given to you in the table parameter. +Ensure proper data flow between functions. Separate operators as if they will run in different files. + +Current UDF operators can only have one output. Build a dataframe to yield all necessary variables +and data. Ensure proper data flow for each UDF and all information is yielded (including training +and testing data) if subsequent UDFs need them. + +Ensure all necessary imports are included in each UDF code block. + +Each UDF operator should be in its own Python code block. Do not combine them into a single block. +Ensure import statements cover all used functions and separate them as necessary. + +It is VERY important that all of the original code in the Jupyter notebook is represented in the generated workflow. +Make sure that nothing in the original is removed and that the semantic meaning of what the original code was doing is retained. +If there are user-defined Python classes, include the entire class definition in the appropriate UDF(s) that use that class. +Always include the code that defines the class inside of every distinct UDF that uses that constructs an object of that class. +Python classes are allowed in Texera UDFs and follow the same semantics as standard Python. +They can be defined outside of ProcessTableOperator, ProcessTupleOperator, and ProcessBatchOperator. + +Return only the JSON formatted response, do not give any explanation. +Make sure the response is a valid JSON structure, including closing all braces and not including commas after the last element. +Follow this JSON format (don't reuse the values, this is just the format). 'code', 'edges', and 'outputs' are all their own key's, do not nest any of these in another one and make sure to close their braces: +{ +"code": { +"UDF1": "code for UDF1 goes here", +"UDF2": "code for UDF2 goes here" +}, +"edges": [ +["UDF1", "UDF2"] +], +"outputs": { +"UDF1": ["min_values", "max_values", "mean_values", "data"], +"UDF2": ["html-content"] +} +} +Make sure only the keys in the code section appear in the edges and outputs sections. Do not include any extraneous fields. +Do not include any extraneous UDF's in the code field that include empty strings. +Give ALL of the code, do not omit anything or use placeholders for code. Make sure ALL code in the original is translated over. +Use only unescaped single quotes inside of the code values for the UDF's, do not use escaped double quotes. +Convert following the instructions and examples given. Here is the code: +`; + +export const MAPPING_PROMPT = ` +Here is an example of a mapping generated between the given example Python code and the Texera UDFs using their CELL and UDF IDs. Cell IDs are designated by the UUID following '# START'. The format should be kept the same. +{ +"UDF1": [ +"CEll3", +"CELL4" +], +"UDF2": [ +"CELL5" +], +"UDF3": [ +"CELL6", +"CELL7" +] +"UDF4": [ +"CELL8" +] +} +Now create a mapping for the UDFs and the original code. Link the code blocks marked by 'START ' and 'END ' with the UDF UUID's. The code between them should be equivalent. Multiple cells can be mapped to the same UDF if the code they contain are the same. There could be any number of cells and UDFs, so only create the correct number in the mapping. Only give the mapping. +`; From 78b9ef33ea4e30c58cd19fb7d02e35a2d708c44c Mon Sep 17 00:00:00 2001 From: Ryan Zhang Date: Wed, 27 May 2026 22:30:16 -0700 Subject: [PATCH 02/20] added flag to disable the service --- .../service/notebook-migration/migration-llm.ts | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/frontend/src/app/workspace/service/notebook-migration/migration-llm.ts b/frontend/src/app/workspace/service/notebook-migration/migration-llm.ts index 0ba919b2a75..5360712b965 100644 --- a/frontend/src/app/workspace/service/notebook-migration/migration-llm.ts +++ b/frontend/src/app/workspace/service/notebook-migration/migration-llm.ts @@ -18,6 +18,7 @@ */ import { Injectable } from "@angular/core"; +import { GuiConfigService } from "../../../common/service/gui-config.service"; import { firstValueFrom, from, Observable, of } from "rxjs"; import { map } from "rxjs/operators"; import { createOpenAI } from "@ai-sdk/openai"; @@ -79,10 +80,17 @@ export class NotebookMigrationLLM { EXAMPLE_OF_MULTIPLE_UDF_CONVERSION, ]; + constructor(private config: GuiConfigService) {} + + private get enabled(): boolean { + return this.config.env.pythonNotebookMigrationEnabled; + } + /** * Initialize a new LLM session with Texera documentation */ public initialize(modelType: string = "gpt-5-mini", apiKey: string = "dummy"): void { + if (!this.enabled) return; this.model = createOpenAI({ baseURL: new URL(`${AppSettings.getApiEndpoint()}`, document.baseURI).toString(), // apiKey is required by the library for creating the OpenAI compatible client; @@ -106,6 +114,7 @@ export class NotebookMigrationLLM { * Verify the connection to the LLM using the given API key */ public async verifyConnection(): Promise { + if (!this.enabled) return false; if (!this.initialized) { throw new Error("LLM session not initialized"); } @@ -164,6 +173,7 @@ export class NotebookMigrationLLM { * Send a Jupyter Notebook to be converted into a workflow and mapping. */ public async convertNotebookToWorkflow(notebook: Notebook): Promise> { + if (!this.enabled) throw new Error("Notebook migration feature is disabled"); if (!this.initialized) { throw new Error("LLM session not initialized"); } From 5a9eee36dd27c90f56fcba465fab65e6e5aef894 Mon Sep 17 00:00:00 2001 From: Ryan Zhang Date: Tue, 16 Jun 2026 16:19:57 -0700 Subject: [PATCH 03/20] refactored to return plain promises from migration LLM client --- .../notebook-migration/migration-llm.ts | 39 ++++++++----------- 1 file changed, 16 insertions(+), 23 deletions(-) diff --git a/frontend/src/app/workspace/service/notebook-migration/migration-llm.ts b/frontend/src/app/workspace/service/notebook-migration/migration-llm.ts index 5360712b965..22889e9162b 100644 --- a/frontend/src/app/workspace/service/notebook-migration/migration-llm.ts +++ b/frontend/src/app/workspace/service/notebook-migration/migration-llm.ts @@ -19,8 +19,6 @@ import { Injectable } from "@angular/core"; import { GuiConfigService } from "../../../common/service/gui-config.service"; -import { firstValueFrom, from, Observable, of } from "rxjs"; -import { map } from "rxjs/operators"; import { createOpenAI } from "@ai-sdk/openai"; import { generateText, type ModelMessage } from "ai"; import { AppSettings } from "../../../common/app-setting"; @@ -142,7 +140,7 @@ export class NotebookMigrationLLM { * Send a prompt and receive a response. * All prior documentation and conversation is preserved. */ - private sendPrompt(prompt: string): Observable { + private async sendPrompt(prompt: string): Promise { if (!this.initialized) { throw new Error("LLM session not initialized"); } @@ -152,27 +150,23 @@ export class NotebookMigrationLLM { content: prompt, }); - return from( - generateText({ - model: this.model, - messages: this.messages, - }) - ).pipe( - map(result => { - this.messages.push({ - role: "assistant", - content: result.text, - }); - - return result.text; - }) - ); + 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> { + public async convertNotebookToWorkflow(notebook: Notebook): Promise { if (!this.enabled) throw new Error("Notebook migration feature is disabled"); if (!this.initialized) { throw new Error("LLM session not initialized"); @@ -186,8 +180,8 @@ export class NotebookMigrationLLM { }) .join("\n\n"); - const workflow = await firstValueFrom(this.sendPrompt(`${WORKFLOW_PROMPT}\n${notebookString}`)); - const mapping = await firstValueFrom(this.sendPrompt(MAPPING_PROMPT)); + const workflow = await this.sendPrompt(`${WORKFLOW_PROMPT}\n${notebookString}`); + const mapping = await this.sendPrompt(MAPPING_PROMPT); // Remove ```json blocks and parse const udfLLMResponse = JSON.parse(workflow.replace(/^```json\s*|```$/g, "").trim()); @@ -293,8 +287,7 @@ export class NotebookMigrationLLM { cell_to_operator: cellToUdf, }; - const result = JSON.stringify({ workflowJSON, workflowNotebookMapping }); - return of(result); + return JSON.stringify({ workflowJSON, workflowNotebookMapping }); } /** From f8f77751e91ed549089a9fade0a6ddc1911805ac Mon Sep 17 00:00:00 2001 From: Ryan Zhang Date: Tue, 16 Jun 2026 16:28:31 -0700 Subject: [PATCH 04/20] improved migration prompt clarity and json validity --- .../notebook-migration/migration-prompts.ts | 22 +++++++++++-------- 1 file changed, 13 insertions(+), 9 deletions(-) diff --git a/frontend/src/app/workspace/service/notebook-migration/migration-prompts.ts b/frontend/src/app/workspace/service/notebook-migration/migration-prompts.ts index 0a752112349..c2375f94f37 100644 --- a/frontend/src/app/workspace/service/notebook-migration/migration-prompts.ts +++ b/frontend/src/app/workspace/service/notebook-migration/migration-prompts.ts @@ -23,7 +23,7 @@ export const TEXERA_OVERVIEW = ` You are a robust compiler that takes python code and translates it to our personal workflow environment Texera that uses python. - Texera is a data analytics tool that uses workflows to do machine learning and data analytics computation. User's are able to drag and drop operators and connect their inputs and outputs in a workflow graphical user interface, which the code we are going to create. + Texera is a data analytics tool that uses workflows to do machine learning and data analytics computation. Users are able to drag and drop operators and connect their inputs and outputs in a workflow graphical user interface, which the code we are going to create. Texera is able to use Python user defined functions. Documentation of a Python UDF in Texera follows: Process Data APIs @@ -67,7 +67,7 @@ Batch API consumes a batch of tuples at a time. Similar to Table, a Batch is als All three APIs can return an empty iterator by yield None. - The template code for a Python UDF follows: MAKE SURE TO USE THE CLASS NAMES AND FUNCTIONS DEFINED, THIS IS A MUST FOR THE PROGRAM TO WORK. SELECT 1 OUT OF THE 3 PROCESSING OPERATOR FUNCTIONS TO BUILD DEPENDINGO ON THE CONTEXT OF CODE TRANSLATION. + The template code for a Python UDF follows: MAKE SURE TO USE THE CLASS NAMES AND FUNCTIONS DEFINED, THIS IS A MUST FOR THE PROGRAM TO WORK. SELECT 1 OUT OF THE 3 PROCESSING OPERATOR FUNCTIONS TO BUILD DEPENDING ON THE CONTEXT OF CODE TRANSLATION. # Choose from the following templates: # # from pytexera import * @@ -219,7 +219,7 @@ Texera requires a unique way of generating visualizations from ML libraries: `; export const EXAMPLE_OF_MULTIPLE_UDF_CONVERSION = ` -Here is an example of breaking up python code into multiple Texera UDFs. Format your response structure exactly like the given example. The "code" key contains a dictionary of the UDF ID's with their respective code. The "edges" key contains a list of pairs that contains the connections between UDFs. The "outputs" key contains a dictionary of the UDF ID's with a list of variable names that they yield in the UDF code. The UDFs can branch and merge, it does not have to be a linear chain depending on your implementation. +Here is an example of breaking up python code into multiple Texera UDFs. Format your response structure exactly like the given example. The "code" key contains a dictionary of the UDF ID's with their respective code. The "edges" key contains a list of pairs that contains the connections between UDFs. The "outputs" key contains a dictionary of the UDF ID's with a list of the output column names of the DataFrame that the UDF yields. The UDFs can branch and merge, it does not have to be a linear chain depending on your implementation. Original Code: \`\`\`python @@ -345,8 +345,10 @@ Use the documentation of Table, Tuple, or Batch to work with parameters within T Do not import other libraries to define these types. There is no need for an __init__ function. Assume all inputs are valid pandas DataFrames, -so do not use .to_pandas(), .to_dataframe(), etc. Do not load data from a file in the first UDF, assume -that the data is already given to you in the table parameter. +so do not use .to_pandas(), .to_dataframe(), etc. Do not load data from a file in the first UDF; +the workflow's source operator supplies the initial data, so assume it is already given to you in the +table parameter. Replacing file-loading code with this input is the one exception to preserving all +original code (see below). Ensure proper data flow between functions. Separate operators as if they will run in different files. Current UDF operators can only have one output. Build a dataframe to yield all necessary variables @@ -360,12 +362,14 @@ Ensure import statements cover all used functions and separate them as necessary It is VERY important that all of the original code in the Jupyter notebook is represented in the generated workflow. Make sure that nothing in the original is removed and that the semantic meaning of what the original code was doing is retained. +The only exception is data-loading code (e.g. pd.read_csv); it is represented by the workflow's input/source operator rather than copied into a UDF. If there are user-defined Python classes, include the entire class definition in the appropriate UDF(s) that use that class. Always include the code that defines the class inside of every distinct UDF that uses that constructs an object of that class. Python classes are allowed in Texera UDFs and follow the same semantics as standard Python. They can be defined outside of ProcessTableOperator, ProcessTupleOperator, and ProcessBatchOperator. Return only the JSON formatted response, do not give any explanation. +Do not wrap the JSON in markdown code fences. Output raw JSON only. Make sure the response is a valid JSON structure, including closing all braces and not including commas after the last element. Follow this JSON format (don't reuse the values, this is just the format). 'code', 'edges', and 'outputs' are all their own key's, do not nest any of these in another one and make sure to close their braces: { @@ -384,7 +388,7 @@ Follow this JSON format (don't reuse the values, this is just the format). 'code Make sure only the keys in the code section appear in the edges and outputs sections. Do not include any extraneous fields. Do not include any extraneous UDF's in the code field that include empty strings. Give ALL of the code, do not omit anything or use placeholders for code. Make sure ALL code in the original is translated over. -Use only unescaped single quotes inside of the code values for the UDF's, do not use escaped double quotes. +The value of each UDF must be a valid JSON string: escape newlines, quotes, and backslashes correctly so that the decoded string is runnable Python. Use whichever quotes the Python code requires. Convert following the instructions and examples given. Here is the code: `; @@ -392,7 +396,7 @@ export const MAPPING_PROMPT = ` Here is an example of a mapping generated between the given example Python code and the Texera UDFs using their CELL and UDF IDs. Cell IDs are designated by the UUID following '# START'. The format should be kept the same. { "UDF1": [ -"CEll3", +"CELL3", "CELL4" ], "UDF2": [ @@ -401,10 +405,10 @@ Here is an example of a mapping generated between the given example Python code "UDF3": [ "CELL6", "CELL7" -] +], "UDF4": [ "CELL8" ] } -Now create a mapping for the UDFs and the original code. Link the code blocks marked by 'START ' and 'END ' with the UDF UUID's. The code between them should be equivalent. Multiple cells can be mapped to the same UDF if the code they contain are the same. There could be any number of cells and UDFs, so only create the correct number in the mapping. Only give the mapping. +Now create a mapping for the UDFs and the original code. Link the code blocks marked by 'START ' and 'END ' with the UDF UUID's. The code between them should be equivalent. Multiple cells can be mapped to the same UDF when that UDF implements the logic of those cells. There could be any number of cells and UDFs, so only create the correct number in the mapping. Only give the mapping. `; From 0c894b396913f03ae6eb69a7304f175cdddda040 Mon Sep 17 00:00:00 2001 From: Ryan Zhang Date: Tue, 16 Jun 2026 16:33:20 -0700 Subject: [PATCH 05/20] refactored to throw consistently when migration feature is disabled --- .../service/notebook-migration/migration-llm.ts | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/frontend/src/app/workspace/service/notebook-migration/migration-llm.ts b/frontend/src/app/workspace/service/notebook-migration/migration-llm.ts index 22889e9162b..0865614b171 100644 --- a/frontend/src/app/workspace/service/notebook-migration/migration-llm.ts +++ b/frontend/src/app/workspace/service/notebook-migration/migration-llm.ts @@ -84,11 +84,17 @@ export class NotebookMigrationLLM { return this.config.env.pythonNotebookMigrationEnabled; } + private assertEnabled(): void { + if (!this.enabled) { + throw new Error("Notebook migration feature is disabled"); + } + } + /** * Initialize a new LLM session with Texera documentation */ public initialize(modelType: string = "gpt-5-mini", apiKey: string = "dummy"): void { - if (!this.enabled) return; + 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; @@ -167,7 +173,7 @@ export class NotebookMigrationLLM { * Send a Jupyter Notebook to be converted into a workflow and mapping. */ public async convertNotebookToWorkflow(notebook: Notebook): Promise { - if (!this.enabled) throw new Error("Notebook migration feature is disabled"); + this.assertEnabled(); if (!this.initialized) { throw new Error("LLM session not initialized"); } From 2448bfed3e2657fae88726224889abe23156ce94 Mon Sep 17 00:00:00 2001 From: Ryan Zhang Date: Tue, 16 Jun 2026 16:46:21 -0700 Subject: [PATCH 06/20] hardened llm json parsing against fences and malformed output --- .../notebook-migration/migration-llm.ts | 18 ++++++++++++++++-- 1 file changed, 16 insertions(+), 2 deletions(-) diff --git a/frontend/src/app/workspace/service/notebook-migration/migration-llm.ts b/frontend/src/app/workspace/service/notebook-migration/migration-llm.ts index 0865614b171..0efa924069a 100644 --- a/frontend/src/app/workspace/service/notebook-migration/migration-llm.ts +++ b/frontend/src/app/workspace/service/notebook-migration/migration-llm.ts @@ -90,6 +90,20 @@ export class NotebookMigrationLLM { } } + 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 */ @@ -190,7 +204,7 @@ export class NotebookMigrationLLM { const mapping = await this.sendPrompt(MAPPING_PROMPT); // Remove ```json blocks and parse - const udfLLMResponse = JSON.parse(workflow.replace(/^```json\s*|```$/g, "").trim()); + const udfLLMResponse = this.parseJsonResponse(workflow, "workflow"); const workflowJSON: WorkflowJSON = { operators: [], @@ -271,7 +285,7 @@ export class NotebookMigrationLLM { }); // Parse mapping - const parsedMapping: Record = JSON.parse(mapping.replace(/^```json\s*|```$/g, "").trim()); + const parsedMapping: Record = this.parseJsonResponse(mapping, "mapping"); const udfToCell: Record = {}; const cellToUdf: Record = {}; From d39c79dd2b534f127d70ed95dcf9d63afd82d93c Mon Sep 17 00:00:00 2001 From: Ryan Zhang Date: Fri, 19 Jun 2026 11:09:57 -0700 Subject: [PATCH 07/20] build UDF operators from live PythonUDFV2 schema --- .../notebook-migration/migration-llm.ts | 53 ++++++------------- 1 file changed, 17 insertions(+), 36 deletions(-) diff --git a/frontend/src/app/workspace/service/notebook-migration/migration-llm.ts b/frontend/src/app/workspace/service/notebook-migration/migration-llm.ts index 0efa924069a..9a229ee3641 100644 --- a/frontend/src/app/workspace/service/notebook-migration/migration-llm.ts +++ b/frontend/src/app/workspace/service/notebook-migration/migration-llm.ts @@ -23,6 +23,8 @@ 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, @@ -47,7 +49,7 @@ export interface Notebook { } interface WorkflowJSON { - operators: any[]; + operators: OperatorPredicate[]; operatorPositions: Record; links: any[]; commentBoxes: any[]; @@ -78,7 +80,10 @@ export class NotebookMigrationLLM { EXAMPLE_OF_MULTIPLE_UDF_CONVERSION, ]; - constructor(private config: GuiConfigService) {} + constructor( + private config: GuiConfigService, + private workflowUtilService: WorkflowUtilService + ) {} private get enabled(): boolean { return this.config.env.pythonNotebookMigrationEnabled; @@ -219,9 +224,6 @@ export class NotebookMigrationLLM { const udfMappingToUUID: Record = {}; Object.entries(udfLLMResponse.code).forEach(([udfId, udfCode], i) => { - const udfUUID = `PythonUDFV2-operator-${uuidv4()}`; - udfMappingToUUID[udfId] = udfUUID; - let udfOutputColumns: { attributeName: string; attributeType: string }[] = []; if (udfLLMResponse.outputs && udfLLMResponse.outputs[udfId]) { udfOutputColumns = udfLLMResponse.outputs[udfId].map((attr: string) => ({ @@ -230,43 +232,22 @@ export class NotebookMigrationLLM { })); } - // Add UDF to operators - workflowJSON.operators.push({ - operatorID: udfUUID, - operatorType: "PythonUDFV2", - operatorVersion: "3d69fdcedbb409b47162c4b55406c77e54abe416", + // Build the operator from the live PythonUDFV2 schema so the operatorVersion, ports, and + // property defaults track the backend definition, then overlay the generated code/outputs. + const base = this.workflowUtilService.getNewOperatorPredicate("PythonUDFV2", udfId); + const operator: OperatorPredicate = { + ...base, operatorProperties: { + ...base.operatorProperties, code: udfCode, - workers: 1, retainInputColumns: false, outputColumns: udfOutputColumns, }, - inputPorts: [ - { - portID: "input-0", - displayName: "", - allowMultiInputs: true, - isDynamicPort: false, - dependencies: [], - }, - ], - outputPorts: [ - { - portID: "output-0", - displayName: "", - allowMultiInputs: false, - isDynamicPort: false, - }, - ], - showAdvanced: false, - isDisabled: false, - customDisplayName: udfId, - dynamicInputPorts: true, - dynamicOutputPorts: true, - }); + }; - // Add UDF to operatorPositions - workflowJSON.operatorPositions[udfUUID] = { x: 140 * (i + 1), y: 0 }; + udfMappingToUUID[udfId] = operator.operatorID; + workflowJSON.operators.push(operator); + workflowJSON.operatorPositions[operator.operatorID] = { x: 140 * (i + 1), y: 0 }; }); // Add links/edges From 90ebaf3f1ba96411246e9ff633f60180c04ebec3 Mon Sep 17 00:00:00 2001 From: Ryan Zhang Date: Mon, 22 Jun 2026 09:26:59 -0700 Subject: [PATCH 08/20] add unit tests for notebook-to-workflow conversion --- .../notebook-migration/migration-llm.spec.ts | 241 ++++++++++++++++++ 1 file changed, 241 insertions(+) create mode 100644 frontend/src/app/workspace/service/notebook-migration/migration-llm.spec.ts diff --git a/frontend/src/app/workspace/service/notebook-migration/migration-llm.spec.ts b/frontend/src/app/workspace/service/notebook-migration/migration-llm.spec.ts new file mode 100644 index 00000000000..faed15fb9a5 --- /dev/null +++ b/frontend/src/app/workspace/service/notebook-migration/migration-llm.spec.ts @@ -0,0 +1,241 @@ +/** + * 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 { NotebookMigrationLLM, Notebook } from "./migration-llm"; +import { GuiConfigService } from "../../../common/service/gui-config.service"; +import { WorkflowUtilService } from "../workflow-graph/util/workflow-util.service"; +import { generateText } from "ai"; +import type { Mock } from "vitest"; + +// The LLM transport and OpenAI client are mocked so the tests exercise only the +// deterministic transformation (parsing, operator/edge construction, cell<->operator mapping). +vi.mock("ai", () => ({ generateText: vi.fn() })); +vi.mock("@ai-sdk/openai", () => ({ + createOpenAI: vi.fn(() => ({ chat: vi.fn(() => ({})) })), +})); + +const mockGenerateText = generateText as unknown as Mock; + +describe("NotebookMigrationLLM", () => { + let opIdCounter = 0; + let stubUtil: WorkflowUtilService; + + // Build a fresh, initialized session with stubbed dependencies. The stubbed + // getNewOperatorPredicate hands out deterministic ids (PythonUDFV2-0, -1, ...). + function makeLLM(): NotebookMigrationLLM { + const stubConfig = { + env: { pythonNotebookMigrationEnabled: true }, + } as unknown as GuiConfigService; + + stubUtil = { + getNewOperatorPredicate: vi.fn((operatorType: string, customDisplayName?: string) => ({ + operatorID: `${operatorType}-${opIdCounter++}`, + operatorType, + operatorVersion: "test-version", + operatorProperties: { workers: 1, defaultEnv: true, envName: "" }, + inputPorts: [{ portID: "input-0", disallowMultiInputs: false }], + outputPorts: [{ portID: "output-0" }], + showAdvanced: false, + isDisabled: false, + customDisplayName, + dynamicInputPorts: true, + dynamicOutputPorts: true, + })), + } as unknown as WorkflowUtilService; + + const llm = new NotebookMigrationLLM(stubConfig, stubUtil); + llm.initialize(); + return llm; + } + + function codeCell(uuid: string | undefined, source: string) { + return { cell_type: "code", metadata: uuid === undefined ? {} : { uuid }, source }; + } + + // Queue the two responses convertNotebookToWorkflow consumes, in order. + function mockResponses(workflowResponse: string, mappingResponse: string) { + mockGenerateText + .mockResolvedValueOnce({ text: workflowResponse }) + .mockResolvedValueOnce({ text: mappingResponse }); + } + + beforeEach(() => { + opIdCounter = 0; + mockGenerateText.mockReset(); + }); + + describe("convertNotebookToWorkflow", () => { + it("builds operators, links, positions, and a bidirectional mapping", async () => { + const notebook: Notebook = { + cells: [codeCell("CELL1", "print(1)"), codeCell("CELL2", "print(2)")], + }; + mockResponses( + JSON.stringify({ + code: { UDF1: "code1", UDF2: "code2" }, + edges: [["UDF1", "UDF2"]], + outputs: { UDF1: ["a", "b"], UDF2: ["c"] }, + }), + JSON.stringify({ UDF1: ["CELL1"], UDF2: ["CELL2"] }) + ); + + const { workflowJSON, workflowNotebookMapping } = JSON.parse( + await makeLLM().convertNotebookToWorkflow(notebook) + ); + + expect(workflowJSON.operators.map((op: any) => op.operatorID)).toEqual([ + "PythonUDFV2-0", + "PythonUDFV2-1", + ]); + expect(workflowJSON.operators[0].operatorProperties).toMatchObject({ + code: "code1", + retainInputColumns: false, + }); + expect(workflowJSON.operatorPositions).toEqual({ + "PythonUDFV2-0": { x: 140, y: 0 }, + "PythonUDFV2-1": { x: 280, y: 0 }, + }); + expect(workflowJSON.links).toHaveLength(1); + expect(workflowJSON.links[0].source).toEqual({ operatorID: "PythonUDFV2-0", portID: "output-0" }); + expect(workflowJSON.links[0].target).toEqual({ operatorID: "PythonUDFV2-1", portID: "input-0" }); + expect(workflowNotebookMapping.operator_to_cell).toEqual({ + "PythonUDFV2-0": ["CELL1"], + "PythonUDFV2-1": ["CELL2"], + }); + expect(workflowNotebookMapping.cell_to_operator).toEqual({ + CELL1: ["PythonUDFV2-0"], + CELL2: ["PythonUDFV2-1"], + }); + }); + + // NOTE (C2): the attributeType is currently hardcoded to "binary". If C2 lands as + // "LLM returns per-column types", update the expected attributeType values here. + it("declares output columns with attributeType binary", async () => { + const notebook: Notebook = { cells: [codeCell("CELL1", "x = 1")] }; + mockResponses( + JSON.stringify({ code: { UDF1: "code1" }, edges: [], outputs: { UDF1: ["a", "b"] } }), + JSON.stringify({ UDF1: ["CELL1"] }) + ); + + const { workflowJSON } = JSON.parse(await makeLLM().convertNotebookToWorkflow(notebook)); + + expect(workflowJSON.operators[0].operatorProperties.outputColumns).toEqual([ + { attributeName: "a", attributeType: "binary" }, + { attributeName: "b", attributeType: "binary" }, + ]); + }); + + it("maps multiple cells onto the same UDF, and one cell onto multiple UDFs", async () => { + const notebook: Notebook = { + cells: [codeCell("CELL1", "a"), codeCell("CELL2", "b")], + }; + mockResponses( + JSON.stringify({ code: { UDF1: "c1", UDF2: "c2" }, edges: [], outputs: {} }), + JSON.stringify({ UDF1: ["CELL1", "CELL2"], UDF2: ["CELL1"] }) + ); + + const { workflowNotebookMapping } = JSON.parse(await makeLLM().convertNotebookToWorkflow(notebook)); + + expect(workflowNotebookMapping.operator_to_cell).toEqual({ + "PythonUDFV2-0": ["CELL1", "CELL2"], + "PythonUDFV2-1": ["CELL1"], + }); + expect(workflowNotebookMapping.cell_to_operator).toEqual({ + CELL1: ["PythonUDFV2-0", "PythonUDFV2-1"], + CELL2: ["PythonUDFV2-0"], + }); + }); + + it("produces a link with an undefined endpoint when an edge references an unknown UDF id", async () => { + const notebook: Notebook = { cells: [codeCell("CELL1", "a")] }; + mockResponses( + JSON.stringify({ code: { UDF1: "c1" }, edges: [["UDF1", "UDFX"]], outputs: {} }), + JSON.stringify({ UDF1: ["CELL1"] }) + ); + + const { workflowJSON } = JSON.parse(await makeLLM().convertNotebookToWorkflow(notebook)); + + // Documents current behavior: udfMappingToUUID["UDFX"] is undefined. + expect(workflowJSON.links[0].source.operatorID).toBe("PythonUDFV2-0"); + expect(workflowJSON.links[0].target.operatorID).toBeUndefined(); + }); + + it("handles empty code, edges, and outputs", async () => { + const notebook: Notebook = { cells: [] }; + mockResponses(JSON.stringify({ code: {}, edges: [], outputs: {} }), JSON.stringify({})); + + const { workflowJSON, workflowNotebookMapping } = JSON.parse( + await makeLLM().convertNotebookToWorkflow(notebook) + ); + + expect(workflowJSON.operators).toEqual([]); + expect(workflowJSON.links).toEqual([]); + expect(workflowNotebookMapping.operator_to_cell).toEqual({}); + expect(workflowNotebookMapping.cell_to_operator).toEqual({}); + }); + + it("emits the 'undefined' cell marker in the prompt when a code cell lacks metadata.uuid", async () => { + const notebook: Notebook = { cells: [codeCell(undefined, "print(1)")] }; + mockResponses( + JSON.stringify({ code: { UDF1: "c1" }, edges: [], outputs: {} }), + JSON.stringify({ UDF1: ["CELL1"] }) + ); + + await makeLLM().convertNotebookToWorkflow(notebook); + + // The notebook string (embedded in the workflow prompt) is sent to generateText. + // messages is a shared, mutated array, so search every message content rather than + // assuming a fixed index. + const allPromptContent = mockGenerateText.mock.calls + .flatMap(call => call[0].messages.map((m: any) => m.content)) + .join("\n"); + expect(allPromptContent).toContain("# START undefined"); + }); + }); + + describe("parseJsonResponse", () => { + // parseJsonResponse is private; cast to access it directly for focused coverage. + const parse = (raw: string) => (makeLLM() as any).parseJsonResponse(raw, "workflow"); + + it("parses bare JSON", () => { + expect(parse('{"a":1}')).toEqual({ a: 1 }); + }); + + it("strips a ```json fence", () => { + expect(parse('```json\n{"a":1}\n```')).toEqual({ a: 1 }); + }); + + it("strips a plain ``` fence", () => { + expect(parse('```\n{"a":1}\n```')).toEqual({ a: 1 }); + }); + + it("tolerates surrounding whitespace and newlines around the fence", () => { + expect(parse('\n\n ```json\n{"a":1}\n``` \n\n')).toEqual({ a: 1 }); + }); + + it("throws a contextual error on malformed JSON", () => { + expect(() => parse("not json")).toThrow("Failed to parse LLM workflow response as JSON"); + }); + + it("does not strip a fence preceded by prose (documents current limitation)", () => { + expect(() => parse('Here is the JSON: ```json\n{"a":1}\n```\nThanks!')).toThrow( + "Failed to parse LLM workflow response as JSON" + ); + }); + }); +}); From cf3e2694b647694cef8c4cd947971e97f6a5d60b Mon Sep 17 00:00:00 2001 From: Ryan Zhang Date: Mon, 22 Jun 2026 09:27:24 -0700 Subject: [PATCH 09/20] add @ai-sdk/openai dependency for the migration LLM client --- frontend/package.json | 1 + frontend/yarn.lock | 13 +++++++++++++ 2 files changed, 14 insertions(+) diff --git a/frontend/package.json b/frontend/package.json index 5c2795da9c8..03360c1d53d 100644 --- a/frontend/package.json +++ b/frontend/package.json @@ -21,6 +21,7 @@ "private": true, "dependencies": { "@abacritt/angularx-social-login": "2.3.0", + "@ai-sdk/openai": "2.0.67", "@ali-hm/angular-tree-component": "12.0.5", "@angular/animations": "21.2.10", "@angular/cdk": "21.2.8", diff --git a/frontend/yarn.lock b/frontend/yarn.lock index 25a59578001..fd981c030ae 100644 --- a/frontend/yarn.lock +++ b/frontend/yarn.lock @@ -30,6 +30,18 @@ __metadata: languageName: node linkType: hard +"@ai-sdk/openai@npm:2.0.67": + version: 2.0.67 + resolution: "@ai-sdk/openai@npm:2.0.67" + dependencies: + "@ai-sdk/provider": "npm:2.0.0" + "@ai-sdk/provider-utils": "npm:3.0.17" + peerDependencies: + zod: ^3.25.76 || ^4.1.8 + checksum: 10c0/7e5c407504d7902c17c816aaccd83f642a3b82012cd8467c8f58aef5f08a49b6c31fff775439d541d40b0c8b5b94cc384f18096d1968e23670e22a56fe82d8bd + languageName: node + linkType: hard + "@ai-sdk/provider-utils@npm:3.0.17": version: 3.0.17 resolution: "@ai-sdk/provider-utils@npm:3.0.17" @@ -11023,6 +11035,7 @@ __metadata: resolution: "gui@workspace:." dependencies: "@abacritt/angularx-social-login": "npm:2.3.0" + "@ai-sdk/openai": "npm:2.0.67" "@ali-hm/angular-tree-component": "npm:12.0.5" "@angular-builders/custom-webpack": "npm:21.0.3" "@angular-devkit/build-angular": "npm:21.2.8" From c0203ff6f756c833c1e6b5db684e7e616a5cb7ed Mon Sep 17 00:00:00 2001 From: Ryan Zhang Date: Mon, 22 Jun 2026 09:33:44 -0700 Subject: [PATCH 10/20] fix format --- .../notebook-migration/migration-llm.spec.ts | 17 ++++------------- 1 file changed, 4 insertions(+), 13 deletions(-) diff --git a/frontend/src/app/workspace/service/notebook-migration/migration-llm.spec.ts b/frontend/src/app/workspace/service/notebook-migration/migration-llm.spec.ts index faed15fb9a5..12c08560ccb 100644 --- a/frontend/src/app/workspace/service/notebook-migration/migration-llm.spec.ts +++ b/frontend/src/app/workspace/service/notebook-migration/migration-llm.spec.ts @@ -70,9 +70,7 @@ describe("NotebookMigrationLLM", () => { // Queue the two responses convertNotebookToWorkflow consumes, in order. function mockResponses(workflowResponse: string, mappingResponse: string) { - mockGenerateText - .mockResolvedValueOnce({ text: workflowResponse }) - .mockResolvedValueOnce({ text: mappingResponse }); + mockGenerateText.mockResolvedValueOnce({ text: workflowResponse }).mockResolvedValueOnce({ text: mappingResponse }); } beforeEach(() => { @@ -94,14 +92,9 @@ describe("NotebookMigrationLLM", () => { JSON.stringify({ UDF1: ["CELL1"], UDF2: ["CELL2"] }) ); - const { workflowJSON, workflowNotebookMapping } = JSON.parse( - await makeLLM().convertNotebookToWorkflow(notebook) - ); + const { workflowJSON, workflowNotebookMapping } = JSON.parse(await makeLLM().convertNotebookToWorkflow(notebook)); - expect(workflowJSON.operators.map((op: any) => op.operatorID)).toEqual([ - "PythonUDFV2-0", - "PythonUDFV2-1", - ]); + expect(workflowJSON.operators.map((op: any) => op.operatorID)).toEqual(["PythonUDFV2-0", "PythonUDFV2-1"]); expect(workflowJSON.operators[0].operatorProperties).toMatchObject({ code: "code1", retainInputColumns: false, @@ -179,9 +172,7 @@ describe("NotebookMigrationLLM", () => { const notebook: Notebook = { cells: [] }; mockResponses(JSON.stringify({ code: {}, edges: [], outputs: {} }), JSON.stringify({})); - const { workflowJSON, workflowNotebookMapping } = JSON.parse( - await makeLLM().convertNotebookToWorkflow(notebook) - ); + const { workflowJSON, workflowNotebookMapping } = JSON.parse(await makeLLM().convertNotebookToWorkflow(notebook)); expect(workflowJSON.operators).toEqual([]); expect(workflowJSON.links).toEqual([]); From 4dc3b1bc0cbd4ad1b834b89200287299b6430b98 Mon Sep 17 00:00:00 2001 From: Ryan Zhang Date: Tue, 23 Jun 2026 11:08:44 -0700 Subject: [PATCH 11/20] fix stale reference urls and cell marker in prompts --- .../service/notebook-migration/migration-prompts.ts | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/frontend/src/app/workspace/service/notebook-migration/migration-prompts.ts b/frontend/src/app/workspace/service/notebook-migration/migration-prompts.ts index c2375f94f37..2594d2d58e6 100644 --- a/frontend/src/app/workspace/service/notebook-migration/migration-prompts.ts +++ b/frontend/src/app/workspace/service/notebook-migration/migration-prompts.ts @@ -19,7 +19,7 @@ // TEXERA DOCUMENTATION -// https://github.com/Texera/texera/wiki/Guide-to-Use-a-Python-UDF +// https://github.com/apache/texera/wiki/Guide-to-Use-a-Python-UDF export const TEXERA_OVERVIEW = ` You are a robust compiler that takes python code and translates it to our personal workflow environment Texera that uses python. @@ -92,7 +92,7 @@ Batch API consumes a batch of tuples at a time. Similar to Table, a Batch is als # yield table `; -// https://github.com/Texera/texera/blob/1fa249a9d55d4dcad36d93e093c2faed5c4434f0/core/amber/src/main/python/core/models/tuple.py +// https://github.com/apache/texera/blob/main/amber/src/main/python/core/models/tuple.py export const TUPLE_DOCUMENTATION = ` ### **Tuple Class Overview** @@ -119,7 +119,7 @@ The \`Tuple\` class is a **lazy-evaluated** data structure designed for efficien * \`tuple.get_partial_tuple(attribute_names)\` returns a new \`Tuple\` instance containing only the specified fields. `; -// https://github.com/Texera/texera/blob/1fa249a9d55d4dcad36d93e093c2faed5c4434f0/core/amber/src/main/python/core/models/table.py +// https://github.com/apache/texera/blob/main/amber/src/main/python/core/models/table.py export const TABLE_DOCUMENTATION = `### **Table Class Overview** The \`Table\` class extends \`pandas.DataFrame\`, providing **structured Tuple-based data management**. It is designed to integrate seamlessly with \`Tuple\` objects. @@ -150,7 +150,7 @@ The \`Table\` class extends \`pandas.DataFrame\`, providing **structured Tuple-b * Provides an **efficient bridge** between \`Tuple\`-based data and \`pandas.DataFrame\`, enabling compatibility with Python's data analysis tools. `; -// https://github.com/Texera/texera/blob/42d803310c180978a9f02992f0e05556796b293c/core/amber/src/main/python/core/models/operator.py +// https://github.com/apache/texera/blob/main/amber/src/main/python/core/models/operator.py export const OPERATOR_DOCUMENTATION = `### **Operator Class Overview** The \`Operator\` class is an **abstract base class (ABC)** for all operators, defining the fundamental structure for processing \`Tuple\`, \`Batch\`, and \`Table\` data in a workflow. @@ -254,7 +254,7 @@ data = data.dropna() print("Minimum values:\n", data.min()) print("\nMaximum values:\n", data.max()) print("\nMean values:\n", data.mean()) -# END CELL 4 +# END CELL4 # START CELL5 # Create a boxplot for the 'Pregnancies' field From 84e87f1aef48862810a1ab94bcdb53a3e34f05c9 Mon Sep 17 00:00:00 2001 From: Ryan Zhang Date: Tue, 23 Jun 2026 11:57:40 -0700 Subject: [PATCH 12/20] reset conversation history between conversions --- .../notebook-migration/migration-llm.spec.ts | 25 +++++++++++++ .../notebook-migration/migration-llm.ts | 36 ++++++++++++++----- 2 files changed, 53 insertions(+), 8 deletions(-) diff --git a/frontend/src/app/workspace/service/notebook-migration/migration-llm.spec.ts b/frontend/src/app/workspace/service/notebook-migration/migration-llm.spec.ts index 12c08560ccb..eecc440aed9 100644 --- a/frontend/src/app/workspace/service/notebook-migration/migration-llm.spec.ts +++ b/frontend/src/app/workspace/service/notebook-migration/migration-llm.spec.ts @@ -197,6 +197,31 @@ describe("NotebookMigrationLLM", () => { .join("\n"); expect(allPromptContent).toContain("# START undefined"); }); + + it("resets conversation history between conversions so a prior notebook does not leak", async () => { + const llm = makeLLM(); + + // First conversion (notebook AAA) on the instance. + mockResponses( + JSON.stringify({ code: { UDF1: "codeAAA" }, edges: [], outputs: {} }), + JSON.stringify({ UDF1: ["AAA"] }) + ); + await llm.convertNotebookToWorkflow({ cells: [codeCell("AAA", "a = 1")] }); + + // Second conversion (notebook BBB) on the SAME instance, no close()/initialize() between. + mockResponses( + JSON.stringify({ code: { UDF1: "codeBBB" }, edges: [], outputs: {} }), + JSON.stringify({ UDF1: ["BBB"] }) + ); + await llm.convertNotebookToWorkflow({ cells: [codeCell("BBB", "b = 2")] }); + + // The 3rd generateText call is the workflow prompt of the second conversion. + const secondConversionMessages = mockGenerateText.mock.calls[2][0].messages.map((m: any) => m.content).join("\n"); + + expect(secondConversionMessages).toContain("# START BBB"); + expect(secondConversionMessages).not.toContain("AAA"); + expect(secondConversionMessages).not.toContain("codeAAA"); + }); }); describe("parseJsonResponse", () => { diff --git a/frontend/src/app/workspace/service/notebook-migration/migration-llm.ts b/frontend/src/app/workspace/service/notebook-migration/migration-llm.ts index 9a229ee3641..37c32980d6e 100644 --- a/frontend/src/app/workspace/service/notebook-migration/migration-llm.ts +++ b/frontend/src/app/workspace/service/notebook-migration/migration-llm.ts @@ -63,6 +63,16 @@ interface CombinedMapping { cell_to_operator: Record; } +/** + * Wraps a single LLM chat session that converts a Jupyter notebook into a Texera + * workflow plus a cell<->operator mapping. + * + * Lifecycle: `initialize()` -> `verifyConnection()` (optional) -> + * `convertNotebookToWorkflow()` -> `close()`. The session keeps a running `messages` + * history shared by the prompts within one conversion. `convertNotebookToWorkflow()` + * resets that history to the documentation prelude at its start, so the same instance + * can convert multiple notebooks without leaking one conversion's context into the next. + */ @Injectable() export class NotebookMigrationLLM { private model: any; @@ -95,6 +105,19 @@ export class NotebookMigrationLLM { } } + /** + * Seed the conversation with the Texera documentation prelude, discarding any + * prior conversation. Used by initialize() and at the start of each conversion. + */ + private seedDocumentation(): void { + this.messages = NotebookMigrationLLM.DOCUMENTATION.map( + (doc): ModelMessage => ({ + role: "system", + content: doc, + }) + ); + } + private parseJsonResponse(raw: string, context: string): any { // Trim first, then strip optional markdown code fences (```json ... ``` or ``` ... ```) const cleaned = raw @@ -121,14 +144,7 @@ export class NotebookMigrationLLM { apiKey: apiKey, }).chat(modelType); - this.messages = [ - ...NotebookMigrationLLM.DOCUMENTATION.map( - (doc): ModelMessage => ({ - role: "system", - content: doc, - }) - ), - ]; + this.seedDocumentation(); this.initialized = true; } @@ -197,6 +213,10 @@ export class NotebookMigrationLLM { throw new Error("LLM session not initialized"); } + // Reset to the documentation prelude so a prior conversion's prompts/responses + // don't leak into this one. The two sendPrompt calls below still share history. + this.seedDocumentation(); + const codeCells = notebook.cells.filter(cell => cell.cell_type === "code"); const notebookString = codeCells .map(cell => { From 6f482c2d7cb43410930b37bd3381865fc9cdc40d Mon Sep 17 00:00:00 2001 From: Ryan Zhang Date: Tue, 23 Jun 2026 13:44:03 -0700 Subject: [PATCH 13/20] type terminal UDF outputs as string for viewable results --- .../notebook-migration/migration-llm.spec.ts | 23 ++++++++++++------- .../notebook-migration/migration-llm.ts | 15 +++++++++++- 2 files changed, 29 insertions(+), 9 deletions(-) diff --git a/frontend/src/app/workspace/service/notebook-migration/migration-llm.spec.ts b/frontend/src/app/workspace/service/notebook-migration/migration-llm.spec.ts index eecc440aed9..0e5ee422c36 100644 --- a/frontend/src/app/workspace/service/notebook-migration/migration-llm.spec.ts +++ b/frontend/src/app/workspace/service/notebook-migration/migration-llm.spec.ts @@ -116,20 +116,27 @@ describe("NotebookMigrationLLM", () => { }); }); - // NOTE (C2): the attributeType is currently hardcoded to "binary". If C2 lands as - // "LLM returns per-column types", update the expected attributeType values here. - it("declares output columns with attributeType binary", async () => { - const notebook: Notebook = { cells: [codeCell("CELL1", "x = 1")] }; + // Intermediate UDFs (a source of some edge) keep "binary" for object passing; terminal + // UDFs (no outgoing edge) default to "string" so the result panel renders typed values. + it("types intermediate UDF outputs as binary and terminal UDF outputs as string", async () => { + const notebook: Notebook = { cells: [codeCell("CELL1", "a"), codeCell("CELL2", "b")] }; mockResponses( - JSON.stringify({ code: { UDF1: "code1" }, edges: [], outputs: { UDF1: ["a", "b"] } }), - JSON.stringify({ UDF1: ["CELL1"] }) + JSON.stringify({ + code: { UDF1: "code1", UDF2: "code2" }, + edges: [["UDF1", "UDF2"]], + outputs: { UDF1: ["x"], UDF2: ["y"] }, + }), + JSON.stringify({ UDF1: ["CELL1"], UDF2: ["CELL2"] }) ); const { workflowJSON } = JSON.parse(await makeLLM().convertNotebookToWorkflow(notebook)); + // UDF1 is a source (intermediate) -> binary; UDF2 is terminal -> string. expect(workflowJSON.operators[0].operatorProperties.outputColumns).toEqual([ - { attributeName: "a", attributeType: "binary" }, - { attributeName: "b", attributeType: "binary" }, + { attributeName: "x", attributeType: "binary" }, + ]); + expect(workflowJSON.operators[1].operatorProperties.outputColumns).toEqual([ + { attributeName: "y", attributeType: "string" }, ]); }); diff --git a/frontend/src/app/workspace/service/notebook-migration/migration-llm.ts b/frontend/src/app/workspace/service/notebook-migration/migration-llm.ts index 37c32980d6e..fbd32895866 100644 --- a/frontend/src/app/workspace/service/notebook-migration/migration-llm.ts +++ b/frontend/src/app/workspace/service/notebook-migration/migration-llm.ts @@ -72,6 +72,11 @@ interface CombinedMapping { * history shared by the prompts within one conversion. `convertNotebookToWorkflow()` * resets that history to the documentation prelude at its start, so the same instance * can convert multiple notebooks without leaking one conversion's context into the next. + * + * Output column types: intermediate UDFs declare their output columns as `binary` so rich + * Python objects (DataFrames, arrays, models) round-trip between operators via pickle. + * Terminal UDFs (no outgoing edge) declare their outputs as `string` so the result panel + * renders viewable values rather than opaque binary blobs. */ @Injectable() export class NotebookMigrationLLM { @@ -243,12 +248,20 @@ export class NotebookMigrationLLM { const udfMappingToUUID: Record = {}; + // UDFs that are never the source of an edge are terminal (result-facing). Their outputs + // default to "string" so the result panel renders typed values; intermediate UDFs keep + // "binary" so rich objects (DataFrames, arrays, models) round-trip between operators via pickle. + const edgeSources = new Set( + (udfLLMResponse.edges || []).map(([source]: [string, string]) => source) + ); + Object.entries(udfLLMResponse.code).forEach(([udfId, udfCode], i) => { let udfOutputColumns: { attributeName: string; attributeType: string }[] = []; if (udfLLMResponse.outputs && udfLLMResponse.outputs[udfId]) { + const attributeType = edgeSources.has(udfId) ? "binary" : "string"; udfOutputColumns = udfLLMResponse.outputs[udfId].map((attr: string) => ({ attributeName: attr, - attributeType: "binary", + attributeType, })); } From 17f87a10f53422898a706629f782e96e762d9d87 Mon Sep 17 00:00:00 2001 From: Ryan Zhang Date: Tue, 23 Jun 2026 13:47:59 -0700 Subject: [PATCH 14/20] format fix --- .../app/workspace/service/notebook-migration/migration-llm.ts | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/frontend/src/app/workspace/service/notebook-migration/migration-llm.ts b/frontend/src/app/workspace/service/notebook-migration/migration-llm.ts index fbd32895866..436c6beb61d 100644 --- a/frontend/src/app/workspace/service/notebook-migration/migration-llm.ts +++ b/frontend/src/app/workspace/service/notebook-migration/migration-llm.ts @@ -251,9 +251,7 @@ export class NotebookMigrationLLM { // UDFs that are never the source of an edge are terminal (result-facing). Their outputs // default to "string" so the result panel renders typed values; intermediate UDFs keep // "binary" so rich objects (DataFrames, arrays, models) round-trip between operators via pickle. - const edgeSources = new Set( - (udfLLMResponse.edges || []).map(([source]: [string, string]) => source) - ); + const edgeSources = new Set((udfLLMResponse.edges || []).map(([source]: [string, string]) => source)); Object.entries(udfLLMResponse.code).forEach(([udfId, udfCode], i) => { let udfOutputColumns: { attributeName: string; attributeType: string }[] = []; From 95a610b41784eed6737bcbb88ccb2c6128963d0c Mon Sep 17 00:00:00 2001 From: Ryan Zhang Date: Tue, 23 Jun 2026 14:00:03 -0700 Subject: [PATCH 15/20] use gui config defaults for workflow settings --- .../service/notebook-migration/migration-llm.spec.ts | 8 +++++++- .../workspace/service/notebook-migration/migration-llm.ts | 8 ++++---- 2 files changed, 11 insertions(+), 5 deletions(-) diff --git a/frontend/src/app/workspace/service/notebook-migration/migration-llm.spec.ts b/frontend/src/app/workspace/service/notebook-migration/migration-llm.spec.ts index 0e5ee422c36..4fb535e0849 100644 --- a/frontend/src/app/workspace/service/notebook-migration/migration-llm.spec.ts +++ b/frontend/src/app/workspace/service/notebook-migration/migration-llm.spec.ts @@ -40,7 +40,11 @@ describe("NotebookMigrationLLM", () => { // getNewOperatorPredicate hands out deterministic ids (PythonUDFV2-0, -1, ...). function makeLLM(): NotebookMigrationLLM { const stubConfig = { - env: { pythonNotebookMigrationEnabled: true }, + env: { + pythonNotebookMigrationEnabled: true, + defaultDataTransferBatchSize: 400, + defaultExecutionMode: "PIPELINED", + }, } as unknown as GuiConfigService; stubUtil = { @@ -114,6 +118,8 @@ describe("NotebookMigrationLLM", () => { CELL1: ["PythonUDFV2-0"], CELL2: ["PythonUDFV2-1"], }); + // Settings come from GUI config defaults, not hardcoded values. + expect(workflowJSON.settings).toEqual({ dataTransferBatchSize: 400, executionMode: "PIPELINED" }); }); // Intermediate UDFs (a source of some edge) keep "binary" for object passing; terminal diff --git a/frontend/src/app/workspace/service/notebook-migration/migration-llm.ts b/frontend/src/app/workspace/service/notebook-migration/migration-llm.ts index 436c6beb61d..be734419d78 100644 --- a/frontend/src/app/workspace/service/notebook-migration/migration-llm.ts +++ b/frontend/src/app/workspace/service/notebook-migration/migration-llm.ts @@ -25,6 +25,7 @@ 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 { WorkflowSettings } from "../../../common/type/workflow"; import { TEXERA_OVERVIEW, TUPLE_DOCUMENTATION, @@ -53,9 +54,7 @@ interface WorkflowJSON { operatorPositions: Record; links: any[]; commentBoxes: any[]; - settings: { - dataTransferBatchSize: number; - }; + settings: WorkflowSettings; } interface CombinedMapping { @@ -242,7 +241,8 @@ export class NotebookMigrationLLM { links: [], commentBoxes: [], settings: { - dataTransferBatchSize: 400, + dataTransferBatchSize: this.config.env.defaultDataTransferBatchSize, + executionMode: this.config.env.defaultExecutionMode, }, }; From c133501f9cbbefbefab2d20d9709e30a235b9354 Mon Sep 17 00:00:00 2001 From: Ryan Zhang Date: Tue, 23 Jun 2026 14:06:23 -0700 Subject: [PATCH 16/20] handle nbformat array-form cell source --- .../notebook-migration/migration-llm.spec.ts | 24 +++++++++++++++++++ .../notebook-migration/migration-llm.ts | 7 ++++-- 2 files changed, 29 insertions(+), 2 deletions(-) diff --git a/frontend/src/app/workspace/service/notebook-migration/migration-llm.spec.ts b/frontend/src/app/workspace/service/notebook-migration/migration-llm.spec.ts index 4fb535e0849..7bcf19c3326 100644 --- a/frontend/src/app/workspace/service/notebook-migration/migration-llm.spec.ts +++ b/frontend/src/app/workspace/service/notebook-migration/migration-llm.spec.ts @@ -211,6 +211,30 @@ describe("NotebookMigrationLLM", () => { expect(allPromptContent).toContain("# START undefined"); }); + it("joins array-form cell source (nbformat lines) without inserting commas", async () => { + const notebook: Notebook = { + cells: [ + { + cell_type: "code", + metadata: { uuid: "CELL1" }, + source: ["import pandas as pd\n", "x = 1\n"], + }, + ], + }; + mockResponses( + JSON.stringify({ code: { UDF1: "c1" }, edges: [], outputs: {} }), + JSON.stringify({ UDF1: ["CELL1"] }) + ); + + await makeLLM().convertNotebookToWorkflow(notebook); + + const allPromptContent = mockGenerateText.mock.calls + .flatMap(call => call[0].messages.map((m: any) => m.content)) + .join("\n"); + expect(allPromptContent).toContain("import pandas as pd\nx = 1\n"); + expect(allPromptContent).not.toContain("import pandas as pd\n,"); + }); + it("resets conversation history between conversions so a prior notebook does not leak", async () => { const llm = makeLLM(); diff --git a/frontend/src/app/workspace/service/notebook-migration/migration-llm.ts b/frontend/src/app/workspace/service/notebook-migration/migration-llm.ts index be734419d78..18b1c5003a4 100644 --- a/frontend/src/app/workspace/service/notebook-migration/migration-llm.ts +++ b/frontend/src/app/workspace/service/notebook-migration/migration-llm.ts @@ -42,7 +42,8 @@ import { interface Cell { cell_type: string; metadata: { [key: string]: any }; - source: string; + // nbformat stores source as either a single string or an array of line strings. + source: string | string[]; } export interface Notebook { @@ -225,7 +226,9 @@ export class NotebookMigrationLLM { const notebookString = codeCells .map(cell => { const uuid = String(cell.metadata.uuid); - return `# START ${uuid}\n${cell.source}\n# END ${uuid}`; + // nbformat line arrays already include trailing newlines, so join with "". + const source = Array.isArray(cell.source) ? cell.source.join("") : cell.source; + return `# START ${uuid}\n${source}\n# END ${uuid}`; }) .join("\n\n"); From 880e2ad416decd3689ad1fd1190073031ea94f04 Mon Sep 17 00:00:00 2001 From: Ryan Zhang Date: Tue, 23 Jun 2026 15:38:12 -0700 Subject: [PATCH 17/20] require metadata.uuid on code cells before conversion --- .../notebook-migration/migration-llm.spec.ts | 18 ++++-------------- .../notebook-migration/migration-llm.ts | 8 ++++++++ 2 files changed, 12 insertions(+), 14 deletions(-) diff --git a/frontend/src/app/workspace/service/notebook-migration/migration-llm.spec.ts b/frontend/src/app/workspace/service/notebook-migration/migration-llm.spec.ts index 7bcf19c3326..16f871f5ce9 100644 --- a/frontend/src/app/workspace/service/notebook-migration/migration-llm.spec.ts +++ b/frontend/src/app/workspace/service/notebook-migration/migration-llm.spec.ts @@ -193,22 +193,12 @@ describe("NotebookMigrationLLM", () => { expect(workflowNotebookMapping.cell_to_operator).toEqual({}); }); - it("emits the 'undefined' cell marker in the prompt when a code cell lacks metadata.uuid", async () => { + it("rejects when a code cell is missing metadata.uuid", async () => { const notebook: Notebook = { cells: [codeCell(undefined, "print(1)")] }; - mockResponses( - JSON.stringify({ code: { UDF1: "c1" }, edges: [], outputs: {} }), - JSON.stringify({ UDF1: ["CELL1"] }) - ); - await makeLLM().convertNotebookToWorkflow(notebook); - - // The notebook string (embedded in the workflow prompt) is sent to generateText. - // messages is a shared, mutated array, so search every message content rather than - // assuming a fixed index. - const allPromptContent = mockGenerateText.mock.calls - .flatMap(call => call[0].messages.map((m: any) => m.content)) - .join("\n"); - expect(allPromptContent).toContain("# START undefined"); + await expect(makeLLM().convertNotebookToWorkflow(notebook)).rejects.toThrow(/metadata\.uuid/); + // It fails before prompting, so the LLM is never called. + expect(mockGenerateText).not.toHaveBeenCalled(); }); it("joins array-form cell source (nbformat lines) without inserting commas", async () => { diff --git a/frontend/src/app/workspace/service/notebook-migration/migration-llm.ts b/frontend/src/app/workspace/service/notebook-migration/migration-llm.ts index 18b1c5003a4..4a9d01421a0 100644 --- a/frontend/src/app/workspace/service/notebook-migration/migration-llm.ts +++ b/frontend/src/app/workspace/service/notebook-migration/migration-llm.ts @@ -223,6 +223,14 @@ export class NotebookMigrationLLM { this.seedDocumentation(); const codeCells = notebook.cells.filter(cell => cell.cell_type === "code"); + + // Every code cell must carry a unique metadata.uuid; it is the join key for the + // cell<->operator mapping. Without it, untagged cells collide on the "undefined" marker. + const untagged = codeCells.find(cell => cell.metadata?.uuid == null || String(cell.metadata.uuid).trim() === ""); + if (untagged) { + throw new Error("Notebook code cells must each have a metadata.uuid before conversion"); + } + const notebookString = codeCells .map(cell => { const uuid = String(cell.metadata.uuid); From 4387d6eb1a4594da014af2d57238479e3dfd755d Mon Sep 17 00:00:00 2001 From: Ryan Zhang Date: Tue, 23 Jun 2026 15:50:59 -0700 Subject: [PATCH 18/20] skip llm edges and mappings with unknown udf ids --- .../notebook-migration/migration-llm.spec.ts | 27 ++++++++++++++++--- .../notebook-migration/migration-llm.ts | 13 ++++++++- 2 files changed, 35 insertions(+), 5 deletions(-) diff --git a/frontend/src/app/workspace/service/notebook-migration/migration-llm.spec.ts b/frontend/src/app/workspace/service/notebook-migration/migration-llm.spec.ts index 16f871f5ce9..72e4c45310c 100644 --- a/frontend/src/app/workspace/service/notebook-migration/migration-llm.spec.ts +++ b/frontend/src/app/workspace/service/notebook-migration/migration-llm.spec.ts @@ -167,7 +167,8 @@ describe("NotebookMigrationLLM", () => { }); }); - it("produces a link with an undefined endpoint when an edge references an unknown UDF id", async () => { + it("skips (with a warning) an edge that references an unknown UDF id", async () => { + const warn = vi.spyOn(console, "warn").mockImplementation(() => {}); const notebook: Notebook = { cells: [codeCell("CELL1", "a")] }; mockResponses( JSON.stringify({ code: { UDF1: "c1" }, edges: [["UDF1", "UDFX"]], outputs: {} }), @@ -176,9 +177,27 @@ describe("NotebookMigrationLLM", () => { const { workflowJSON } = JSON.parse(await makeLLM().convertNotebookToWorkflow(notebook)); - // Documents current behavior: udfMappingToUUID["UDFX"] is undefined. - expect(workflowJSON.links[0].source.operatorID).toBe("PythonUDFV2-0"); - expect(workflowJSON.links[0].target.operatorID).toBeUndefined(); + // The dangling edge is dropped rather than producing an undefined endpoint. + expect(workflowJSON.links).toEqual([]); + expect(warn).toHaveBeenCalledWith(expect.stringContaining("UDFX")); + warn.mockRestore(); + }); + + it("skips (with a warning) a mapping entry that references an unknown UDF id", async () => { + const warn = vi.spyOn(console, "warn").mockImplementation(() => {}); + const notebook: Notebook = { cells: [codeCell("CELL1", "a")] }; + mockResponses( + JSON.stringify({ code: { UDF1: "c1" }, edges: [], outputs: {} }), + JSON.stringify({ UDF1: ["CELL1"], UDFTYPO: ["CELL1"] }) + ); + + const { workflowNotebookMapping } = JSON.parse(await makeLLM().convertNotebookToWorkflow(notebook)); + + // Only the valid UDF id survives in the mapping. + expect(workflowNotebookMapping.operator_to_cell).toEqual({ "PythonUDFV2-0": ["CELL1"] }); + expect(workflowNotebookMapping.cell_to_operator).toEqual({ CELL1: ["PythonUDFV2-0"] }); + expect(warn).toHaveBeenCalledWith(expect.stringContaining("UDFTYPO")); + warn.mockRestore(); }); it("handles empty code, edges, and outputs", async () => { diff --git a/frontend/src/app/workspace/service/notebook-migration/migration-llm.ts b/frontend/src/app/workspace/service/notebook-migration/migration-llm.ts index 4a9d01421a0..959dd9c553c 100644 --- a/frontend/src/app/workspace/service/notebook-migration/migration-llm.ts +++ b/frontend/src/app/workspace/service/notebook-migration/migration-llm.ts @@ -292,8 +292,15 @@ export class NotebookMigrationLLM { workflowJSON.operatorPositions[operator.operatorID] = { x: 140 * (i + 1), y: 0 }; }); - // Add links/edges + const knownUdfIds = new Set(Object.keys(udfMappingToUUID)); + + // Add links/edges. Skip (with a warning) any edge that references a UDF id the LLM + // never defined in `code`, rather than emitting a link with an undefined endpoint. (udfLLMResponse.edges || []).forEach(([source, target]: [string, string]) => { + if (!knownUdfIds.has(source) || !knownUdfIds.has(target)) { + console.warn(`Skipping edge with unknown UDF id: ${source} -> ${target}`); + return; + } workflowJSON.links.push({ linkID: `link-${uuidv4()}`, source: { @@ -314,6 +321,10 @@ export class NotebookMigrationLLM { const cellToUdf: Record = {}; Object.entries(parsedMapping).forEach(([udf, cells]) => { + if (!knownUdfIds.has(udf)) { + console.warn(`Skipping mapping entry with unknown UDF id: ${udf}`); + return; + } const udfUUID = udfMappingToUUID[udf]; udfToCell[udfUUID] = cells; cells.forEach(cell => { From 406c77b079d123936db91733edf766099009153a Mon Sep 17 00:00:00 2001 From: Ryan Zhang Date: Tue, 23 Jun 2026 16:00:18 -0700 Subject: [PATCH 19/20] tolerate prose-wrapped json in llm responses --- .../notebook-migration/migration-llm.spec.ts | 10 ++++---- .../notebook-migration/migration-llm.ts | 23 +++++++++++++------ 2 files changed, 22 insertions(+), 11 deletions(-) diff --git a/frontend/src/app/workspace/service/notebook-migration/migration-llm.spec.ts b/frontend/src/app/workspace/service/notebook-migration/migration-llm.spec.ts index 72e4c45310c..87dfd6d6e05 100644 --- a/frontend/src/app/workspace/service/notebook-migration/migration-llm.spec.ts +++ b/frontend/src/app/workspace/service/notebook-migration/migration-llm.spec.ts @@ -294,10 +294,12 @@ describe("NotebookMigrationLLM", () => { expect(() => parse("not json")).toThrow("Failed to parse LLM workflow response as JSON"); }); - it("does not strip a fence preceded by prose (documents current limitation)", () => { - expect(() => parse('Here is the JSON: ```json\n{"a":1}\n```\nThanks!')).toThrow( - "Failed to parse LLM workflow response as JSON" - ); + it("extracts fenced JSON even when surrounded by prose", () => { + expect(parse('Here is the JSON: ```json\n{"a":1}\n```\nThanks!')).toEqual({ a: 1 }); + }); + + it("extracts the outermost object from fence-less prose", () => { + expect(parse('Sure! {"a":1} hope that helps')).toEqual({ a: 1 }); }); }); }); diff --git a/frontend/src/app/workspace/service/notebook-migration/migration-llm.ts b/frontend/src/app/workspace/service/notebook-migration/migration-llm.ts index 959dd9c553c..7c9aa62a3b9 100644 --- a/frontend/src/app/workspace/service/notebook-migration/migration-llm.ts +++ b/frontend/src/app/workspace/service/notebook-migration/migration-llm.ts @@ -124,14 +124,23 @@ export class NotebookMigrationLLM { } 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(); + let text = raw.trim(); + + // Prefer the contents of a fenced code block if present (```json ... ``` or ``` ... ```), + // even when wrapped in prose. Otherwise fall back to the outermost {...} object. + const fenced = text.match(/```(?:[a-zA-Z]+)?\s*([\s\S]*?)```/); + if (fenced) { + text = fenced[1].trim(); + } else { + const firstBrace = text.indexOf("{"); + const lastBrace = text.lastIndexOf("}"); + if (firstBrace !== -1 && lastBrace > firstBrace) { + text = text.slice(firstBrace, lastBrace + 1); + } + } + try { - return JSON.parse(cleaned); + return JSON.parse(text); } catch (err) { throw new Error(`Failed to parse LLM ${context} response as JSON: ${(err as Error).message}`); } From 4bb8623e8c129c46869b46d0f4ceaac8bb9073c9 Mon Sep 17 00:00:00 2001 From: Ryan Zhang Date: Wed, 24 Jun 2026 09:53:45 -0700 Subject: [PATCH 20/20] authenticate llm proxy with the texera jwt --- .../service/notebook-migration/migration-llm.spec.ts | 3 ++- .../service/notebook-migration/migration-llm.ts | 12 +++++++----- 2 files changed, 9 insertions(+), 6 deletions(-) diff --git a/frontend/src/app/workspace/service/notebook-migration/migration-llm.spec.ts b/frontend/src/app/workspace/service/notebook-migration/migration-llm.spec.ts index 87dfd6d6e05..58c17cdfc3b 100644 --- a/frontend/src/app/workspace/service/notebook-migration/migration-llm.spec.ts +++ b/frontend/src/app/workspace/service/notebook-migration/migration-llm.spec.ts @@ -64,7 +64,8 @@ describe("NotebookMigrationLLM", () => { } as unknown as WorkflowUtilService; const llm = new NotebookMigrationLLM(stubConfig, stubUtil); - llm.initialize(); + // Pass an explicit token so tests don't depend on AuthService/localStorage state. + llm.initialize("gpt-5-mini", "test-token"); return llm; } diff --git a/frontend/src/app/workspace/service/notebook-migration/migration-llm.ts b/frontend/src/app/workspace/service/notebook-migration/migration-llm.ts index 7c9aa62a3b9..2922c3ee0e1 100644 --- a/frontend/src/app/workspace/service/notebook-migration/migration-llm.ts +++ b/frontend/src/app/workspace/service/notebook-migration/migration-llm.ts @@ -19,6 +19,7 @@ import { Injectable } from "@angular/core"; import { GuiConfigService } from "../../../common/service/gui-config.service"; +import { AuthService } from "../../../common/service/user/auth.service"; import { createOpenAI } from "@ai-sdk/openai"; import { generateText, type ModelMessage } from "ai"; import { AppSettings } from "../../../common/app-setting"; @@ -149,13 +150,14 @@ export class NotebookMigrationLLM { /** * Initialize a new LLM session with Texera documentation */ - public initialize(modelType: string = "gpt-5-mini", apiKey: string = "dummy"): void { + public initialize(modelType: string = "gpt-5-mini", accessToken: string = AuthService.getAccessToken() ?? ""): 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, + // The /api/chat/* LiteLLM proxy authenticates the caller with the Texera JWT. The AI SDK + // sends this value verbatim as `Authorization: Bearer `, so we pass the user's + // access token; the backend validates it, then substitutes the LiteLLM master key upstream. + apiKey: accessToken, }).chat(modelType); this.seedDocumentation(); @@ -164,7 +166,7 @@ export class NotebookMigrationLLM { } /** - * Verify the connection to the LLM using the given API key + * Verify the connection to the LLM using the current access token */ public async verifyConnection(): Promise { if (!this.enabled) return false;