@jupyterlite/ai 0.6.2 → 0.8.0

package/lib/default-providers/index.js

@@ -1,26 +1,40 @@
+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 { ChatGoogleGenerativeAI } from '@langchain/google-genai';
 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 { GeminiCompleter } from './Gemini/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 GeminiSettings from './Gemini/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 GeminiInstructions from './Gemini/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 = [
     {
         name: 'Anthropic',
-        chatModel: ChatAnthropic,
+        chat: ChatAnthropic,
         completer: AnthropicCompleter,
         settingsSchema: AnthropicSettings,
         errorMessage: (error) => error.error.error.message
@@ -29,26 +43,127 @@ const AIProviders = [
         name: 'ChromeAI',
         // TODO: fix
         // @ts-expect-error: missing properties
-        chatModel: ChromeAI,
+        chat: ChromeAI,
         completer: ChromeCompleter,
         instructions: ChromeAIInstructions,
-        settingsSchema: ChromeAISettings
+        settingsSchema: ChromeAISettings,
+        compatibilityCheck: chromeAICompatibilityCheck
     },
     {
         name: 'MistralAI',
-        chatModel: ChatMistralAI,
+        chat: ChatMistralAI,
         completer: CodestralCompleter,
         instructions: MistralAIInstructions,
         settingsSchema: MistralAISettings
     },
+    {
+        name: 'Ollama',
+        chat: ChatOllama,
+        completer: OllamaCompleter,
+        instructions: OllamaInstructions,
+        settingsSchema: OllamaAISettings
+    },
+    {
+        name: 'Gemini',
+        chat: ChatGoogleGenerativeAI,
+        completer: GeminiCompleter,
+        instructions: GeminiInstructions,
+        settingsSchema: GeminiSettings
+    },
     {
         name: 'OpenAI',
-        chatModel: ChatOpenAI,
+        chat: ChatOpenAI,
         completer: OpenAICompleter,
         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',
+            chat: ChatWebLLM,
+            completer: WebLLMCompleter,
+            settingsSchema: WebLLMSettings,
+            instructions: WebLLMInstructions,
+            compatibilityCheck: webLLMCompatibilityCheck,
+            exposeChatModel: true
+        });
+        registry.providerChanged.connect(async (sender, role) => {
+            const { currentChatModel } = registry;
+            // 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 (registry.currentName(role) === 'WebLLM') {
+                // Leaving this check here, but it should never happen, this check is done in
+                // the provider registry, and the current name is set to 'None' if there is a
+                // compatibility error.
+                const compatibilityError = await webLLMCompatibilityCheck();
+                if (compatibilityError) {
+                    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 +173,7 @@ export const defaultProviderPlugins = AIProviders.map(provider => {
         }
     };
 });
+export const defaultProviderPlugins = [
+    webLLMProviderPlugin,
+    ...aiProviderPlugins
+];

package/lib/index.d.ts

@@ -1,6 +1,7 @@
 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';
+export * from './base-completer';

package/lib/index.js

@@ -1,16 +1,17 @@
 import { ActiveCellManager, buildChatSidebar, buildErrorWidget, ChatCommandRegistry, IChatCommandRegistry, InputToolbarRegistry } from '@jupyter/chat';
+import { ILayoutRestorer } from '@jupyterlab/application';
 import { IThemeManager } from '@jupyterlab/apputils';
 import { ICompletionProviderManager } from '@jupyterlab/completer';
 import { INotebookTracker } from '@jupyterlab/notebook';
 import { IRenderMimeRegistry } from '@jupyterlab/rendermime';
-import { ISettingConnector, ISettingRegistry } from '@jupyterlab/settingregistry';
+import { ISettingRegistry } from '@jupyterlab/settingregistry';
 import { IFormRendererRegistry } from '@jupyterlab/ui-components';
 import { ISecretsManager, SecretsManager } from 'jupyter-secrets-manager';
-import { ChatHandler } from './chat-handler';
+import { ChatHandler, welcomeMessage } from './chat-handler';
 import { CompletionProvider } from './completion-provider';
 import { defaultProviderPlugins } from './default-providers';
 import { AIProviderRegistry } from './provider';
-import { aiSettingsRenderer, SettingConnector } from './settings';
+import { aiSettingsRenderer, textArea } from './settings';
 import { IAIProviderRegistry, PLUGIN_IDS } from './tokens';
 import { stopItem } from './components/stop-button';
 const chatCommandRegistryPlugin = {
@@ -29,8 +30,13 @@ const chatPlugin = {
     description: 'LLM chat extension',
     autoStart: true,
     requires: [IAIProviderRegistry, IRenderMimeRegistry, IChatCommandRegistry],
-    optional: [INotebookTracker, ISettingRegistry, IThemeManager],
-    activate: async (app, providerRegistry, rmRegistry, chatCommandRegistry, notebookTracker, settingsRegistry, themeManager) => {
+    optional: [
+        INotebookTracker,
+        ISettingRegistry,
+        IThemeManager,
+        ILayoutRestorer
+    ],
+    activate: async (app, providerRegistry, rmRegistry, chatCommandRegistry, notebookTracker, settingsRegistry, themeManager, restorer) => {
         let activeCellManager = null;
         if (notebookTracker) {
             activeCellManager = new ActiveCellManager({
@@ -54,7 +60,7 @@ const chatPlugin = {
             chatHandler.config = { sendWithShiftEnter, enableCodeToolbar };
             chatHandler.personaName = personaName;
         }
-        Promise.all([app.restored, settingsRegistry === null || settingsRegistry === void 0 ? void 0 : settingsRegistry.load(chatPlugin.id)])
+        Promise.all([app.restored, settingsRegistry?.load(chatPlugin.id)])
             .then(([, settings]) => {
             if (!settings) {
                 console.warn('The SettingsRegistry is not loaded for the chat extension');
@@ -70,8 +76,8 @@ const chatPlugin = {
         const inputToolbarRegistry = InputToolbarRegistry.defaultToolbarRegistry();
         const stopButton = stopItem(() => chatHandler.stopStreaming());
         inputToolbarRegistry.addItem('stop', stopButton);
-        chatHandler.writersChanged.connect((_, users) => {
-            if (users.filter(user => user.username === chatHandler.personaName).length) {
+        chatHandler.writersChanged.connect((_, writers) => {
+            if (writers.filter(writer => writer.user.username === chatHandler.personaName).length) {
                 inputToolbarRegistry.hide('send');
                 inputToolbarRegistry.show('stop');
             }
@@ -86,15 +92,19 @@ const chatPlugin = {
                 themeManager,
                 rmRegistry,
                 chatCommandRegistry,
-                inputToolbarRegistry
+                inputToolbarRegistry,
+                welcomeMessage: welcomeMessage(providerRegistry.providers)
             });
-            chatWidget.title.caption = 'Jupyterlite AI Chat';
         }
         catch (e) {
             chatWidget = buildErrorWidget(themeManager);
         }
+        chatWidget.title.caption = 'Jupyterlite AI Chat';
+        chatWidget.id = '@jupyterlite/ai:chat-widget';
         app.shell.add(chatWidget, 'left', { rank: 2000 });
-        console.log('Chat extension initialized');
+        if (restorer) {
+            restorer.add(chatWidget, chatWidget.id);
+        }
     }
 };
 const completerPlugin = {
@@ -113,33 +123,43 @@ const providerRegistryPlugin = SecretsManager.sign(PLUGIN_IDS.providerRegistry,
     id: PLUGIN_IDS.providerRegistry,
     autoStart: true,
     requires: [IFormRendererRegistry, ISettingRegistry],
-    optional: [IRenderMimeRegistry, ISecretsManager, ISettingConnector],
+    optional: [IRenderMimeRegistry, ISecretsManager],
     provides: IAIProviderRegistry,
-    activate: (app, editorRegistry, settingRegistry, rmRegistry, secretsManager, settingConnector) => {
+    activate: (app, editorRegistry, settingRegistry, rmRegistry, secretsManager) => {
         const providerRegistry = new AIProviderRegistry({
             token,
             secretsManager
         });
-        editorRegistry.addRenderer(`${PLUGIN_IDS.providerRegistry}.AIprovider`, aiSettingsRenderer({
+        editorRegistry.addRenderer(`${PLUGIN_IDS.providerRegistry}.AIproviders`, aiSettingsRenderer({
             providerRegistry,
             secretsToken: token,
             rmRegistry,
-            secretsManager,
-            settingConnector
+            secretsManager
         }));
         settingRegistry
             .load(providerRegistryPlugin.id)
             .then(settings => {
+            if (!secretsManager) {
+                delete settings.schema.properties?.['UseSecretsManager'];
+            }
             const updateProvider = () => {
-                var _a;
-                // Update the settings to the AI providers.
-                const providerSettings = ((_a = settings.get('AIprovider').composite) !== null && _a !== void 0 ? _a : {
-                    provider: 'None'
-                });
-                providerRegistry.setProvider({
-                    name: providerSettings.provider,
-                    settings: providerSettings
-                });
+                // Get the Ai provider settings.
+                const providerSettings = settings.get('AIproviders')
+                    .composite;
+                // Update completer provider.
+                if (Object.keys(providerSettings).includes('completer')) {
+                    providerRegistry.setCompleterProvider(providerSettings['completer']);
+                }
+                else {
+                    providerRegistry.setCompleterProvider({});
+                }
+                // Update chat provider.
+                if (Object.keys(providerSettings).includes('chat')) {
+                    providerRegistry.setChatProvider(providerSettings['chat']);
+                }
+                else {
+                    providerRegistry.setChatProvider({});
+                }
             };
             settings.changed.connect(() => updateProvider());
             updateProvider();
@@ -150,23 +170,47 @@ const providerRegistryPlugin = SecretsManager.sign(PLUGIN_IDS.providerRegistry,
         return providerRegistry;
     }
 }));
-/**
- * Provides the settings connector as a separate plugin to allow for alternative
- * implementations that may want to fetch settings from a different source or
- * endpoint.
- */
-const settingsConnector = {
-    id: PLUGIN_IDS.settingsConnector,
-    description: 'Provides a settings connector which does not save passwords.',
+const systemPromptsPlugin = {
+    id: PLUGIN_IDS.systemPrompts,
     autoStart: true,
-    provides: ISettingConnector,
-    activate: (app) => new SettingConnector(app.serviceManager.settings)
+    requires: [IAIProviderRegistry, ISettingRegistry],
+    optional: [IFormRendererRegistry],
+    activate: (app, providerRegistry, settingsRegistry, editorRegistry) => {
+        // Set textarea renderer for the prompt setting.
+        editorRegistry?.addRenderer(`${PLUGIN_IDS.systemPrompts}.chatSystemPrompt`, textArea);
+        editorRegistry?.addRenderer(`${PLUGIN_IDS.systemPrompts}.completionSystemPrompt`, textArea);
+        /**
+         * Update the prompts in the provider registry.
+         */
+        function loadSetting(setting) {
+            providerRegistry.chatSystemPrompt = setting.get('chatSystemPrompt')
+                .composite;
+            providerRegistry.completerSystemPrompt = setting.get('completionSystemPrompt').composite;
+        }
+        Promise.all([
+            app.restored,
+            settingsRegistry?.load(PLUGIN_IDS.systemPrompts)
+        ])
+            .then(([, settings]) => {
+            if (!settings) {
+                console.warn('The SettingsRegistry is not loaded for the chat extension');
+                return;
+            }
+            loadSetting(settings);
+            settings.changed.connect(loadSetting);
+        })
+            .catch(reason => {
+            console.error(`Something went wrong when reading the settings.\n${reason}`);
+        });
+    }
 };
 export default [
     providerRegistryPlugin,
     chatCommandRegistryPlugin,
     chatPlugin,
     completerPlugin,
-    settingsConnector,
+    systemPromptsPlugin,
     ...defaultProviderPlugins
 ];
+export { IAIProviderRegistry } from './tokens';
+export * from './base-completer';

package/lib/provider.d.ts

@@ -1,13 +1,10 @@
 import { BaseLanguageModel } from '@langchain/core/language_models/base';
-import { BaseChatModel } from '@langchain/core/language_models/chat_models';
-import { ISignal } from '@lumino/signaling';
 import { ReadonlyPartialJSONObject } from '@lumino/coreutils';
+import { ISignal } from '@lumino/signaling';
 import { JSONSchema7 } from 'json-schema';
 import { ISecretsManager } from 'jupyter-secrets-manager';
-import { IBaseCompleter } from './base-completer';
-import { IAIProvider, IAIProviderRegistry, ISetProviderOptions } from './tokens';
-export declare const chatSystemPrompt: (options: AIProviderRegistry.IPromptOptions) => string;
-export declare const COMPLETION_SYSTEM_PROMPT = "\nYou are an application built to provide helpful code completion suggestions.\nYou should only produce code. Keep comments to minimum, use the\nprogramming language comment syntax. Produce clean code.\nThe code is written in JupyterLab, a data analysis and code development\nenvironment which can execute code extended with additional syntax for\ninteractive features, such as magics.\nOnly give raw strings back, do not format the response using backticks.\nThe output should be a single string, and should correspond to what a human users\nwould write.\nDo not include the prompt in the output, only the string that should be appended to the current input.\n";
+import { IAIProvider, IAIProviderRegistry, ModelRole } from './tokens';
+import { AIChatModel, AICompleter } from './types/ai-model';
 export declare class AIProviderRegistry implements IAIProviderRegistry {
     /**
      * The constructor of the provider registry.
@@ -24,15 +21,25 @@ export declare class AIProviderRegistry implements IAIProviderRegistry {
     /**
      * Get the current provider name.
      */
-    get currentName(): string;
+    currentName(role: ModelRole): string;
+    /**
+     * Get the current AICompleter.
+     */
+    get currentCompleter(): AICompleter | null;
     /**
-     * Get the current completer of the completion provider.
+     * Getter/setter for the completer system prompt.
      */
-    get currentCompleter(): IBaseCompleter | null;
+    get completerSystemPrompt(): string;
+    set completerSystemPrompt(value: string);
     /**
-     * Get the current llm chat model.
+     * Get the current AIChatModel.
      */
-    get currentChatModel(): BaseChatModel | null;
+    get currentChatModel(): AIChatModel | null;
+    /**
+     * Getter/setter for the chat system prompt.
+     */
+    get chatSystemPrompt(): string;
+    set chatSystemPrompt(value: string);
     /**
      * Get the settings schema of a given provider.
      */
@@ -41,39 +48,58 @@ export declare class AIProviderRegistry implements IAIProviderRegistry {
      * Get the instructions of a given provider.
      */
     getInstructions(provider: string): string | undefined;
+    /**
+     * Get the compatibility check function of a given provider.
+     */
+    getCompatibilityCheck(provider: string): (() => Promise<string | null>) | undefined;
     /**
      * Format an error message from the current provider.
      */
     formatErrorMessage(error: any): string;
     /**
-     * Get the current chat error;
+     * Get/set the current chat error;
      */
     get chatError(): string;
+    private set chatError(value);
     /**
-     * get the current completer error.
+     * Get/set the current completer error.
      */
     get completerError(): string;
+    private set completerError(value);
     /**
-     * Set the providers (chat model and completer).
-     * Creates the providers if the name has changed, otherwise only updates their config.
+     * A function to emit a notification error.
+     */
+    private _emitErrorNotification;
+    /**
+     * 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.
+     */
+    setCompleterProvider(settings: ReadonlyPartialJSONObject): Promise<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.
      */
-    setProvider(options: ISetProviderOptions): Promise<void>;
+    setChatProvider(settings: ReadonlyPartialJSONObject): Promise<void>;
     /**
      * A signal emitting when the provider or its settings has changed.
      */
-    get providerChanged(): ISignal<IAIProviderRegistry, void>;
+    get providerChanged(): ISignal<IAIProviderRegistry, ModelRole>;
+    /**
+     * Build a new settings object containing the secrets.
+     */
+    private _buildFullSettings;
     private _secretsManager;
-    private _currentProvider;
-    private _completer;
-    private _chatModel;
-    private _name;
     private _providerChanged;
     private _chatError;
     private _completerError;
-    private _providers;
+    private _notifications;
     private _deferredProvider;
+    private _chatPrompt;
+    private _completerPrompt;
 }
 export declare namespace AIProviderRegistry {
     /**