@jupyterlite/ai 0.10.0 → 0.11.0

package/lib/tools/notebook.js

@@ -1,704 +0,0 @@
-import { CodeCell, MarkdownCell } from '@jupyterlab/cells';
-import { DocumentWidget } from '@jupyterlab/docregistry';
-import { NotebookPanel } from '@jupyterlab/notebook';
-import { tool } from 'ai';
-import { z } from 'zod';
-/**
- * Find a kernel name that matches the specified language
- */
-async function findKernelByLanguage(kernelSpecManager, language) {
-    try {
-        await kernelSpecManager.ready;
-        const specs = kernelSpecManager.specs;
-        if (!specs || !specs.kernelspecs) {
-            return 'python3'; // Final fallback
-        }
-        // If no language specified, return the default kernel
-        if (!language) {
-            return specs.default || Object.keys(specs.kernelspecs)[0] || 'python3';
-        }
-        // Normalize the language name for comparison
-        const normalizedLanguage = language.toLowerCase().trim();
-        // Find kernels that match the requested language
-        for (const [kernelName, kernelSpec] of Object.entries(specs.kernelspecs)) {
-            if (!kernelSpec) {
-                continue;
-            }
-            const kernelLanguage = kernelSpec.language?.toLowerCase() || '';
-            // Direct language match
-            if (kernelLanguage === normalizedLanguage) {
-                return kernelName;
-            }
-        }
-        // No matching kernel found, return default
-        console.warn(`No kernel found for language '${language}', using default`);
-        return specs.default || Object.keys(specs.kernelspecs)[0] || 'python3';
-    }
-    catch (error) {
-        console.warn('Failed to find kernel by language:', error);
-        return 'python3';
-    }
-}
-/**
- * Helper function to get a notebook widget by path or use the active one
- */
-async function getNotebookWidget(notebookPath, docManager, notebookTracker) {
-    if (notebookPath) {
-        // Open specific notebook by path using document manager
-        let widget = docManager.findWidget(notebookPath);
-        if (!widget) {
-            widget = docManager.openOrReveal(notebookPath);
-        }
-        if (!(widget instanceof NotebookPanel)) {
-            throw new Error(`Widget for ${notebookPath} is not a notebook panel`);
-        }
-        return widget ?? null;
-    }
-    else {
-        // Use current active notebook
-        return notebookTracker?.currentWidget || null;
-    }
-}
-/**
- * Create a notebook creation tool
- */
-export function createNotebookCreationTool(docManager, kernelSpecManager) {
-    return tool({
-        title: 'Create Notebook',
-        description: 'Create a new Jupyter notebook with a kernel for the specified programming language',
-        inputSchema: z.object({
-            language: z
-                .string()
-                .optional()
-                .nullable()
-                .describe('The programming language for the notebook (e.g., python, r, julia, javascript, etc.). Will use system default if not specified.'),
-            name: z
-                .string()
-                .optional()
-                .nullable()
-                .describe('Optional name for the notebook file (without .ipynb extension)')
-        }),
-        execute: async (input) => {
-            const kernel = await findKernelByLanguage(kernelSpecManager, input.language);
-            const { name } = input;
-            if (!name) {
-                return {
-                    success: false,
-                    error: 'A name must be provided to create a notebook'
-                };
-            }
-            // TODO: handle cwd / path?
-            const fileName = name.endsWith('.ipynb') ? name : `${name}.ipynb`;
-            // Create untitled notebook first
-            const notebookModel = await docManager.newUntitled({
-                type: 'notebook'
-            });
-            // Rename to desired filename
-            await docManager.services.contents.rename(notebookModel.path, fileName);
-            // Create widget with specific kernel
-            const notebook = docManager.createNew(fileName, 'default', {
-                name: kernel
-            });
-            if (!(notebook instanceof DocumentWidget)) {
-                return {
-                    success: false,
-                    error: 'Failed to create notebook widget'
-                };
-            }
-            await notebook.context.ready;
-            await notebook.context.save();
-            docManager.openOrReveal(fileName);
-            return {
-                success: true,
-                message: `Successfully created notebook ${fileName} with ${kernel} kernel${input.language ? ` for ${input.language}` : ''}`,
-                notebookPath: fileName,
-                notebookName: fileName,
-                kernel,
-                language: input.language
-            };
-        }
-    });
-}
-/**
- * Create a tool for adding cells to a specific notebook
- */
-export function createAddCellTool(docManager, notebookTracker) {
-    return tool({
-        title: 'Add Cell',
-        description: 'Add a cell to the current notebook with optional content',
-        inputSchema: z.object({
-            notebookPath: z
-                .string()
-                .optional()
-                .nullable()
-                .describe('Path to the notebook file. If not provided, uses the currently active notebook'),
-            content: z
-                .string()
-                .optional()
-                .nullable()
-                .describe('Content to add to the cell'),
-            cellType: z
-                .enum(['code', 'markdown', 'raw'])
-                .default('code')
-                .describe('Type of cell to add'),
-            position: z
-                .enum(['above', 'below'])
-                .optional()
-                .default('below')
-                .describe('Position relative to current cell')
-        }),
-        execute: async ({ notebookPath, content, cellType = 'code', position = 'below' }) => {
-            const currentWidget = await getNotebookWidget(notebookPath, docManager, notebookTracker);
-            if (!currentWidget) {
-                return {
-                    success: false,
-                    error: notebookPath
-                        ? `Failed to open notebook at path: ${notebookPath}`
-                        : 'No active notebook and no notebook path provided'
-                };
-            }
-            const notebook = currentWidget.content;
-            const model = notebook.model;
-            if (!model) {
-                return {
-                    success: false,
-                    error: 'No notebook model available'
-                };
-            }
-            // Check if we should replace the first empty cell instead of adding
-            const shouldReplaceFirstCell = model.cells.length === 1 &&
-                model.cells.get(0).sharedModel.getSource().trim() === '';
-            if (shouldReplaceFirstCell) {
-                // Replace the first empty cell by removing it and adding new one
-                model.sharedModel.deleteCell(0);
-            }
-            // Create the new cell using shared model
-            const newCellData = {
-                cell_type: cellType,
-                source: content || '',
-                metadata: cellType === 'code' ? { trusted: true } : {}
-            };
-            model.sharedModel.addCell(newCellData);
-            // Execute markdown cells after creation to render them
-            if (cellType === 'markdown' && content) {
-                const cellIndex = model.cells.length - 1;
-                const cellWidget = notebook.widgets[cellIndex];
-                if (cellWidget && cellWidget instanceof MarkdownCell) {
-                    try {
-                        await cellWidget.ready;
-                        cellWidget.rendered = true;
-                    }
-                    catch (error) {
-                        console.warn('Failed to render markdown cell:', error);
-                    }
-                }
-            }
-            return {
-                success: true,
-                message: `${cellType} cell added successfully`,
-                content: content || '',
-                cellType,
-                position
-            };
-        }
-    });
-}
-/**
- * Create a tool for getting notebook information
- */
-export function createGetNotebookInfoTool(docManager, notebookTracker) {
-    return tool({
-        title: 'Get Notebook Info',
-        description: 'Get information about a notebook including number of cells and active cell index',
-        inputSchema: z.object({
-            notebookPath: z
-                .string()
-                .optional()
-                .nullable()
-                .describe('Path to the notebook file. If not provided, uses the currently active notebook')
-        }),
-        execute: async (input) => {
-            const { notebookPath } = input;
-            const currentWidget = await getNotebookWidget(notebookPath, docManager, notebookTracker);
-            if (!currentWidget) {
-                return JSON.stringify({
-                    success: false,
-                    error: notebookPath
-                        ? `Failed to open notebook at path: ${notebookPath}`
-                        : 'No active notebook and no notebook path provided'
-                });
-            }
-            const notebook = currentWidget.content;
-            const model = notebook.model;
-            if (!model) {
-                return JSON.stringify({
-                    success: false,
-                    error: 'No notebook model available'
-                });
-            }
-            const cellCount = model.cells.length;
-            const activeCellIndex = notebook.activeCellIndex;
-            const activeCell = notebook.activeCell;
-            const activeCellType = activeCell?.model.type || 'unknown';
-            return JSON.stringify({
-                success: true,
-                notebookName: currentWidget.title.label,
-                notebookPath: currentWidget.context.path,
-                cellCount,
-                activeCellIndex,
-                activeCellType,
-                isDirty: model.dirty
-            });
-        }
-    });
-}
-/**
- * Create a tool for getting cell information by index
- */
-export function createGetCellInfoTool(docManager, notebookTracker) {
-    return tool({
-        title: 'Get Cell Info',
-        description: 'Get information about a specific cell including its type, source content, and outputs',
-        inputSchema: z.object({
-            notebookPath: z
-                .string()
-                .optional()
-                .nullable()
-                .describe('Path to the notebook file. If not provided, uses the currently active notebook'),
-            cellIndex: z
-                .number()
-                .optional()
-                .nullable()
-                .describe('Index of the cell to get information for (0-based). If not provided, uses the currently active cell')
-        }),
-        execute: async (input) => {
-            const { notebookPath } = input;
-            let { cellIndex } = input;
-            const currentWidget = await getNotebookWidget(notebookPath, docManager, notebookTracker);
-            if (!currentWidget) {
-                return JSON.stringify({
-                    success: false,
-                    error: notebookPath
-                        ? `Failed to open notebook at path: ${notebookPath}`
-                        : 'No active notebook and no notebook path provided'
-                });
-            }
-            const notebook = currentWidget.content;
-            const model = notebook.model;
-            if (!model) {
-                return JSON.stringify({
-                    success: false,
-                    error: 'No notebook model available'
-                });
-            }
-            if (cellIndex === undefined || cellIndex === null) {
-                cellIndex = notebook.activeCellIndex;
-            }
-            if (cellIndex < 0 || cellIndex >= model.cells.length) {
-                return JSON.stringify({
-                    success: false,
-                    error: `Invalid cell index: ${cellIndex}. Notebook has ${model.cells.length} cells.`
-                });
-            }
-            const cell = model.cells.get(cellIndex);
-            const cellType = cell.type;
-            const sharedModel = cell.sharedModel;
-            const source = sharedModel.getSource();
-            // Get outputs for code cells
-            let outputs = [];
-            if (cellType === 'code') {
-                const rawOutputs = sharedModel.toJSON().outputs;
-                outputs = Array.isArray(rawOutputs) ? rawOutputs : [];
-            }
-            return JSON.stringify({
-                success: true,
-                cellId: cell.id,
-                cellIndex,
-                cellType,
-                source,
-                outputs,
-                executionCount: cellType === 'code' ? cell.executionCount : null
-            });
-        }
-    });
-}
-/**
- * Create a tool for setting cell content and type
- */
-export function createSetCellContentTool(docManager, notebookTracker, diffManager) {
-    return tool({
-        title: 'Set Cell Content',
-        description: 'Set the content of a specific cell and return both the previous and new content',
-        inputSchema: z.object({
-            notebookPath: z
-                .string()
-                .optional()
-                .nullable()
-                .describe('Path to the notebook file. If not provided, uses the currently active notebook'),
-            cellId: z
-                .string()
-                .optional()
-                .nullable()
-                .describe('ID of the cell to modify. If provided, takes precedence over cellIndex'),
-            cellIndex: z
-                .number()
-                .optional()
-                .nullable()
-                .describe('Index of the cell to modify (0-based). Used if cellId is not provided. If neither is provided, targets the active cell'),
-            content: z.string().describe('New content for the cell')
-        }),
-        execute: async (input) => {
-            const { notebookPath, cellId, cellIndex, content } = input;
-            const notebookWidget = await getNotebookWidget(notebookPath, docManager, notebookTracker);
-            if (!notebookWidget) {
-                return JSON.stringify({
-                    success: false,
-                    error: notebookPath
-                        ? `Failed to open notebook at path: ${notebookPath}`
-                        : 'No active notebook and no notebook path provided'
-                });
-            }
-            const notebook = notebookWidget.content;
-            const targetNotebookPath = notebookWidget.context.path;
-            const model = notebook.model;
-            if (!model) {
-                return JSON.stringify({
-                    success: false,
-                    error: 'No notebook model available'
-                });
-            }
-            // Determine target cell index
-            let targetCellIndex;
-            if (cellId !== undefined && cellId !== null) {
-                // Find cell by ID
-                targetCellIndex = -1;
-                for (let i = 0; i < model.cells.length; i++) {
-                    if (model.cells.get(i).id === cellId) {
-                        targetCellIndex = i;
-                        break;
-                    }
-                }
-                if (targetCellIndex === -1) {
-                    return JSON.stringify({
-                        success: false,
-                        error: `Cell with ID '${cellId}' not found in notebook`
-                    });
-                }
-            }
-            else if (cellIndex !== undefined && cellIndex !== null) {
-                // Use provided cell index
-                if (cellIndex < 0 || cellIndex >= model.cells.length) {
-                    return JSON.stringify({
-                        success: false,
-                        error: `Invalid cell index: ${cellIndex}. Notebook has ${model.cells.length} cells.`
-                    });
-                }
-                targetCellIndex = cellIndex;
-            }
-            else {
-                // Use active cell
-                targetCellIndex = notebook.activeCellIndex;
-                if (targetCellIndex === -1 || targetCellIndex >= model.cells.length) {
-                    return JSON.stringify({
-                        success: false,
-                        error: 'No active cell or invalid active cell index'
-                    });
-                }
-            }
-            // Get the target cell
-            const targetCell = model.cells.get(targetCellIndex);
-            if (!targetCell) {
-                return JSON.stringify({
-                    success: false,
-                    error: `Cell at index ${targetCellIndex} not found`
-                });
-            }
-            const sharedModel = targetCell.sharedModel;
-            // Get previous content and type
-            const previousContent = sharedModel.getSource();
-            const previousCellType = targetCell.type;
-            const retrievedCellId = targetCell.id;
-            sharedModel.setSource(content);
-            // Show the cell diff using the diff manager if available
-            if (diffManager) {
-                await diffManager.showCellDiff({
-                    original: previousContent,
-                    modified: content,
-                    cellId: retrievedCellId,
-                    notebookPath: targetNotebookPath
-                });
-            }
-            return JSON.stringify({
-                success: true,
-                message: cellId !== undefined && cellId !== null
-                    ? `Cell with ID '${cellId}' content replaced successfully`
-                    : cellIndex !== undefined && cellIndex !== null
-                        ? `Cell ${targetCellIndex} content replaced successfully`
-                        : 'Active cell content replaced successfully',
-                notebookPath: targetNotebookPath,
-                cellId: retrievedCellId,
-                cellIndex: targetCellIndex,
-                previousContent,
-                previousCellType,
-                newContent: content,
-                wasActiveCell: cellId === undefined && cellIndex === undefined
-            });
-        }
-    });
-}
-/**
- * Create a tool for running a specific cell
- */
-export function createRunCellTool(docManager, notebookTracker) {
-    return tool({
-        title: 'Run Cell',
-        description: 'Run a specific cell in the notebook by index',
-        inputSchema: z.object({
-            notebookPath: z
-                .string()
-                .optional()
-                .nullable()
-                .describe('Path to the notebook file. If not provided, uses the currently active notebook'),
-            cellIndex: z.number().describe('Index of the cell to run (0-based)'),
-            recordTiming: z
-                .boolean()
-                .default(true)
-                .describe('Whether to record execution timing')
-        }),
-        needsApproval: true,
-        execute: async (input) => {
-            const { notebookPath, cellIndex, recordTiming = true } = input;
-            const currentWidget = await getNotebookWidget(notebookPath, docManager, notebookTracker);
-            if (!currentWidget) {
-                return JSON.stringify({
-                    success: false,
-                    error: notebookPath
-                        ? `Failed to open notebook at path: ${notebookPath}`
-                        : 'No active notebook and no notebook path provided'
-                });
-            }
-            const notebook = currentWidget.content;
-            const model = notebook.model;
-            if (!model) {
-                return JSON.stringify({
-                    success: false,
-                    error: 'No notebook model available'
-                });
-            }
-            if (cellIndex < 0 || cellIndex >= model.cells.length) {
-                return JSON.stringify({
-                    success: false,
-                    error: `Invalid cell index: ${cellIndex}. Notebook has ${model.cells.length} cells.`
-                });
-            }
-            // Get the target cell widget
-            const cellWidget = notebook.widgets[cellIndex];
-            if (!cellWidget) {
-                return JSON.stringify({
-                    success: false,
-                    error: `Cell widget at index ${cellIndex} not found`
-                });
-            }
-            // Execute using shared model approach (non-disruptive)
-            if (cellWidget instanceof CodeCell) {
-                // Use direct CodeCell.execute() method
-                const sessionCtx = currentWidget.sessionContext;
-                await CodeCell.execute(cellWidget, sessionCtx, {
-                    recordTiming,
-                    deletedCells: model.deletedCells
-                });
-                const codeModel = cellWidget.model;
-                return JSON.stringify({
-                    success: true,
-                    message: `Cell ${cellIndex} executed successfully`,
-                    cellIndex,
-                    executionCount: codeModel.executionCount,
-                    hasOutput: codeModel.outputs.length > 0
-                });
-            }
-            else {
-                // For non-code cells, just return success
-                return JSON.stringify({
-                    success: true,
-                    message: `Cell ${cellIndex} is not a code cell, no execution needed`,
-                    cellIndex,
-                    cellType: cellWidget.model.type
-                });
-            }
-        }
-    });
-}
-/**
- * Create a tool for deleting a specific cell
- */
-export function createDeleteCellTool(docManager, notebookTracker) {
-    return tool({
-        title: 'Delete Cell',
-        description: 'Delete a specific cell from the notebook by index',
-        inputSchema: z.object({
-            notebookPath: z
-                .string()
-                .optional()
-                .nullable()
-                .describe('Path to the notebook file. If not provided, uses the currently active notebook'),
-            cellIndex: z.number().describe('Index of the cell to delete (0-based)')
-        }),
-        execute: async (input) => {
-            const { notebookPath, cellIndex } = input;
-            const currentWidget = await getNotebookWidget(notebookPath, docManager, notebookTracker);
-            if (!currentWidget) {
-                return JSON.stringify({
-                    success: false,
-                    error: notebookPath
-                        ? `Failed to open notebook at path: ${notebookPath}`
-                        : 'No active notebook and no notebook path provided'
-                });
-            }
-            const notebook = currentWidget.content;
-            const model = notebook.model;
-            if (!model) {
-                return JSON.stringify({
-                    success: false,
-                    error: 'No notebook model available'
-                });
-            }
-            if (cellIndex < 0 || cellIndex >= model.cells.length) {
-                return JSON.stringify({
-                    success: false,
-                    error: `Invalid cell index: ${cellIndex}. Notebook has ${model.cells.length} cells.`
-                });
-            }
-            // Validate cell exists
-            const targetCell = model.cells.get(cellIndex);
-            if (!targetCell) {
-                return JSON.stringify({
-                    success: false,
-                    error: `Cell at index ${cellIndex} not found`
-                });
-            }
-            // Delete cell using shared model (non-disruptive)
-            model.sharedModel.deleteCell(cellIndex);
-            return JSON.stringify({
-                success: true,
-                message: `Cell ${cellIndex} deleted successfully`,
-                cellIndex,
-                remainingCells: model.cells.length
-            });
-        }
-    });
-}
-/**
- * Create a tool for executing code in the active cell (non-disruptive alternative to mcp__ide__executeCode)
- */
-export function createExecuteActiveCellTool(docManager, notebookTracker) {
-    return tool({
-        title: 'Execute Active Cell',
-        description: 'Execute the currently active cell in the notebook without disrupting user focus',
-        inputSchema: z.object({
-            notebookPath: z
-                .string()
-                .optional()
-                .nullable()
-                .describe('Path to the notebook file. If not provided, uses the currently active notebook'),
-            code: z
-                .string()
-                .optional()
-                .nullable()
-                .describe('Optional: set cell content before executing'),
-            recordTiming: z
-                .boolean()
-                .default(true)
-                .describe('Whether to record execution timing')
-        }),
-        execute: async (input) => {
-            const { notebookPath, code, recordTiming = true } = input;
-            const currentWidget = await getNotebookWidget(notebookPath, docManager, notebookTracker);
-            if (!currentWidget) {
-                return JSON.stringify({
-                    success: false,
-                    error: notebookPath
-                        ? `Failed to open notebook at path: ${notebookPath}`
-                        : 'No active notebook and no notebook path provided'
-                });
-            }
-            const notebook = currentWidget.content;
-            const model = notebook.model;
-            const activeCellIndex = notebook.activeCellIndex;
-            if (!model || activeCellIndex === -1) {
-                return JSON.stringify({
-                    success: false,
-                    error: 'No notebook model or active cell available'
-                });
-            }
-            const activeCell = model.cells.get(activeCellIndex);
-            if (!activeCell) {
-                return JSON.stringify({
-                    success: false,
-                    error: 'Active cell not found'
-                });
-            }
-            // Set code content if provided
-            if (code) {
-                activeCell.sharedModel.setSource(code);
-            }
-            // Get the cell widget for execution
-            const cellWidget = notebook.widgets[activeCellIndex];
-            if (!cellWidget || !(cellWidget instanceof CodeCell)) {
-                return JSON.stringify({
-                    success: false,
-                    error: 'Active cell is not a code cell'
-                });
-            }
-            // Execute using shared model approach (non-disruptive)
-            const sessionCtx = currentWidget.sessionContext;
-            await CodeCell.execute(cellWidget, sessionCtx, {
-                recordTiming,
-                deletedCells: model.deletedCells
-            });
-            const codeModel = cellWidget.model;
-            return JSON.stringify({
-                success: true,
-                message: 'Code executed successfully in active cell',
-                cellIndex: activeCellIndex,
-                executionCount: codeModel.executionCount,
-                hasOutput: codeModel.outputs.length > 0,
-                code: code || activeCell.sharedModel.getSource()
-            });
-        }
-    });
-}
-/**
- * Create a tool for saving a specific notebook
- */
-export function createSaveNotebookTool(docManager, notebookTracker) {
-    return tool({
-        title: 'Save Notebook',
-        description: 'Save a specific notebook to disk',
-        inputSchema: z.object({
-            notebookPath: z
-                .string()
-                .optional()
-                .nullable()
-                .describe('Path to the notebook file. If not provided, uses the currently active notebook')
-        }),
-        execute: async (input) => {
-            const { notebookPath } = input;
-            const currentWidget = await getNotebookWidget(notebookPath, docManager, notebookTracker);
-            if (!currentWidget) {
-                return JSON.stringify({
-                    success: false,
-                    error: notebookPath
-                        ? `Failed to open notebook at path: ${notebookPath}`
-                        : 'No active notebook and no notebook path provided'
-                });
-            }
-            await currentWidget.context.save();
-            return JSON.stringify({
-                success: true,
-                message: 'Notebook saved successfully',
-                notebookName: currentWidget.title.label,
-                notebookPath: currentWidget.context.path
-            });
-        }
-    });
-}