@jupyterlite/ai 0.6.2 → 0.7.0

package/lib/default-providers/Ollama/settings-schema.json

@@ -0,0 +1,146 @@
+{
+    "$schema": "http://json-schema.org/draft-07/schema#",
+    "type": "object",
+    "properties": {
+        "numa": {
+            "type": "boolean"
+        },
+        "numCtx": {
+            "type": "number"
+        },
+        "numBatch": {
+            "type": "number"
+        },
+        "numGpu": {
+            "type": "number"
+        },
+        "mainGpu": {
+            "type": "number"
+        },
+        "lowVram": {
+            "type": "boolean"
+        },
+        "f16Kv": {
+            "type": "boolean"
+        },
+        "logitsAll": {
+            "type": "boolean"
+        },
+        "vocabOnly": {
+            "type": "boolean"
+        },
+        "useMmap": {
+            "type": "boolean"
+        },
+        "useMlock": {
+            "type": "boolean"
+        },
+        "embeddingOnly": {
+            "type": "boolean"
+        },
+        "numThread": {
+            "type": "number"
+        },
+        "numKeep": {
+            "type": "number"
+        },
+        "seed": {
+            "type": "number"
+        },
+        "numPredict": {
+            "type": "number"
+        },
+        "topK": {
+            "type": "number"
+        },
+        "topP": {
+            "type": "number"
+        },
+        "tfsZ": {
+            "type": "number"
+        },
+        "typicalP": {
+            "type": "number"
+        },
+        "repeatLastN": {
+            "type": "number"
+        },
+        "temperature": {
+            "type": "number"
+        },
+        "repeatPenalty": {
+            "type": "number"
+        },
+        "presencePenalty": {
+            "type": "number"
+        },
+        "frequencyPenalty": {
+            "type": "number"
+        },
+        "mirostat": {
+            "type": "number"
+        },
+        "mirostatTau": {
+            "type": "number"
+        },
+        "mirostatEta": {
+            "type": "number"
+        },
+        "penalizeNewline": {
+            "type": "boolean"
+        },
+        "keepAlive": {
+            "type": [
+                "string",
+                "number"
+            ],
+            "default": "5m"
+        },
+        "stop": {
+            "type": "array",
+            "items": {
+                "type": "string"
+            }
+        },
+        "disableStreaming": {
+            "type": "boolean",
+            "description": "Whether to disable streaming.\n\nIf streaming is bypassed, then `stream()` will defer to `invoke()`.\n\n- If true, will always bypass streaming case.\n- If false (default), will always use streaming case if available."
+        },
+        "model": {
+            "type": "string",
+            "description": "The model to invoke. If the model does not exist, it will be pulled.",
+            "default": ""
+        },
+        "baseUrl": {
+            "type": "string",
+            "description": "The host URL of the Ollama server.",
+            "default": ""
+        },
+        "headers": {
+            "type": "object",
+            "additionalProperties": false,
+            "description": "Optional HTTP Headers to include in the request."
+        },
+        "checkOrPullModel": {
+            "type": "boolean",
+            "description": "Whether or not to check the model exists on the local machine before invoking it. If set to `true`, the model will be pulled if it does not exist.",
+            "default": false
+        },
+        "streaming": {
+            "type": "boolean"
+        },
+        "format": {
+            "anyOf": [
+                {
+                    "type": "string"
+                },
+                {
+                    "type": "object"
+                }
+            ]
+        }
+    },
+    "additionalProperties": false,
+    "description": "Input to chat model class.",
+    "definitions": {}
+}

package/lib/default-providers/OpenAI/completer.d.ts

@@ -1,9 +1,7 @@
 import { CompletionHandler, IInlineCompletionContext } from '@jupyterlab/completer';
-import { BaseChatModel } from '@langchain/core/language_models/chat_models';
 import { BaseCompleter, IBaseCompleter } from '../../base-completer';
 export declare class OpenAICompleter implements IBaseCompleter {
     constructor(options: BaseCompleter.IOptions);
-    get completer(): BaseChatModel;
     /**
      * Getter and setter for the initial prompt.
      */

package/lib/default-providers/OpenAI/completer.js

@@ -3,12 +3,8 @@ import { ChatOpenAI } from '@langchain/openai';
 import { COMPLETION_SYSTEM_PROMPT } from '../../provider';
 export class OpenAICompleter {
     constructor(options) {
-        this._prompt = COMPLETION_SYSTEM_PROMPT;
         this._completer = new ChatOpenAI({ ...options.settings });
     }
-    get completer() {
-        return this._completer;
-    }
     /**
      * Getter and setter for the initial prompt.
      */
@@ -48,4 +44,6 @@ export class OpenAICompleter {
             return { items: [] };
         }
     }
+    _completer;
+    _prompt = COMPLETION_SYSTEM_PROMPT;
 }

