@jupyterlite/ai 0.8.0 → 0.9.0-a0

package/src/tokens.ts

@@ -1,158 +1,271 @@
-import { BaseChatModel } from '@langchain/core/language_models/chat_models';
-import { ReadonlyPartialJSONObject, Token } from '@lumino/coreutils';
+import { Token } from '@lumino/coreutils';
 import { ISignal } from '@lumino/signaling';
-import { JSONSchema7 } from 'json-schema';
+import { FunctionTool } from '@openai/agents';
+import { LanguageModel } from 'ai';
+import type { AISettingsModel } from './models/settings-model';
+import type { IModelOptions } from './providers/models';
 
-import { IBaseCompleter } from './base-completer';
-import { AIChatModel, AICompleter } from './types/ai-model';
+/**
+ * Type definition for a tool
+ */
+export type ITool = FunctionTool<any, any, any>;
 
-export const PLUGIN_IDS = {
-  chat: '@jupyterlite/ai:chat',
-  chatCommandRegistry: '@jupyterlite/ai:autocompletion-registry',
-  completer: '@jupyterlite/ai:completer',
-  providerRegistry: '@jupyterlite/ai:provider-registry',
-  settingsConnector: '@jupyterlite/ai:settings-connector',
-  systemPrompts: '@jupyterlite/ai:system-prompts'
-};
+/**
+ * Interface for token usage statistics from AI model interactions
+ */
+export interface ITokenUsage {
+  /**
+   * Number of input tokens consumed (prompt tokens)
+   */
+  inputTokens: number;
 
-export type ModelRole = 'chat' | 'completer';
+  /**
+   * Number of output tokens generated (completion tokens)
+   */
+  outputTokens: number;
+}
 
