@jupyterlite/ai 0.9.0-a3 → 0.9.0

package/src/tools/file.ts

@@ -1,11 +1,14 @@
+import { PathExt } from '@jupyterlab/coreutils';
 import { CommandRegistry } from '@lumino/commands';
 import { IDocumentManager } from '@jupyterlab/docmanager';
+import { IDocumentWidget } from '@jupyterlab/docregistry';
+import { IEditorTracker } from '@jupyterlab/fileeditor';
 
 import { tool } from '@openai/agents';
 
 import { z } from 'zod';
 
-import { ITool } from '../tokens';
+import { IDiffManager, ITool } from '../tokens';
 
 /**
  * Create a tool for creating new files of various types
@@ -18,16 +21,11 @@ export function createNewFileTool(docManager: IDocumentManager): ITool {
     parameters: z.object({
       fileName: z.string().describe('Name of the file to create'),
       fileType: z
-        .enum([
-          'text',
-          'python',
-          'markdown',
-          'json',
-          'javascript',
-          'typescript'
-        ])
+        .string()
         .default('text')
-        .describe('Type of file to create'),
+        .describe(
+          'Type of file to create. Common examples: text, python, markdown, json, javascript, typescript, yaml, julia, r, csv'
+        ),
       content: z
         .string()
         .optional()
@@ -39,96 +37,66 @@ export function createNewFileTool(docManager: IDocumentManager): ITool {
         .nullable()
         .describe('Directory where to create the file (optional)')
     }),
+    errorFunction: (context, error) => {
+      return JSON.stringify({
+        success: false,
+        error: `Failed to create file: ${error instanceof Error ? error.message : String(error)}`
+      });
+    },
     execute: async (input: {
       fileName: string;
-      fileType?:
-        | 'text'
-        | 'python'
-        | 'markdown'
-        | 'json'
-        | 'javascript'
-        | 'typescript';
+      fileType?: string;
       content?: string | null;
       cwd?: string | null;
     }) => {
       const { fileName, content = '', cwd, fileType = 'text' } = input;
 
-      try {
-        // Determine file extension based on type
-        const extensions: Record<string, string> = {
-          python: 'py',
-          markdown: 'md',
-          json: 'json',
-          text: 'txt',
-          javascript: 'js',
-          typescript: 'ts'
-        };
-
-        const ext = extensions[fileType] || 'txt';
-
-        // If fileName already has an extension, use it as-is, otherwise add the extension
-        const fullFileName = fileName.includes('.')
-          ? fileName
-          : `${fileName}.${ext}`;
-
-        // For Python files, ensure .py extension if fileType is python
-        const finalFileName =
-          fileType === 'python' &&
-          !fileName.endsWith('.py') &&
-          !fileName.includes('.')
-            ? `${fileName}.py`
-            : fullFileName;
-
-        const fullPath = cwd ? `${cwd}/${finalFileName}` : finalFileName;
-
-        // Create file with content using document manager
-        const model = await docManager.services.contents.newUntitled({
-          path: cwd || '',
-          type: 'file',
-          ext
-        });
+      const registeredFileType = docManager.registry.getFileType(fileType);
+      const ext = registeredFileType?.extensions[0] || '.txt';
 
-        // Rename to desired name if needed
-        let finalPath = model.path;
-        if (model.name !== finalFileName) {
-          const renamed = await docManager.services.contents.rename(
-            model.path,
-            fullPath
-          );
-          finalPath = renamed.path;
-        }
+      const existingExt = PathExt.extname(fileName);
+      const fullFileName = existingExt ? fileName : `${fileName}${ext}`;
 
-        // Set content if provided
-        if (content) {
-          await docManager.services.contents.save(finalPath, {
-            type: 'file',
-            format: 'text',
-            content
-          });
-        }
+      const fullPath = cwd ? `${cwd}/${fullFileName}` : fullFileName;
 
-        // Open the newly created file
-        let opened = false;
-        if (!docManager.findWidget(finalPath)) {
-          docManager.openOrReveal(finalPath);
-          opened = true;
-        }
+      const model = await docManager.services.contents.newUntitled({
+        path: cwd || '',
+        type: 'file',
+        ext
+      });
+
+      let finalPath = model.path;
+      if (model.name !== fullFileName) {
+        const renamed = await docManager.services.contents.rename(
+          model.path,
+          fullPath
+        );
+        finalPath = renamed.path;
+      }
 
-        return {
-          success: true,
-          message: `${fileType} file '${finalFileName}' created and opened successfully`,
-          fileName: finalFileName,
-          filePath: finalPath,
-          fileType,
-          hasContent: !!content,
-          opened
-        };
-      } catch (error) {
-        return {
-          success: false,
-          error: `Failed to create file: ${(error as Error).message}`
-        };
+      if (content) {
+        await docManager.services.contents.save(finalPath, {
+          type: 'file',
+          format: 'text',
+          content
+        });
       }
+
+      let opened = false;
+      if (!docManager.findWidget(finalPath)) {
+        docManager.openOrReveal(finalPath);
+        opened = true;
+      }
+
+      return {
+        success: true,
+        message: `${fileType} file '${fullFileName}' created and opened successfully`,
+        fileName: fullFileName,
+        filePath: finalPath,
+        fileType,
+        hasContent: !!content,
+        opened
+      };
     }
   });
 }
@@ -143,31 +111,27 @@ export function createOpenFileTool(docManager: IDocumentManager): ITool {
     parameters: z.object({
       filePath: z.string().describe('Path to the file to open')
     }),
+    errorFunction: (context, error) => {
+      return JSON.stringify({
+        success: false,
+        error: `Failed to open file: ${error instanceof Error ? error.message : String(error)}`
+      });
+    },
     execute: async (input: { filePath: string }) => {
       const { filePath } = input;
 
-      try {
-        const widget = docManager.openOrReveal(filePath);
-
-        if (!widget) {
-          return {
-            success: false,
-            error: `Failed to open file: ${filePath}`
-          };
-        }
+      const widget = docManager.openOrReveal(filePath);
 
-        return {
-          success: true,
-          message: `File '${filePath}' opened successfully`,
-          filePath,
-          widgetId: widget.id
-        };
-      } catch (error) {
-        return {
-          success: false,
-          error: `Failed to open file: ${(error as Error).message}`
-        };
+      if (!widget) {
+        throw new Error(`Could not open file: ${filePath}`);
       }
+
+      return {
+        success: true,
+        message: `File '${filePath}' opened successfully`,
+        filePath,
+        widgetId: widget.id
+      };
     }
   });
 }
@@ -182,23 +146,22 @@ export function createDeleteFileTool(docManager: IDocumentManager): ITool {
     parameters: z.object({
       filePath: z.string().describe('Path to the file to delete')
     }),
+    errorFunction: (context, error) => {
+      return JSON.stringify({
+        success: false,
+        error: `Failed to delete file: ${error instanceof Error ? error.message : String(error)}`
+      });
+    },
     execute: async (input: { filePath: string }) => {
       const { filePath } = input;
 
-      try {
-        await docManager.services.contents.delete(filePath);
+      await docManager.services.contents.delete(filePath);
 
-        return {
-          success: true,
-          message: `File '${filePath}' deleted successfully`,
-          filePath
-        };
-      } catch (error) {
-        return {
-          success: false,
-          error: `Failed to delete file: ${(error as Error).message}`
-        };
-      }
+      return {
+        success: true,
+        message: `File '${filePath}' deleted successfully`,
+        filePath
+      };
     }
   });
 }
@@ -214,24 +177,23 @@ export function createRenameFileTool(docManager: IDocumentManager): ITool {
       oldPath: z.string().describe('Current path of the file'),
       newPath: z.string().describe('New path/name for the file')
     }),
+    errorFunction: (context, error) => {
+      return JSON.stringify({
+        success: false,
+        error: `Failed to rename file: ${error instanceof Error ? error.message : String(error)}`
+      });
+    },
     execute: async (input: { oldPath: string; newPath: string }) => {
       const { oldPath, newPath } = input;
 
-      try {
-        await docManager.services.contents.rename(oldPath, newPath);
-
-        return {
-          success: true,
-          message: `File renamed from '${oldPath}' to '${newPath}' successfully`,
-          oldPath,
-          newPath
-        };
-      } catch (error) {
-        return {
-          success: false,
-          error: `Failed to rename file: ${(error as Error).message}`
-        };
-      }
+      await docManager.services.contents.rename(oldPath, newPath);
+
+      return {
+        success: true,
+        message: `File renamed from '${oldPath}' to '${newPath}' successfully`,
+        oldPath,
+        newPath
+      };
     }
   });
 }
@@ -249,24 +211,23 @@ export function createCopyFileTool(docManager: IDocumentManager): ITool {
         .string()
         .describe('Destination path for the copied file')
     }),
+    errorFunction: (context, error) => {
+      return JSON.stringify({
+        success: false,
+        error: `Failed to copy file: ${error instanceof Error ? error.message : String(error)}`
+      });
+    },
     execute: async (input: { sourcePath: string; destinationPath: string }) => {
       const { sourcePath, destinationPath } = input;
 
-      try {
-        await docManager.services.contents.copy(sourcePath, destinationPath);
-
-        return {
-          success: true,
-          message: `File copied from '${sourcePath}' to '${destinationPath}' successfully`,
-          sourcePath,
-          destinationPath
-        };
-      } catch (error) {
-        return {
-          success: false,
-          error: `Failed to copy file: ${(error as Error).message}`
-        };
-      }
+      await docManager.services.contents.copy(sourcePath, destinationPath);
+
+      return {
+        success: true,
+        message: `File copied from '${sourcePath}' to '${destinationPath}' successfully`,
+        sourcePath,
+        destinationPath
+      };
     }
   });
 }
@@ -283,25 +244,195 @@ export function createNavigateToDirectoryTool(
     parameters: z.object({
       directoryPath: z.string().describe('Path to the directory to navigate to')
     }),
+    errorFunction: (context, error) => {
+      return JSON.stringify({
+        success: false,
+        error: `Failed to navigate to directory: ${error instanceof Error ? error.message : String(error)}`
+      });
+    },
     execute: async (input: { directoryPath: string }) => {
       const { directoryPath } = input;
 
-      try {
-        await commands.execute('filebrowser:go-to-path', {
-          path: directoryPath
+      await commands.execute('filebrowser:go-to-path', {
+        path: directoryPath
+      });
+
+      return {
+        success: true,
+        message: `Navigated to directory '${directoryPath}' successfully`,
+        directoryPath
+      };
+    }
+  });
+}
+
+/**
+ * Create a tool for getting file information and content
+ */
+export function createGetFileInfoTool(
+  docManager: IDocumentManager,
+  editorTracker?: IEditorTracker
+): ITool {
+  return tool({
+    name: 'get_file_info',
+    description:
+      'Get information about a file including its path, name, extension, and content. Works with text-based files like Python files, markdown, JSON, etc. For Jupyter notebooks, use dedicated notebook tools instead. If no file path is provided, returns information about the currently active file in the editor.',
+    parameters: z.object({
+      filePath: z
+        .string()
+        .optional()
+        .nullable()
+        .describe(
+          'Path to the file to read (e.g., "script.py", "README.md", "config.json"). If not provided, uses the currently active file in the editor.'
+        )
+    }),
+    errorFunction: (context, error) => {
+      return JSON.stringify({
+        success: false,
+        error: `Failed to get file info: ${error instanceof Error ? error.message : String(error)}`
+      });
+    },
+    execute: async (input: { filePath?: string | null }) => {
+      const { filePath } = input;
+
+      let widget: IDocumentWidget | null = null;
+
+      if (filePath) {
+        widget =
+          docManager.findWidget(filePath) ??
+          docManager.openOrReveal(filePath) ??
+          null;
+
+        if (!widget) {
+          throw new Error(`Failed to open file at path: ${filePath}`);
+        }
+      } else {
+        widget = editorTracker?.currentWidget ?? null;
+
+        if (!widget) {
+          throw new Error(
+            'No active file in the editor and no file path provided'
+          );
+        }
+      }
+
+      if (!widget.context) {
+        throw new Error('Widget is not a document');
+      }
+
+      await widget.context.ready;
+
+      const model = widget.context.model;
+
+      if (!model) {
+        throw new Error('File model not available');
+      }
+
+      const sharedModel = model.sharedModel;
+      const content = sharedModel.getSource();
+      const resolvedFilePath = widget.context.path;
+      const fileName = widget.title.label;
+      const fileExtension = PathExt.extname(resolvedFilePath) || 'unknown';
+
+      return JSON.stringify({
+        success: true,
+        filePath: resolvedFilePath,
+        fileName,
+        fileExtension,
+        content,
+        isDirty: model.dirty,
+        readOnly: model.readOnly,
+        widgetType: widget.constructor.name
+      });
+    }
+  });
+}
+
+/**
+ * Create a tool for setting the content of a file
+ */
+export function createSetFileContentTool(
+  docManager: IDocumentManager,
+  diffManager?: IDiffManager
+): ITool {
+  return tool({
+    name: 'set_file_content',
+    description:
+      'Set or update the content of an existing file. This will replace the entire content of the file. For Jupyter notebooks, use dedicated notebook tools instead.',
+    parameters: z.object({
+      filePath: z
+        .string()
+        .describe(
+          'Path to the file to update (e.g., "script.py", "README.md", "config.json")'
+        ),
+      content: z.string().describe('The new content to set for the file'),
+      save: z
+        .boolean()
+        .optional()
+        .default(true)
+        .describe('Whether to save the file after updating (default: true)')
+    }),
+    errorFunction: (context, error) => {
+      return JSON.stringify({
+        success: false,
+        error: `Failed to set file content: ${error instanceof Error ? error.message : String(error)}`
+      });
+    },
+    execute: async (input: {
+      filePath: string;
+      content: string;
+      save?: boolean;
+    }) => {
+      const { filePath, content, save = true } = input;
+
+      let widget = docManager.findWidget(filePath);
+
+      if (!widget) {
+        widget = docManager.openOrReveal(filePath);
+      }
+
+      if (!widget) {
+        throw new Error(`Failed to open file at path: ${filePath}`);
+      }
+
+      await widget.context.ready;
+
+      const model = widget.context.model;
+
+      if (!model) {
+        throw new Error('File model not available');
+      }
+
+      if (model.readOnly) {
+        throw new Error('File is read-only and cannot be modified');
+      }
+
+      const sharedModel = model.sharedModel;
+      const originalContent = sharedModel.getSource();
+
+      sharedModel.setSource(content);
+
+      // Show the file diff using the diff manager if available
+      if (diffManager) {
+        await diffManager.showFileDiff({
+          original: String(originalContent),
+          modified: content,
+          filePath
         });
+      }
 
-        return {
-          success: true,
-          message: `Navigated to directory '${directoryPath}' successfully`,
-          directoryPath
-        };
-      } catch (error) {
-        return {
-          success: false,
-          error: `Failed to navigate to directory: ${(error as Error).message}`
-        };
+      if (save) {
+        await widget.context.save();
       }
+
+      return JSON.stringify({
+        success: true,
+        filePath,
+        fileName: widget.title.label,
+        contentLength: content.length,
+        saved: save,
+        isDirty: model.dirty
+      });
     }
   });
 }