package/lib/default-providers/WebLLM/completer.d.ts

@@ -0,0 +1,27 @@
+import { CompletionHandler, IInlineCompletionContext } from '@jupyterlab/completer';
+import { ChatWebLLM } from '@langchain/community/chat_models/webllm';
+import { BaseCompleter, IBaseCompleter } from '../../base-completer';
+export declare class WebLLMCompleter implements IBaseCompleter {
+    constructor(options: BaseCompleter.IOptions);
+    /**
+     * Initialize the WebLLM model
+     */
+    private _initializeModel;
+    /**
+     * Getter and setter for the initial prompt.
+     */
+    get prompt(): string;
+    set prompt(value: string);
+    get provider(): ChatWebLLM;
+    fetch(request: CompletionHandler.IRequest, context: IInlineCompletionContext): Promise<{
+        items: {
+            insertText: string;
+        }[];
+    }>;
+    private _completer;
+    private _prompt;
+    private _isInitialized;
+    private _isInitializing;
+    private _initError;
+    private _abortController;
+}

package/lib/default-providers/WebLLM/completer.js

@@ -0,0 +1,136 @@
+import { HumanMessage, SystemMessage } from '@langchain/core/messages';
+import { ChatWebLLM } from '@langchain/community/chat_models/webllm';
+import { COMPLETION_SYSTEM_PROMPT } from '../../provider';
+/**
+ * Regular expression to match the '```' string at the start of a string.
+ * So the completions returned by the LLM can still be kept after removing the code block formatting.
+ *
+ * For example, if the response contains the following content after typing `import pandas`:
+ *
+ * ```python
+ * as pd
+ * ```
+ *
+ * The formatting string after removing the code block delimiters will be:
+ *
+ * as pd
+ */
+const CODE_BLOCK_START_REGEX = /^```(?:[a-zA-Z]+)?\n?/;
+/**
+ * Regular expression to match the '```' string at the end of a string.
+ */
+const CODE_BLOCK_END_REGEX = /```$/;
+export class WebLLMCompleter {
+    constructor(options) {
+        const model = options.settings.model;
+        // provide model separately since ChatWebLLM expects it
+        this._completer = new ChatWebLLM({
+            ...options.settings,
+            model
+        });
+        // Initialize the model and track its status
+        this._isInitialized = false;
+        this._isInitializing = false;
+        this._initError = null;
+        void this._initializeModel();
+    }
+    /**
+     * Initialize the WebLLM model
+     */
+    async _initializeModel() {
+        if (this._isInitialized || this._isInitializing) {
+            return;
+        }
+        this._isInitializing = true;
+        try {
+            await this._completer.initialize((progress) => {
+                console.log('WebLLM initialization progress:', progress);
+            });
+            this._isInitialized = true;
+            this._isInitializing = false;
+            console.log('WebLLM model successfully initialized');
+        }
+        catch (error) {
+            this._initError =
+                error instanceof Error ? error : new Error(String(error));
+            this._isInitializing = false;
+            console.error('Failed to initialize WebLLM model:', error);
+        }
+    }
+    /**
+     * Getter and setter for the initial prompt.
+     */
+    get prompt() {
+        return this._prompt;
+    }
+    set prompt(value) {
+        this._prompt = value;
+    }
+    get provider() {
+        return this._completer;
+    }
+    async fetch(request, context) {
+        // Abort any pending request
+        if (this._abortController) {
+            this._abortController.abort();
+        }
+        // Create a new abort controller for this request
+        this._abortController = new AbortController();
+        const signal = this._abortController.signal;
+        if (!this._isInitialized) {
+            if (this._initError) {
+                console.error('WebLLM model failed to initialize:', this._initError);
+                return { items: [] };
+            }
+            if (!this._isInitializing) {
+                // Try to initialize again if it's not currently initializing
+                await this._initializeModel();
+            }
+            else {
+                console.log('WebLLM model is still initializing, please try again later');
+                return { items: [] };
+            }
+            // Return empty if still not initialized
+            if (!this._isInitialized) {
+                return { items: [] };
+            }
+        }
+        const { text, offset: cursorOffset } = request;
+        const prompt = text.slice(0, cursorOffset);
+        const trimmedPrompt = prompt.trim();
+        const messages = [
+            new SystemMessage(this._prompt),
+            new HumanMessage(trimmedPrompt)
+        ];
+        try {
+            console.log('Trigger invoke');
+            const response = await this._completer.invoke(messages, { signal });
+            let content = response.content;
+            console.log('Response content:', content);
+            if (CODE_BLOCK_START_REGEX.test(content)) {
+                content = content
+                    .replace(CODE_BLOCK_START_REGEX, '')
+                    .replace(CODE_BLOCK_END_REGEX, '');
+            }
+            const items = [{ insertText: content }];
+            return {
+                items
+            };
+        }
+        catch (error) {
+            if (error instanceof Error) {
+                console.error('Error fetching completion from WebLLM:', error.message);
+            }
+            else {
+                console.error('Unknown error fetching completion from WebLLM:', error);
+            }
+            return { items: [] };
+        }
+    }
+    _completer;
+    _prompt = COMPLETION_SYSTEM_PROMPT;
+    _isInitialized = false;
+    _isInitializing = false;
+    _initError = null;
+    _abortController = null;
+}

package/lib/default-providers/WebLLM/instructions.d.ts

@@ -0,0 +1,6 @@
+declare const _default: "\nWebLLM enables running LLMs directly in your browser, making it possible to use AI features without sending data to external servers.\n\n<i class=\"fas fa-info-circle\"></i> WebLLM runs models entirely in your browser, so initial model download may be large (100MB-2GB depending on the model).\n\n<i class=\"fas fa-exclamation-triangle\"></i> <strong>Requirements:</strong> WebLLM requires a browser with WebGPU support (Chrome 113+, Edge 113+, or Safari 17+). It will not work on older browsers or browsers without WebGPU enabled.\n\n1. Enter a model in the JupyterLab settings under the **Ai providers** section. Select the `WebLLM` provider and type the model you want to use.\n2. When you first use WebLLM, your browser will download the model. A progress notification will appear:\n3. Once loaded, use the chat\n4. Example of available models:\n  - Llama-3.2-1B-Instruct-q4f32_1-MLC\n  - Mistral-7B-Instruct-v0.3-q4f32_1-MLC\n  - Qwen3-0.6B-q4f32_1-MLC\n5. See the full list of models: https://github.com/mlc-ai/web-llm/blob/632d34725629b480b5b2772379ef5c150b1286f0/src/config.ts#L303-L309\n\n<i class=\"fas fa-exclamation-triangle\"></i> Model performance depends on your device's hardware capabilities. More powerful devices will run models faster. Some larger models may not work well on devices with limited GPU memory or may experience slow response times.\n";
+export default _default;
+/**
+ * Check if the browser supports WebLLM.
+ */
+export declare function compatibilityCheck(): Promise<string | null>;

package/lib/default-providers/WebLLM/instructions.js

@@ -0,0 +1,32 @@
+export default `
+WebLLM enables running LLMs directly in your browser, making it possible to use AI features without sending data to external servers.
+
+<i class="fas fa-info-circle"></i> WebLLM runs models entirely in your browser, so initial model download may be large (100MB-2GB depending on the model).
+
+<i class="fas fa-exclamation-triangle"></i> <strong>Requirements:</strong> WebLLM requires a browser with WebGPU support (Chrome 113+, Edge 113+, or Safari 17+). It will not work on older browsers or browsers without WebGPU enabled.
+
+1. Enter a model in the JupyterLab settings under the **Ai providers** section. Select the \`WebLLM\` provider and type the model you want to use.
+2. When you first use WebLLM, your browser will download the model. A progress notification will appear:
+3. Once loaded, use the chat
+4. Example of available models:
+  - Llama-3.2-1B-Instruct-q4f32_1-MLC
+  - Mistral-7B-Instruct-v0.3-q4f32_1-MLC
+  - Qwen3-0.6B-q4f32_1-MLC
+5. See the full list of models: https://github.com/mlc-ai/web-llm/blob/632d34725629b480b5b2772379ef5c150b1286f0/src/config.ts#L303-L309
+
+<i class="fas fa-exclamation-triangle"></i> Model performance depends on your device's hardware capabilities. More powerful devices will run models faster. Some larger models may not work well on devices with limited GPU memory or may experience slow response times.
+`;
+/**
+ * Check if the browser supports WebLLM.
+ */
+export async function compatibilityCheck() {
+    // Check if the browser supports the ChromeAI model
+    if (typeof navigator === 'undefined' || !('gpu' in navigator)) {
+        return 'Your browser does not support WebLLM, it does not support required WebGPU.';
+    }
+    if ((await navigator.gpu.requestAdapter()) === null) {
+        return 'You may need to enable WebGPU, `await navigator.gpu.requestAdapter()` is null.';
+    }
+    // If the model is available, return null to indicate compatibility
+    return null;
+}

