@jupyterlite/ai 0.9.0-a2 → 0.9.0-a4

package/README.md

@@ -5,15 +5,10 @@
 
 AI code completions and chat for JupyterLab, Notebook 7 and JupyterLite ✨
 
-[a screencast showing the Jupyterlite AI extension in JupyterLite](https://github.com/jupyterlite/ai/assets/591645/855c4e3e-3a63-4868-8052-5c9909922c21)
+[a screencast showing the Jupyterlite AI extension in JupyterLite](https://github.com/user-attachments/assets/e33d7d84-53ca-4835-a034-b6757476c98b)
 
 ## Requirements
 
-> [!NOTE]
-> This extension is meant to be used in JupyterLite to enable AI code completions and chat in the browser, with a specific provider.
-> To enable more AI providers in JupyterLab and Jupyter Notebook, we recommend using the [Jupyter AI](https://github.com/jupyterlab/jupyter-ai) extension directly.
-> At the moment Jupyter AI is not compatible with JupyterLite, but might be to some extent in the future.
-
 - JupyterLab >= 4.4.0 or Notebook >= 7.4.0
 
 ## ✨ Try it in your browser ✨
@@ -32,63 +27,221 @@ To install the extension, execute:
 pip install jupyterlite-ai
 ```
 
-To install requirements (jupyterlab, jupyterlite and notebook), there is an optional dependencies argument:
+To install requirements (JupyterLab, JupyterLite and Notebook):
 
 ```bash
 pip install jupyterlite-ai[jupyter]
 ```
 
-# Usage
+## Usage
+
+> [!NOTE]
+> This documentation applies to the upcoming **0.9.0** release.
+> For the latest stable version, please refer to the [0.8.x branch](https://github.com/jupyterlite/ai/tree/0.8.x).
 
 AI providers typically require using an API key to access their models.
 
-The process is different for each provider, so you may refer to their documentation to learn how to generate new API keys, if they are not covered in the sections below.
+The process is different for each provider, so you may refer to their documentation to learn how to generate new API keys.
 
-## Using MistralAI
+### Using a provider with an API key (e.g. Anthropic, MistralAI, OpenAI)
 
-> [!WARNING]
-> This extension is still very much experimental. It is not an official MistralAI extension.
+1. Open the AI settings and
+2. Click on "Add a new provider"
+3. Enter the details for the provider
+4. In the chat, select the new provider
+
+![screenshot showing the dialog to add a new provider](https://github.com/user-attachments/assets/823c71c6-5807-44c8-80b6-2e59379a65d5)
 
-1. Go to https://console.mistral.ai/api-keys/ and create an API key.
+### Using Ollama
 
-![Screenshot showing how to create an API key](./img/1-api-key.png)
+[Ollama](https://ollama.com/) allows you to run open-weight LLMs locally on your machine.
 
-2. Open the JupyterLab settings and go to the **Ai providers** section to select the `MistralAI`
-   provider and the API key (required).
+#### Setting up Ollama
 
-![Screenshot showing how to add the API key to the settings](./img/2-jupyterlab-settings.png)
+1. Install Ollama following the instructions at https://ollama.com/download
+2. Pull a model, for example:
 
-3. Open the chat, or use the inline completer
+```bash
+ollama pull llama3.2
+```
 
-![Screenshot showing how to use the chat](./img/3-usage.png)
+3. Start the Ollama server (it typically runs on `http://localhost:11434`)
 
-## Using ChromeAI
+#### Configuring `jupyterlite-ai` to use Ollama
 
-> [!WARNING]
-> Support for ChromeAI is still experimental and only available in Google Chrome.
+1. In JupyterLab, open the AI settings panel and go to the **Providers** section
+2. Click on "Add a new provider"
+3. Select the **Ollama** provider
+4. Configure the following settings:
+   - **Model**: The model name you pulled (e.g., `llama3.2`)
 
-You can test ChromeAI is enabled in your browser by going to the following URL: https://chromeai.org/
+### Using a generic OpenAI-compatible provider
 
-Enable the proper flags in Google Chrome.
+The Generic provider allows you to connect to any OpenAI-compatible API endpoint.
 
-- chrome://flags/#prompt-api-for-gemini-nano
-  - Select: `Enabled`
-- chrome://flags/#optimization-guide-on-device-model
-  - Select: `Enabled BypassPrefRequirement`
-- chrome://components
-  - Click `Check for Update` on Optimization Guide On Device Model to download the model
-- [Optional] chrome://flags/#text-safety-classifier
+1. In JupyterLab, open the AI settings panel and go to the **Providers** section
+2. Click on "Add a new provider"
+3. Select the **Generic** provider
+4. Configure the following settings:
+   - **Base URL**: The base URL of your API endpoint
+   - **Model**: The model name to use
+   - **API Key**: Your API key (if required by the provider)
 
-![a screenshot showing how to enable the ChromeAI flag in Google Chrome](https://github.com/user-attachments/assets/d48f46cc-52ee-4ce5-9eaf-c763cdbee04c)
+### Using LiteLLM Proxy
 
-Then restart Chrome for these changes to take effect.
+[LiteLLM Proxy](https://docs.litellm.ai/docs/simple_proxy) is an OpenAI-compatible proxy server that allows you to call 100+ LLMs through a unified interface.
 
-> [!WARNING]
-> On first use, Chrome will download the on-device model, which can be as large as 22GB (according to their docs and at the time of writing).
-> During the download, ChromeAI may not be available via the extension.
+Using LiteLLM Proxy with jupyterlite-ai provides flexibility to switch between different AI providers (OpenAI, Anthropic, Google, Azure, local models, etc.) without changing your JupyterLite configuration. It's particularly useful for enterprise deployments where the proxy can be hosted within private infrastructure to manage external API calls and keep API keys server-side.
+
+#### Setting up LiteLLM Proxy
+
+1. Install LiteLLM:
+
+Follow the instructions at https://docs.litellm.ai/docs/simple_proxy.
+
+2. Create a `litellm_config.yaml` file with your model configuration:
+
+```yaml
+model_list:
+  - model_name: gpt-5
+    litellm_params:
+      model: gpt-5
+      api_key: os.environ/OPENAI_API_KEY
+
+  - model_name: claude-sonnet
+    litellm_params:
+      model: claude-sonnet-4-5-20250929
+      api_key: os.environ/ANTHROPIC_API_KEY
+```
+
+3. Start the proxy server, for example:
+
+```bash
+litellm --config litellm_config.yaml
+```
+
+The proxy will start on `http://0.0.0.0:4000` by default.
+
+#### Configuring `jupyterlite-ai` to use LiteLLM Proxy
+
+Configure the [Generic provider (OpenAI-compatible)](#using-a-generic-openai-compatible-provider) with the following settings:
+
+- **Base URL**: `http://0.0.0.0:4000` (or your proxy server URL)
+- **Model**: The model name from your `litellm_config.yaml` (e.g., `gpt-5`, `claude-sonnet`)
+- **API Key (optional)**: If the LiteLLM Proxy server requires an API key, provide it here.
+
+> [!IMPORTANT]
+> The API key must be configured on the LiteLLM Proxy server (in the `litellm_config.yaml` file). Providing an API key via the AI provider settings UI will not have any effect, as the proxy server handles authentication with the upstream AI providers.
 
 > [!NOTE]
-> For more information about Chrome Built-in AI: https://developer.chrome.com/docs/ai/get-started
+> For more information about LiteLLM Proxy configuration, see the [LiteLLM documentation](https://docs.litellm.ai/docs/simple_proxy).
+
+## Custom Providers
+
+`jupyterlite-ai` supports custom AI providers through its provider registry system. Third-party providers can be registered programmatically in a JupyterLab extension.
+
+Providers are based on the [Vercel AI SDK](https://sdk.vercel.ai/docs/introduction), which provides a unified interface for working with different AI models.
+
+### Registering a Custom Provider
+
+**Example: Registering a custom OpenAI-compatible provider**
+
+```typescript
+import {
+  JupyterFrontEnd,
+  JupyterFrontEndPlugin
+} from '@jupyterlab/application';
+import { IProviderRegistry } from '@jupyterlite/ai';
+import { createOpenAI } from '@ai-sdk/openai';
+
+const plugin: JupyterFrontEndPlugin<void> = {
+  id: 'my-extension:custom-provider',
+  autoStart: true,
+  requires: [IProviderRegistry],
+  activate: (app: JupyterFrontEnd, registry: IProviderRegistry) => {
+    const providerInfo = {
+      id: 'my-custom-provider',
+      name: 'My Custom Provider',
+      apiKeyRequirement: 'required' as const,
+      defaultModels: ['my-model'],
+      supportsBaseURL: true,
+      factory: (options: {
+        apiKey: string;
+        baseURL?: string;
+        model?: string;
+      }) => {
+        const provider = createOpenAI({
+          apiKey: options.apiKey,
+          baseURL: options.baseURL || 'https://api.example.com/v1'
+        });
+        return provider(options.model || 'my-model');
+      }
+    };
+
+    registry.registerProvider(providerInfo);
+  }
+};
+```
+
+The provider configuration object requires the following properties:
+
+- `id`: Unique identifier for the provider
+- `name`: Display name shown in the settings UI
+- `apiKeyRequirement`: Whether an API key is `'required'`, `'optional'`, or `'none'`
+- `defaultModels`: Array of model names to show in the settings
+- `supportsBaseURL`: Whether the provider supports a custom base URL
+- `factory`: Function that creates and returns a language model (the registry automatically wraps it for chat usage)
+
+**Example: Using a custom fetch function**
+
+You can provide a custom `fetch` function to the provider, which is useful for adding custom headers, handling authentication, or routing requests through a proxy:
+
+```typescript
+factory: (options: { apiKey: string; baseURL?: string; model?: string }) => {
+  const provider = createOpenAI({
+    apiKey: options.apiKey,
+    baseURL: options.baseURL || 'https://api.example.com/v1',
+    fetch: async (url, init) => {
+      // Custom fetch implementation
+      const modifiedInit = {
+        ...init,
+        headers: {
+          ...init?.headers,
+          'X-Custom-Header': 'custom-value'
+        }
+      };
+      return fetch(url, modifiedInit);
+    }
+  });
+  return provider(options.model || 'my-model');
+};
+```
+
+## API key management
+
+To avoid storing the API keys in the settings, `jupyterlite-ai` uses [jupyter-secrets-manager](https://github.com/jupyterlab-contrib/jupyter-secrets-manager) by default.
+
+The secrets manager get the API keys from a connector in a secure way.\
+The default connector of the secrets manager is _in memory_, which means that **the API keys are reset when reloading the page**.
+
+To prevent the keys from being reset on reload, there are two options:
+
+1. use a connector that fetches the keys on a remote server (using secure rest API, or web socket)
+
+This is the recommended method, as it ensures the security of the keys and makes them accessible only to logged-in users. \
+But it requires some frontend and backend deployments:
+
+- a server that can store and send the keys on demand
+- a way to get authenticated to the server
+- a frontend extension providing the connector, able to connect to the server side
+
+2. disable the use of the secrets manager from the AI settings panel
+
+> [!WARNING]
+> The API keys will be stored in plain text using the settings system of Jupyterlab
+>
+> - using Jupyterlab, the settings are stored in a [directory](https://jupyterlab.readthedocs.io/en/stable/user/directories.html#jupyterlab-user-settings-directory) on the server
+> - using Jupyterlite, the settings are stored in the [browser](https://jupyterlite.readthedocs.io/en/latest/howto/configure/storage.html#configure-the-browser-storage)
 
 ## Uninstall
 
@@ -136,6 +289,29 @@ By default, the `jlpm build` command generates the source maps for this extensio
 jupyter lab build --minimize=False
 ```
 
+### Running UI tests
+
+The UI tests use Playwright and can be configured with environment variables:
+
+- `PWVIDEO`: Controls video recording during tests (default: `retain-on-failure`)
+  - `on`: Record video for all tests
+  - `off`: Do not record video
+  - `retain-on-failure`: Only keep videos for failed tests
+- `PWSLOWMO`: Adds a delay (in milliseconds) between Playwright actions for debugging (default: `0`)
+
+Example usage:
+
+```bash
+# Record all test videos
+PWVIDEO=on jlpm playwright test
+
+# Slow down test execution by 500ms per action
+PWSLOWMO=500 jlpm playwright test
+
+# Combine both options
+PWVIDEO=on PWSLOWMO=1000 jlpm playwright test
+```
+
 ### Development uninstall
 
 ```bash

package/lib/agent.d.ts

@@ -1,8 +1,59 @@
 import { ISignal } from '@lumino/signaling';
 import { ISecretsManager } from 'jupyter-secrets-manager';
+import { BrowserMCPServerStreamableHttp } from './mcp/browser';
 import { AISettingsModel } from './models/settings-model';
-import type { IChatProviderRegistry } from './tokens';
+import type { IProviderRegistry } from './tokens';
 import { ITool, IToolRegistry, ITokenUsage } from './tokens';
+export declare namespace AgentManagerFactory {
+    interface IOptions {
+        /**
+         * The settings model.
+         */
+        settingsModel: AISettingsModel;
+        /**
+         * The secrets manager.
+         */
+        secretsManager?: ISecretsManager;
+        /**
+         * The token used to request the secrets manager.
+         */
+        token: symbol;
+    }
+}
+export declare class AgentManagerFactory {
+    constructor(options: AgentManagerFactory.IOptions);
+    createAgent(options: IAgentManagerOptions): AgentManager;
+    /**
+     * Signal emitted when MCP connection status changes
+     */
+    get mcpConnectionChanged(): ISignal<this, boolean>;
+    /**
+     * Checks if a specific MCP server is connected by server name.
+     * @param serverName The name of the MCP server to check
+     * @returns True if the server is connected, false otherwise
+     */
+    isMCPServerConnected(serverName: string): boolean;
+    /**
+     * Handles settings changes and reinitializes the agent.
+     */
+    private _onSettingsChanged;
+    /**
+     * Initializes MCP (Model Context Protocol) servers based on current settings.
+     * Closes existing servers and connects to enabled servers from configuration.
+     */
+    private _initializeMCPServers;
+    /**
+     * Initializes the AI agent with current settings and tools.
+     * Sets up the agent with model configuration, tools, and MCP servers.
+     */
+    private _initializeAgents;
+    private _agentManagers;
+    private _settingsModel;
+    private _secretsManager?;
+    private _mcpServers;
+    private _mcpConnectionChanged;
+    private _isInitializing;
+}
 /**
  * Event type mapping for type safety with inlined interface definitions
  */
@@ -22,18 +73,18 @@ export interface IAgentEventTypeMap {
     tool_call_start: {
         callId: string;
         toolName: string;
-        input: any;
+        input: string;
     };
     tool_call_complete: {
         callId: string;
         toolName: string;
-        output: any;
+        output: string;
         isError: boolean;
     };
     tool_approval_required: {
         interruptionId: string;
         toolName: string;
-        toolInput: any;
+        toolInput: string;
         callId?: string;
     };
     grouped_approval_required: {
@@ -41,7 +92,7 @@ export interface IAgentEventTypeMap {
         approvals: Array<{
             interruptionId: string;
             toolName: string;
-            toolInput: any;
+            toolInput: string;
         }>;
     };
     error: {
@@ -68,17 +119,21 @@ export interface IAgentManagerOptions {
      */
     toolRegistry?: IToolRegistry;
     /**
-     * Optional chat provider registry for model creation
+     * Optional provider registry for model creation
      */
-    chatProviderRegistry?: IChatProviderRegistry;
+    providerRegistry?: IProviderRegistry;
     /**
      * The secrets manager.
      */
     secretsManager?: ISecretsManager;
     /**
-     * The token used to request the secrets manager.
+     * The active provider to use with this agent.
      */
-    token: symbol;
+    activeProvider?: string;
+    /**
+     * Initial token usage.
+     */
+    tokenUsage?: ITokenUsage;
 }
 /**
  * Manages the AI agent lifecycle and execution loop.
@@ -97,9 +152,9 @@ export declare class AgentManager {
      */
     get agentEvent(): ISignal<this, IAgentEvent>;
     /**
-     * Signal emitted when MCP connection status changes
+     * Signal emitted when the active provider has changed.
      */
-    get mcpConnectionChanged(): ISignal<this, boolean>;
+    get activeProviderChanged(): ISignal<this, string | undefined>;
     /**
      * Gets the current token usage statistics.
      */
@@ -108,6 +163,11 @@ export declare class AgentManager {
      * Signal emitted when token usage statistics change.
      */
     get tokenUsageChanged(): ISignal<this, ITokenUsage>;
+    /**
+     * The active provider for this agent.
+     */
+    get activeProvider(): string;
+    set activeProvider(value: string);
     /**
      * Sets the selected tools by name and reinitializes the agent.
      * @param toolNames Array of tool names to select
@@ -118,12 +178,6 @@ export declare class AgentManager {
      * @returns Array of selected tools formatted for OpenAI agents
      */
     get selectedAgentTools(): ITool[];
-    /**
-     * Checks if a specific MCP server is connected by server name.
-     * @param serverName The name of the MCP server to check
-     * @returns True if the server is connected, false otherwise
-     */
-    isMCPServerConnected(serverName: string): boolean;
     /**
      * Checks if the current configuration is valid for agent operations.
      * Uses the provider registry to determine if an API key is required.
@@ -167,26 +221,23 @@ export declare class AgentManager {
      * @param interruptionIds Array of interruption IDs to reject
      */
     rejectGroupedToolCalls(groupId: string, interruptionIds: string[]): Promise<void>;
-    /**
-     * Handles settings changes and reinitializes the agent.
-     */
-    private _onSettingsChanged;
     /**
      * Initializes the AI agent with current settings and tools.
      * Sets up the agent with model configuration, tools, and MCP servers.
      */
-    private _initializeAgent;
-    /**
-     * Initializes MCP (Model Context Protocol) servers based on current settings.
-     * Closes existing servers and connects to enabled servers from configuration.
-     */
-    private _initializeMCPServers;
+    initializeAgent: (mcpServers?: BrowserMCPServerStreamableHttp[]) => Promise<void>;
     /**
      * Processes the result stream from agent execution.
      * Handles message streaming, tool calls, and emits appropriate events.
      * @param result The async iterable result from agent execution
      */
     private _processRunResult;
+    /**
+     * Formats tool input for display by pretty-printing JSON strings.
+     * @param input The tool input string to format
+     * @returns Pretty-printed JSON string
+     */
+    private _formatToolInput;
     /**
      * Handles the start of a tool call from the model event.
      * @param modelEvent The model event containing tool call information
@@ -225,7 +276,7 @@ export declare class AgentManager {
     private _getEnhancedSystemPrompt;
     private _settingsModel;
     private _toolRegistry?;
-    private _chatProviderRegistry?;
+    private _providerRegistry?;
     private _secretsManager?;
     private _selectedToolNames;
     private _agent;
@@ -237,7 +288,8 @@ export declare class AgentManager {
     private _pendingApprovals;
     private _interruptedState;
     private _agentEvent;
-    private _mcpConnectionChanged;
     private _tokenUsage;
     private _tokenUsageChanged;
+    private _activeProvider;
+    private _activeProviderChanged;
 }