-export interface IDict<T = any> {
-  [key: string]: T;
+/**
+ * Interface for a named tool (tool with a name identifier)
+ */
+export interface INamedTool {
+  /**
+   * The unique name of the tool
+   */
+  name: string;
+  /**
+   * The tool instance
+   */
+  tool: ITool;
 }
 
-export interface IType<T> {
-  new (...args: any[]): T;
+/**
+ * The tool registry interface for managing AI tools
+ */
+export interface IToolRegistry {
+  /**
+   * The registered tools as a record (name -> tool mapping).
+   */
+  readonly tools: Record<string, ITool>;
+
+  /**
+   * The registered named tools array.
+   */
+  readonly namedTools: INamedTool[];
+
+  /**
+   * A signal triggered when the tools have changed.
+   */
+  readonly toolsChanged: ISignal<IToolRegistry, void>;
+
+  /**
+   * Add a new tool to the registry.
+   */
+  add(name: string, tool: ITool): void;
+
+  /**
+   * Get a tool for a given name.
+   * Return null if the name is not provided or if there is no registered tool with the
+   * given name.
+   */
+  get(name: string | null): ITool | null;
+
+  /**
+   * Remove a tool from the registry by name.
+   */
+  remove(name: string): boolean;
 }
 
 /**
- * The provider interface.
+ * The tool registry token.
+ */
+export const IToolRegistry = new Token<IToolRegistry>(
+  '@jupyterlite/ai:tool-registry',
+  'Tool registry for AI agent functionality'
+);
+
+/**
+ * Token for the chat provider registry.
  */
-export interface IAIProvider {
+export const IChatProviderRegistry = new Token<IChatProviderRegistry>(
+  '@jupyterlite/ai:chat-provider-registry',
+  'Registry for chat AI providers'
+);
+
+/**
+ * Token for the completion provider registry.
+ */
+export const ICompletionProviderRegistry =
+  new Token<ICompletionProviderRegistry>(
+    '@jupyterlite/ai:completion-provider-registry',
+    'Registry for completion providers'
+  );
+
+/**
+ * Interface for a provider factory function that creates chat models
+ */
+export interface IChatProviderFactory {
+  (options: IModelOptions): any; // Returns the model instance for @openai/agents
+}
+
+/**
+ * Interface for a provider factory function that creates completion models
+ */
+export interface ICompletionProviderFactory {
+  (options: IModelOptions): LanguageModel;
+}
+
+/**
+ * Base information about a registered provider
+ */
+export interface IBaseProviderInfo {
   /**
-   * The name of the provider.
+   * Unique identifier for the provider
    */
-  name: string;
+  id: string;
+
   /**
-   * The chat model class to use.
+   * Display name for the provider
    */
-  chat?: IType<BaseChatModel>;
+  name: string;
+
   /**
-   * The completer class to use.
+   * Whether this provider requires an API key
    */
-  completer?: IType<IBaseCompleter>;
+  requiresApiKey: boolean;
+
   /**
-   * the settings schema for the provider.
+   * Default model names for this provider
    */
-  settingsSchema?: any;
+  defaultModels: string[];
+
   /**
-   * The instructions to be displayed in the settings, as helper to use the provider.
-   * A markdown renderer is used to render the instructions.
+   * Whether this provider supports custom base URLs
    */
-  instructions?: string;
+  supportsBaseURL?: boolean;
+
   /**
-   * A function that extract the error message from the provider API error.
-   * Default to `(error) => error.message`.
+   * Whether this provider supports custom headers
    */
-  errorMessage?: (error: any) => string;
+  supportsHeaders?: boolean;
+
   /**
-   * Compatibility check function, to determine if the provider is compatible with the
-   * current environment.
+   * Whether this provider supports tool calling
    */
-  compatibilityCheck?: () => Promise<string | null>;
+  supportsToolCalling?: boolean;
+
   /**
-   * Whether to expose or not the chat model.
-   *
-   * ### CAUTION
-   * This flag will expose the whole chat model API, which may contain private keys.
-   * Be sure to use it with a model that does not expose sensitive information in the
-   * API.
+   * Additional provider-specific configuration schema
    */
-  exposeChatModel?: boolean;
+  customSettings?: Record<string, any>;
 }
 
 /**
- * The provider registry interface.
+ * Information about a chat provider
  */
-export interface IAIProviderRegistry {
+export interface IChatProviderInfo extends IBaseProviderInfo {
   /**
-   * Get the list of provider names.
+   * Factory function for creating chat models
    */
-  readonly providers: string[];
+  factory: IChatProviderFactory;
+}
+
+/**
+ * Information about a completion provider
+ */
+export interface ICompletionProviderInfo extends IBaseProviderInfo {
   /**
-   * Add a new provider.
+   * Factory function for creating completion models
    */
-  add(provider: IAIProvider): void;
+  factory: ICompletionProviderFactory;
+}
+
+/**
+ * Registry for chat AI providers
+ */
+export interface IChatProviderRegistry {
   /**
-   * Get the current provider name.
+   * The registered providers as a record (id -> info mapping).
    */
-  currentName(role: ModelRole): string;
+  readonly providers: Record<string, IChatProviderInfo>;
+
   /**
-   * Get the current completer of the completion provider.
+   * A signal triggered when providers have changed.
    */
-  currentCompleter: AICompleter | null;
+  readonly providersChanged: ISignal<IChatProviderRegistry, void>;
+
   /**
-   * Getter/setter for the completer system prompt.
+   * Register a new chat provider.
    */
-  completerSystemPrompt: string;
+  registerProvider(info: IChatProviderInfo): void;
+
   /**
-   * Get the current llm chat model.
+   * Unregister a chat provider.
    */
-  currentChatModel: AIChatModel | null;
+  unregisterProvider(id: string): boolean;
+
   /**
-   * Getter/setter for the chat system prompt.
+   * Get provider info by id.
    */
-  chatSystemPrompt: string;
+  getProviderInfo(id: string): IChatProviderInfo | null;
+
   /**
-   * Get the settings schema of a given provider.
+   * Create a chat model instance for the given provider.
    */
-  getSettingsSchema(provider: string): JSONSchema7;
+  createChatModel(id: string, options: IModelOptions): any | null;
+
   /**
-   * Get the instructions of a given provider.
+   * Get all available provider IDs.
    */
-  getInstructions(provider: string): string | undefined;
+  getAvailableProviders(): string[];
+}
+
+/**
+ * Registry for completion providers
+ */
+export interface ICompletionProviderRegistry {
   /**
-   * Get the compatibility check function of a given provider.
+   * The registered providers as a record (id -> info mapping).
    */
-  getCompatibilityCheck(
-    provider: string
-  ): (() => Promise<string | null>) | undefined;
+  readonly providers: Record<string, ICompletionProviderInfo>;
+
   /**
-   * Format an error message from the current provider.
+   * A signal triggered when providers have changed.
    */
-  formatErrorMessage(error: any): string;
+  readonly providersChanged: ISignal<ICompletionProviderRegistry, void>;
+
   /**
-   * Set the completer provider.
-   * Creates the provider if the name has changed, otherwise only updates its config.
-   *
-   * @param options - An object with the name and the settings of the provider to use.
+   * Register a new completion provider.
    */
-  setCompleterProvider(settings: ReadonlyPartialJSONObject): void;
+  registerProvider(info: ICompletionProviderInfo): void;
+
   /**
-   * Set the chat provider.
-   * Creates the provider if the name has changed, otherwise only updates its config.
-   *
-   * @param options - An object with the name and the settings of the provider to use.
+   * Unregister a completion provider.
    */
-  setChatProvider(settings: ReadonlyPartialJSONObject): void;
+  unregisterProvider(id: string): boolean;
+
   /**
-   * A signal emitting when the provider or its settings has changed.
+   * Get provider info by id.
    */
-  readonly providerChanged: ISignal<IAIProviderRegistry, ModelRole>;
+  getProviderInfo(id: string): ICompletionProviderInfo | null;
+
   /**
-   * Get the current chat error;
+   * Create a completion model instance for the given provider.
    */
-  readonly chatError: string;
+  createCompletionModel(
+    id: string,
+    options: IModelOptions
+  ): LanguageModel | null;
+
   /**
-   * get the current completer error.
+   * Get all available provider IDs.
    */
-  readonly completerError: string;
+  getAvailableProviders(): string[];
 }
 
 /**
- * The provider registry token.
+ * Token for the AI settings model.
  */
-export const IAIProviderRegistry = new Token<IAIProviderRegistry>(
-  '@jupyterlite/ai:provider-registry',
-  'Provider for chat and completion LLM provider'
+export const IAISettingsModel = new Token<AISettingsModel>(
+  '@jupyterlite/ai:IAISettingsModel'
 );

package/src/tools/commands.ts

@@ -0,0 +1,151 @@
+import { CommandRegistry } from '@lumino/commands';
+import { tool } from '@openai/agents';
+import { z } from 'zod';
+import { ITool } from '../tokens';
+import { AISettingsModel } from '../models/settings-model';
+
+/**
+ * Create a tool to discover all available commands and their metadata
+ */
+export function createDiscoverCommandsTool(commands: CommandRegistry): ITool {
+  return tool({
+    name: 'discover_commands',
+    description:
+      'Discover all available JupyterLab commands with their metadata, arguments, and descriptions',
+    parameters: z.object({
+      // currently unused, but could be used to filter commands by a search term
+      query: z
+        .string()
+        .optional()
+        .nullable()
+        .describe('Optional search query to filter commands')
+    }),
+    execute: async () => {
+      try {
+        const commandList: Array<{
+          id: string;
+          label?: string;
+          caption?: string;
+          description?: string;
+          args?: any;
+        }> = [];
+
+        // Get all command IDs
+        const commandIds = commands.listCommands();
+
+        for (const id of commandIds) {
+          try {
+            // Get command metadata using various CommandRegistry methods
+            const description = await commands.describedBy(id);
+            const label = commands.label(id);
+            const caption = commands.caption(id);
+            const usage = commands.usage(id);
+
+            commandList.push({
+              id,
+              label: label || undefined,
+              caption: caption || undefined,
+              description: usage || undefined,
+              args: description?.args || undefined
+            });
+          } catch (error) {
+            // Some commands might not have descriptions, skip them
+            commandList.push({ id });
+          }
+        }
+
+        return {
+          success: true,
+          commandCount: commandList.length,
+          commands: commandList
+        };
+      } catch (error) {
+        return {
+          success: false,
+          error: `Failed to discover commands: ${error instanceof Error ? error.message : String(error)}`
+        };
+      }
+    }
+  });
+}
+
+/**
+ * Create a tool to execute a specific JupyterLab command
+ */
+export function createExecuteCommandTool(
+  commands: CommandRegistry,
+  settingsModel: AISettingsModel
+): ITool {
+  return tool({
+    name: 'execute_command',
+    description:
+      'Execute a specific JupyterLab command with optional arguments',
+    parameters: z.object({
+      commandId: z.string().describe('The ID of the command to execute'),
+      args: z
+        .any()
+        .optional()
+        .describe('Optional arguments to pass to the command')
+    }),
+    needsApproval: async (_context, { commandId }) => {
+      // Use configurable list of commands requiring approval
+      const commandsRequiringApproval =
+        settingsModel.config.commandsRequiringApproval;
+
+      return commandsRequiringApproval.some(
+        cmd => commandId.includes(cmd) || cmd.includes(commandId)
+      );
+    },
+    execute: async (input: { commandId: string; args?: any }) => {
+      const { commandId, args } = input;
+
+      // Check if command exists
+      if (!commands.hasCommand(commandId)) {
+        return {
+          success: false,
+          error: `Command '${commandId}' does not exist. Use 'discover_commands' to see available commands.`
+        };
+      }
+
+      try {
+        // Execute the command
+        const result = await commands.execute(commandId, args);
+
+        // Handle Widget objects specially (including subclasses like DocumentWidget)
+        let serializedResult;
+        if (
+          result &&
+          typeof result === 'object' &&
+          (result.constructor?.name?.includes('Widget') || result.id)
+        ) {
+          serializedResult = {
+            type: result.constructor?.name || 'Widget',
+            id: result.id,
+            title: result.title?.label || result.title,
+            className: result.className
+          };
+        } else {
+          // For other objects, try JSON serialization with fallback
+          try {
+            serializedResult = JSON.parse(JSON.stringify(result));
+          } catch {
+            serializedResult = result
+              ? '[Complex object - cannot serialize]'
+              : 'Command executed successfully';
+          }
+        }
+
+        return {
+          success: true,
+          commandId,
+          result: serializedResult
+        };
+      } catch (error) {
+        return {
+          success: false,
+          error: `Failed to execute command '${commandId}': ${error instanceof Error ? error.message : String(error)}`
+        };
+      }
+    }
+  });
+}