@jupyterlite/ai 0.6.2 → 0.7.0

package/src/default-providers/index.ts

@@ -2,28 +2,44 @@ import {
   JupyterFrontEnd,
   JupyterFrontEndPlugin
 } from '@jupyterlab/application';
+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 { IAIProvider, 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 { IAIProvider, IAIProviderRegistry } from '../tokens';
 
 // Build the AIProvider list
 const AIProviders: IAIProvider[] = [
@@ -41,7 +57,8 @@ const AIProviders: IAIProvider[] = [
     chatModel: ChromeAI,
     completer: ChromeCompleter,
     instructions: ChromeAIInstructions,
-    settingsSchema: ChromeAISettings
+    settingsSchema: ChromeAISettings,
+    compatibilityCheck: chromeAICompatibilityCheck
   },
   {
     name: 'MistralAI',
@@ -50,6 +67,13 @@ const AIProviders: IAIProvider[] = [
     instructions: MistralAIInstructions,
     settingsSchema: MistralAISettings
   },
+  {
+    name: 'Ollama',
+    chatModel: ChatOllama,
+    completer: OllamaCompleter,
+    instructions: OllamaInstructions,
+    settingsSchema: OllamaAISettings
+  },
   {
     name: 'OpenAI',
     chatModel: ChatOpenAI,
@@ -58,14 +82,127 @@ const AIProviders: IAIProvider[] = [
   }
 ];
 
-export const defaultProviderPlugins: JupyterFrontEndPlugin<void>[] =
-  AIProviders.map(provider => {
-    return {
-      id: `@jupyterlite/ai:${provider.name}`,
-      autoStart: true,
-      requires: [IAIProviderRegistry],
-      activate: (app: JupyterFrontEnd, registry: IAIProviderRegistry) => {
-        registry.add(provider);
+/**
+ * Register the WebLLM provider in a separate plugin since it creates notifications
+ * when the model is changed in the settings.
+ */
+const webLLMProviderPlugin: JupyterFrontEndPlugin<void> = {
+  id: '@jupyterlite/ai:webllm',
+  description: 'Register the WebLLM provider',
+  autoStart: true,
+  requires: [IAIProviderRegistry],
+  activate: (app: JupyterFrontEnd, registry: IAIProviderRegistry) => {
+    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 as ChatWebLLM;
+        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,
+    requires: [IAIProviderRegistry],
+    activate: (app: JupyterFrontEnd, registry: IAIProviderRegistry) => {
+      registry.add(provider);
+    }
+  };
+});
+
+export const defaultProviderPlugins: JupyterFrontEndPlugin<void>[] = [
+  webLLMProviderPlugin,
+  ...aiProviderPlugins
+];

package/src/index.ts

@@ -15,19 +15,16 @@ import { ReactWidget, 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 { ReadonlyPartialJSONObject } from '@lumino/coreutils';
 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 } from './settings';
 import { IAIProviderRegistry, PLUGIN_IDS } from './tokens';
 import { stopItem } from './components/stop-button';
 
@@ -109,9 +106,11 @@ const chatPlugin: JupyterFrontEndPlugin<void> = {
     const stopButton = stopItem(() => chatHandler.stopStreaming());
     inputToolbarRegistry.addItem('stop', stopButton);
 
-    chatHandler.writersChanged.connect((_, users) => {
+    chatHandler.writersChanged.connect((_, writers) => {
       if (
-        users.filter(user => user.username === chatHandler.personaName).length
+        writers.filter(
+          writer => writer.user.username === chatHandler.personaName
+        ).length
       ) {
         inputToolbarRegistry.hide('send');
         inputToolbarRegistry.show('stop');
@@ -127,7 +126,8 @@ const chatPlugin: JupyterFrontEndPlugin<void> = {
         themeManager,
         rmRegistry,
         chatCommandRegistry,
-        inputToolbarRegistry
+        inputToolbarRegistry,
+        welcomeMessage: welcomeMessage(providerRegistry.providers)
       });
       chatWidget.title.caption = 'Jupyterlite AI Chat';
     } catch (e) {
@@ -162,15 +162,14 @@ const providerRegistryPlugin: JupyterFrontEndPlugin<IAIProviderRegistry> =
     id: PLUGIN_IDS.providerRegistry,
     autoStart: true,
     requires: [IFormRendererRegistry, ISettingRegistry],
-    optional: [IRenderMimeRegistry, ISecretsManager, ISettingConnector],
+    optional: [IRenderMimeRegistry, ISecretsManager],
     provides: IAIProviderRegistry,
     activate: (
       app: JupyterFrontEnd,
       editorRegistry: IFormRendererRegistry,
       settingRegistry: ISettingRegistry,
       rmRegistry?: IRenderMimeRegistry,
-      secretsManager?: ISecretsManager,
-      settingConnector?: ISettingConnector
+      secretsManager?: ISecretsManager
     ): IAIProviderRegistry => {
       const providerRegistry = new AIProviderRegistry({
         token,
@@ -183,14 +182,16 @@ const providerRegistryPlugin: JupyterFrontEndPlugin<IAIProviderRegistry> =
           providerRegistry,
           secretsToken: token,
           rmRegistry,
-          secretsManager,
-          settingConnector
+          secretsManager
         })
       );
 
       settingRegistry
         .load(providerRegistryPlugin.id)
         .then(settings => {
+          if (!secretsManager) {
+            delete settings.schema.properties?.['UseSecretsManager'];
+          }
           const updateProvider = () => {
             // Update the settings to the AI providers.
             const providerSettings = (settings.get('AIprovider').composite ?? {
@@ -216,25 +217,12 @@ const providerRegistryPlugin: JupyterFrontEndPlugin<IAIProviderRegistry> =
     }
   }));
 
-/**
- * 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: JupyterFrontEndPlugin<ISettingConnector> = {
-  id: PLUGIN_IDS.settingsConnector,
-  description: 'Provides a settings connector which does not save passwords.',
-  autoStart: true,
-  provides: ISettingConnector,
-  activate: (app: JupyterFrontEnd) =>
-    new SettingConnector(app.serviceManager.settings)
-};
-
 export default [
   providerRegistryPlugin,
   chatCommandRegistryPlugin,
   chatPlugin,
   completerPlugin,
-  settingsConnector,
   ...defaultProviderPlugins
 ];
+
+export { IAIProviderRegistry } from './tokens';

package/src/provider.ts

@@ -1,3 +1,7 @@
+import {
+  CompletionHandler,
+  IInlineCompletionContext
+} from '@jupyterlab/completer';
 import { BaseLanguageModel } from '@langchain/core/language_models/base';
 import { BaseChatModel } from '@langchain/core/language_models/chat_models';
 import { ISignal, Signal } from '@lumino/signaling';
@@ -14,6 +18,7 @@ import {
   ISetProviderOptions,
   PLUGIN_IDS
 } from './tokens';
+import { AIChatModel, AICompleter } from './types/ai-model';
 
 const SECRETS_NAMESPACE = PLUGIN_IDS.providerRegistry;
 
@@ -42,9 +47,10 @@ The code is written in JupyterLab, a data analysis and code development
 environment which can execute code extended with additional syntax for
 interactive features, such as magics.
 Only give raw strings back, do not format the response using backticks.
-The output should be a single string, and should correspond to what a human users
-would write.
+The output should be a single string, and should only contain the code that will complete the
+give code passed as input, no explanation whatsoever.
 Do not include the prompt in the output, only the string that should be appended to the current input.
+Here is the code to complete:
 `;
 
 export class AIProviderRegistry implements IAIProviderRegistry {
@@ -60,19 +66,19 @@ export class AIProviderRegistry implements IAIProviderRegistry {
    * Get the list of provider names.
    */
   get providers(): string[] {
-    return Array.from(this._providers.keys());
+    return Array.from(Private.providers.keys());
   }
 
   /**
    * Add a new provider.
    */
   add(provider: IAIProvider): void {
-    if (this._providers.has(provider.name)) {
+    if (Private.providers.has(provider.name)) {
       throw new Error(
         `A AI provider named '${provider.name}' is already registered`
       );
     }
-    this._providers.set(provider.name, provider);
+    Private.providers.set(provider.name, provider);
 
     // Set the provider if the loading has been deferred.
     if (provider.name === this._deferredProvider?.name) {
@@ -84,34 +90,57 @@ export class AIProviderRegistry implements IAIProviderRegistry {
    * Get the current provider name.
    */
   get currentName(): string {
-    return this._name;
+    return Private.getName();
   }
 
   /**
-   * Get the current completer of the completion provider.
+   * Get the current AICompleter.
    */
-  get currentCompleter(): IBaseCompleter | null {
-    if (this._name === 'None') {
+  get currentCompleter(): AICompleter | null {
+    if (Private.getName() === 'None') {
+      return null;
+    }
+    const completer = Private.getCompleter();
+    if (completer === null) {
       return null;
     }
-    return this._completer;
+    return {
+      fetch: (
+        request: CompletionHandler.IRequest,
+        context: IInlineCompletionContext
+      ) => completer.fetch(request, context)
+    };
   }
 
   /**
-   * Get the current llm chat model.
+   * Get the current AIChatModel.
    */
-  get currentChatModel(): BaseChatModel | null {
-    if (this._name === 'None') {
+  get currentChatModel(): AIChatModel | null {
+    if (Private.getName() === 'None') {
+      return null;
+    }
+    const currentProvider = Private.providers.get(Private.getName()) ?? null;
+
+    const chatModel = Private.getChatModel();
+    if (chatModel === null) {
       return null;
     }
-    return this._chatModel;
+    if (currentProvider?.exposeChatModel ?? false) {
+      // Expose the full chat model if expected.
+      return chatModel as AIChatModel;
+    }
+
+    // Otherwise, we create a reduced AIChatModel interface.
+    return {
+      stream: (input: any, options?: any) => chatModel.stream(input, options)
+    };
   }
 
   /**
    * Get the settings schema of a given provider.
    */
   getSettingsSchema(provider: string): JSONSchema7 {
-    return (this._providers.get(provider)?.settingsSchema?.properties ||
+    return (Private.providers.get(provider)?.settingsSchema?.properties ||
       {}) as JSONSchema7;
   }
 
@@ -119,15 +148,25 @@ export class AIProviderRegistry implements IAIProviderRegistry {
    * Get the instructions of a given provider.
    */
   getInstructions(provider: string): string | undefined {
-    return this._providers.get(provider)?.instructions;
+    return Private.providers.get(provider)?.instructions;
+  }
+
+  /**
+   * Get the compatibility check function of a given provider.
+   */
+  getCompatibilityCheck(
+    provider: string
+  ): (() => Promise<string | null>) | undefined {
+    return Private.providers.get(provider)?.compatibilityCheck;
   }
 
   /**
    * Format an error message from the current provider.
    */
   formatErrorMessage(error: any): string {
-    if (this._currentProvider?.errorMessage) {
-      return this._currentProvider?.errorMessage(error);
+    const currentProvider = Private.providers.get(Private.getName()) ?? null;
+    if (currentProvider?.errorMessage) {
+      return currentProvider?.errorMessage(error);
     }
     if (error.message) {
       return error.message;
@@ -143,7 +182,7 @@ export class AIProviderRegistry implements IAIProviderRegistry {
   }
 
   /**
-   * get the current completer error.
+   * Get the current completer error.
    */
   get completerError(): string {
     return this._completerError;
@@ -157,8 +196,8 @@ export class AIProviderRegistry implements IAIProviderRegistry {
    */
   async setProvider(options: ISetProviderOptions): Promise<void> {
     const { name, settings } = options;
-    this._currentProvider = this._providers.get(name) ?? null;
-    if (this._currentProvider === null) {
+    const currentProvider = Private.providers.get(name) ?? null;
+    if (currentProvider === null) {
       // The current provider may not be loaded when the settings are first loaded.
       // Let's defer the provider loading.
       this._deferredProvider = options;
@@ -166,6 +205,23 @@ export class AIProviderRegistry implements IAIProviderRegistry {
       this._deferredProvider = null;
     }
 
+    const compatibilityCheck = this.getCompatibilityCheck(name);
+    if (compatibilityCheck !== undefined) {
+      const error = await compatibilityCheck();
+      if (error !== null) {
+        this._chatError = error.trim();
+        this._completerError = error.trim();
+        Private.setName('None');
+        this._providerChanged.emit();
+        return;
+      }
+    }
+
+    if (name === 'None') {
+      this._chatError = '';
+      this._completerError = '';
+    }
+
     // Build a new settings object containing the secrets.
     const fullSettings: IDict = {};
     for (const key of Object.keys(settings)) {
@@ -184,33 +240,37 @@ export class AIProviderRegistry implements IAIProviderRegistry {
       fullSettings[key] = settings[key];
     }
 
-    if (this._currentProvider?.completer !== undefined) {
+    if (currentProvider?.completer !== undefined) {
       try {
-        this._completer = new this._currentProvider.completer({
-          settings: fullSettings
-        });
+        Private.setCompleter(
+          new currentProvider.completer({
+            settings: fullSettings
+          })
+        );
         this._completerError = '';
       } catch (e: any) {
         this._completerError = e.message;
       }
     } else {
-      this._completer = null;
+      Private.setCompleter(null);
     }
 
-    if (this._currentProvider?.chatModel !== undefined) {
+    if (currentProvider?.chatModel !== undefined) {
       try {
-        this._chatModel = new this._currentProvider.chatModel({
-          ...fullSettings
-        });
+        Private.setChatModel(
+          new currentProvider.chatModel({
+            ...fullSettings
+          })
+        );
         this._chatError = '';
       } catch (e: any) {
         this._chatError = e.message;
-        this._chatModel = null;
+        Private.setChatModel(null);
       }
     } else {
-      this._chatModel = null;
+      Private.setChatModel(null);
     }
-    this._name = name;
+    Private.setName(name);
     this._providerChanged.emit();
   }
 
@@ -222,14 +282,9 @@ export class AIProviderRegistry implements IAIProviderRegistry {
   }
 
   private _secretsManager: ISecretsManager | null;
-  private _currentProvider: IAIProvider | null = null;
-  private _completer: IBaseCompleter | null = null;
-  private _chatModel: BaseChatModel | null = null;
-  private _name: string = 'None';
   private _providerChanged = new Signal<IAIProviderRegistry, void>(this);
   private _chatError: string = '';
   private _completerError: string = '';
-  private _providers = new Map<string, IAIProvider>();
   private _deferredProvider: ISetProviderOptions | null = null;
 }
 
@@ -303,21 +358,53 @@ export namespace AIProviderRegistry {
 
 namespace Private {
   /**
-   * The token to use with the secrets manager.
+   * The token to use with the secrets manager, setter and getter.
    */
   let secretsToken: symbol;
+  export function setToken(value: symbol): void {
+    secretsToken = value;
+  }
+  export function getToken(): symbol {
+    return secretsToken;
+  }
 
   /**
-   * Set of the token.
+   * The providers map, in private namespace to prevent updating the 'exposeChatModel'
+   * flag.
    */
-  export function setToken(value: symbol): void {
-    secretsToken = value;
+  export const providers = new Map<string, IAIProvider>();
+
+  /**
+   * The name of the current provider, setter and getter.
+   * It is in a private namespace to prevent updating it without updating the models.
+   */
+  let name: string = 'None';
+  export function setName(value: string): void {
+    name = value;
+  }
+  export function getName(): string {
+    return name;
   }
 
   /**
-   * get the token.
+   * The chat model setter and getter.
    */
-  export function getToken(): symbol {
-    return secretsToken;
+  let chatModel: BaseChatModel | null = null;
+  export function setChatModel(model: BaseChatModel | null): void {
+    chatModel = model;
+  }
+  export function getChatModel(): BaseChatModel | null {
+    return chatModel;
+  }
+
+  /**
+   * The completer setter and getter.
+   */
+  let completer: IBaseCompleter | null = null;
+  export function setCompleter(model: IBaseCompleter | null): void {
+    completer = model;
+  }
+  export function getCompleter(): IBaseCompleter | null {
+    return completer;
   }
 }

package/src/settings/index.ts

@@ -1,3 +1,2 @@
 export * from './panel';
-export * from './settings-connector';
 export * from './utils';