package/lib/default-providers/WebLLM/settings-schema.json

@@ -0,0 +1,21 @@
+{
+    "$schema": "http://json-schema.org/draft-07/schema#",
+    "type": "object",
+    "properties": {
+        "disableStreaming": {
+            "type": "boolean",
+            "description": "Whether to disable streaming.\n\nIf streaming is bypassed, then `stream()` will defer to `invoke()`.\n\n- If true, will always bypass streaming case.\n- If false (default), will always use streaming case if available."
+        },
+        "temperature": {
+            "type": "number"
+        },
+        "model": {
+            "type": "string"
+        }
+    },
+    "required": [
+        "model"
+    ],
+    "additionalProperties": false,
+    "definitions": {}
+}

package/lib/default-providers/index.js

@@ -1,21 +1,31 @@
+import { Notification } from '@jupyterlab/apputils';
 import { ChatAnthropic } from '@langchain/anthropic';
+import { ChatWebLLM } from '@langchain/community/chat_models/webllm';
 import { ChromeAI } from '@langchain/community/experimental/llms/chrome_ai';
 import { ChatMistralAI } from '@langchain/mistralai';
+import { ChatOllama } from '@langchain/ollama';
 import { ChatOpenAI } from '@langchain/openai';
-import { IAIProviderRegistry } from '../tokens';
 // Import completers
 import { AnthropicCompleter } from './Anthropic/completer';
 import { ChromeCompleter } from './ChromeAI/completer';
 import { CodestralCompleter } from './MistralAI/completer';