package/src/tools/notebook.ts

@@ -3,13 +3,12 @@ import { IDocumentManager } from '@jupyterlab/docmanager';
 import { DocumentWidget } from '@jupyterlab/docregistry';
 import { INotebookTracker, NotebookPanel } from '@jupyterlab/notebook';
 import { KernelSpec } from '@jupyterlab/services';
-import { CommandRegistry } from '@lumino/commands';
 
 import { tool } from '@openai/agents';
 
 import { z } from 'zod';
 
-import { ITool } from '../tokens';
+import { IDiffManager, ITool } from '../tokens';
 
 /**
  * Find a kernel name that matches the specified language
@@ -198,6 +197,7 @@ export function createAddCellTool(
         .describe('Type of cell to add'),
       position: z
         .enum(['above', 'below'])
+        .optional()
         .default('below')
         .describe('Position relative to current cell')
     }),
@@ -460,8 +460,8 @@ export function createGetCellInfoTool(
  */
 export function createSetCellContentTool(
   docManager: IDocumentManager,
-  commands: CommandRegistry,
-  notebookTracker?: INotebookTracker
+  notebookTracker?: INotebookTracker,
+  diffManager?: IDiffManager
 ): ITool {
   return tool({
     name: 'set_cell_content',
@@ -581,16 +581,15 @@ export function createSetCellContentTool(
 
         sharedModel.setSource(content);
 
-        // Show the cell diff using jupyterlab-cell-diff if available
-        const showDiffCommandId = 'jupyterlab-cell-diff:show-codemirror';
-        void commands.execute(showDiffCommandId, {
-          originalSource: previousContent,
-          newSource: content,
-          cellId: retrievedCellId,
-          showActionButtons: true,
-          openDiff: true,
-          notebookPath: targetNotebookPath
-        });
+        // 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,