+import { OllamaCompleter } from './Ollama/completer';
 import { OpenAICompleter } from './OpenAI/completer';
+import { WebLLMCompleter } from './WebLLM/completer';
 // Import Settings
 import AnthropicSettings from './Anthropic/settings-schema.json';
 import ChromeAISettings from './ChromeAI/settings-schema.json';
 import MistralAISettings from './MistralAI/settings-schema.json';
+import OllamaAISettings from './Ollama/settings-schema.json';
 import OpenAISettings from './OpenAI/settings-schema.json';
+import WebLLMSettings from './WebLLM/settings-schema.json';
 // Import instructions
-import ChromeAIInstructions from './ChromeAI/instructions';
+import ChromeAIInstructions, { compatibilityCheck as chromeAICompatibilityCheck } from './ChromeAI/instructions';
 import MistralAIInstructions from './MistralAI/instructions';
+import OllamaInstructions from './Ollama/instructions';
+import WebLLMInstructions, { compatibilityCheck as webLLMCompatibilityCheck } from './WebLLM/instructions';
+import { prebuiltAppConfig } from '@mlc-ai/web-llm';
+import { IAIProviderRegistry } from '../tokens';
 // Build the AIProvider list
 const AIProviders = [
     {
@@ -32,7 +42,8 @@ const AIProviders = [
         chatModel: ChromeAI,
         completer: ChromeCompleter,
         instructions: ChromeAIInstructions,
-        settingsSchema: ChromeAISettings
+        settingsSchema: ChromeAISettings,
+        compatibilityCheck: chromeAICompatibilityCheck
     },
     {
         name: 'MistralAI',
@@ -41,6 +52,13 @@ const AIProviders = [
         instructions: MistralAIInstructions,
         settingsSchema: MistralAISettings
     },
+    {
+        name: 'Ollama',
+        chatModel: ChatOllama,
+        completer: OllamaCompleter,
+        instructions: OllamaInstructions,
+        settingsSchema: OllamaAISettings
+    },
     {
         name: 'OpenAI',
         chatModel: ChatOpenAI,
@@ -48,7 +66,100 @@ const AIProviders = [
         settingsSchema: OpenAISettings
     }
 ];
-export const defaultProviderPlugins = AIProviders.map(provider => {
+/**
+ * Register the WebLLM provider in a separate plugin since it creates notifications
+ * when the model is changed in the settings.
+ */
+const webLLMProviderPlugin = {
+    id: '@jupyterlite/ai:webllm',
+    description: 'Register the WebLLM provider',
+    autoStart: true,
+    requires: [IAIProviderRegistry],
+    activate: (app, registry) => {
+        registry.add({
+            name: 'WebLLM',
+            chatModel: ChatWebLLM,
+            completer: WebLLMCompleter,
+            settingsSchema: WebLLMSettings,
+            instructions: WebLLMInstructions,
+            compatibilityCheck: webLLMCompatibilityCheck,
+            exposeChatModel: true
+        });
+        registry.providerChanged.connect(async (sender, args) => {
+            const { currentName, currentChatModel, chatError } = registry;
+            if (currentChatModel === null) {
+                Notification.emit(chatError, 'error', {
+                    autoClose: 2000
+                });
+                return;
+            }
+            // TODO: implement a proper way to handle models that may need to be initialized before being used.
+            // Mostly applies to WebLLM and ChromeAI as they may need to download the model in the browser first.
+            if (currentName === 'WebLLM') {
+                const compatibilityError = await webLLMCompatibilityCheck();
+                if (compatibilityError) {
+                    Notification.dismiss();
+                    Notification.emit(compatibilityError, 'error', {
+                        autoClose: 2000
+                    });
+                    return;
+                }
+                const model = currentChatModel;
+                if (model === null || !model.model) {
+                    return;
+                }
+                // Find if the model is part of the prebuiltAppConfig
+                const modelRecord = prebuiltAppConfig.model_list.find(modelRecord => modelRecord.model_id === model.model);
+                if (!modelRecord) {
+                    Notification.dismiss();
+                    Notification.emit(`Model ${model.model} not found in the prebuiltAppConfig`, 'error', {
+                        autoClose: 2000
+                    });
+                    return;
+                }
+                // create a notification
+                const notification = Notification.emit('Loading model...', 'in-progress', {
+                    autoClose: false,
+                    progress: 0
+                });
+                try {
+                    void model.initialize(report => {
+                        const { progress, text } = report;
+                        if (progress === 1) {
+                            Notification.update({
+                                id: notification,
+                                progress: 1,
+                                message: `Model ${model.model} loaded successfully`,
+                                type: 'success',
+                                autoClose: 2000
+                            });
+                            return;
+                        }
+                        Notification.update({
+                            id: notification,
+                            progress: progress / 1,
+                            message: text,
+                            type: 'in-progress'
+                        });
+                    });
+                }
+                catch (err) {
+                    Notification.update({
+                        id: notification,
+                        progress: 1,
+                        message: `Error loading model ${model.model}`,
+                        type: 'error',
+                        autoClose: 2000
+                    });
+                }
+            }
+        });
+    }
+};
+/**
+ * Register all default AI providers.
+ */
+const aiProviderPlugins = AIProviders.map(provider => {
     return {
         id: `@jupyterlite/ai:${provider.name}`,
         autoStart: true,
@@ -58,3 +169,7 @@ export const defaultProviderPlugins = AIProviders.map(provider => {
         }
     };
 });
+export const defaultProviderPlugins = [
+    webLLMProviderPlugin,
+    ...aiProviderPlugins
+];

package/lib/index.d.ts

@@ -1,6 +1,6 @@
 import { IChatCommandRegistry } from '@jupyter/chat';
 import { JupyterFrontEndPlugin } from '@jupyterlab/application';
-import { ISettingConnector } from '@jupyterlab/settingregistry';
 import { IAIProviderRegistry } from './tokens';
-declare const _default: (JupyterFrontEndPlugin<void> | JupyterFrontEndPlugin<IChatCommandRegistry> | JupyterFrontEndPlugin<IAIProviderRegistry> | JupyterFrontEndPlugin<ISettingConnector>)[];
+declare const _default: (JupyterFrontEndPlugin<void> | JupyterFrontEndPlugin<IChatCommandRegistry> | JupyterFrontEndPlugin<IAIProviderRegistry>)[];
 export default _default;
+export { IAIProviderRegistry } from './tokens';