@jupyterlite/ai 0.3.0 → 0.4.0

package/schema/provider-registry.json

@@ -0,0 +1,17 @@
+{
+  "title": "AI provider",
+  "description": "Provider registry settings",
+  "jupyter.lab.setting-icon": "@jupyterlite/ai:jupyternaut-lite",
+  "jupyter.lab.setting-icon-label": "JupyterLite AI Chat",
+  "type": "object",
+  "properties": {
+    "AIprovider": {
+      "type": "object",
+      "title": "AI provider",
+      "description": "The AI provider configuration",
+      "default": {},
+      "additionalProperties": true
+    }
+  },
+  "additionalProperties": false
+}

package/src/chat-handler.ts

@@ -17,9 +17,8 @@ import {
   SystemMessage
 } from '@langchain/core/messages';
 import { UUID } from '@lumino/coreutils';
-import { getErrorMessage } from './llm-models';
 import { chatSystemPrompt } from './provider';
-import { IAIProvider } from './token';
+import { IAIProviderRegistry } from './tokens';
 import { jupyternautLiteIcon } from './icons';
 
 /**
@@ -37,17 +36,21 @@ export type ConnectionMessage = {
 export class ChatHandler extends ChatModel {
   constructor(options: ChatHandler.IOptions) {
     super(options);
-    this._aiProvider = options.aiProvider;
-    this._prompt = chatSystemPrompt({ provider_name: this._aiProvider.name });
+    this._providerRegistry = options.providerRegistry;
+    this._prompt = chatSystemPrompt({
+      provider_name: this._providerRegistry.currentName
+    });
 
-    this._aiProvider.modelChange.connect(() => {
-      this._errorMessage = this._aiProvider.chatError;
-      this._prompt = chatSystemPrompt({ provider_name: this._aiProvider.name });
+    this._providerRegistry.providerChanged.connect(() => {
+      this._errorMessage = this._providerRegistry.chatError;
+      this._prompt = chatSystemPrompt({
+        provider_name: this._providerRegistry.currentName
+      });
     });
   }
 
   get provider(): BaseChatModel | null {
-    return this._aiProvider.chatModel;
+    return this._providerRegistry.currentChatModel;
   }
 
   /**
@@ -95,7 +98,7 @@ export class ChatHandler extends ChatModel {
     };
     this.messageAdded(msg);
 
-    if (this._aiProvider.chatModel === null) {
+    if (this._providerRegistry.currentChatModel === null) {
       const errorMsg: IChatMessage = {
         id: UUID.uuid4(),
         body: `**${this._errorMessage ? this._errorMessage : this._defaultErrorMessage}**`,
@@ -134,7 +137,7 @@ export class ChatHandler extends ChatModel {
     let content = '';
 
     try {
-      for await (const chunk of await this._aiProvider.chatModel.stream(
+      for await (const chunk of await this._providerRegistry.currentChatModel.stream(
         messages
       )) {
         content += chunk.content ?? chunk;
@@ -144,7 +147,7 @@ export class ChatHandler extends ChatModel {
       this._history.messages.push(botMsg);
       return true;
     } catch (reason) {
-      const error = getErrorMessage(this._aiProvider.name, reason);
+      const error = this._providerRegistry.formatErrorMessage(reason);
       const errorMsg: IChatMessage = {
         id: UUID.uuid4(),
         body: `**${error}**`,
@@ -171,7 +174,7 @@ export class ChatHandler extends ChatModel {
     super.messageAdded(message);
   }
 
-  private _aiProvider: IAIProvider;
+  private _providerRegistry: IAIProviderRegistry;
   private _personaName = 'AI';
   private _prompt: string;
   private _errorMessage: string = '';
@@ -181,6 +184,6 @@ export class ChatHandler extends ChatModel {
 
 export namespace ChatHandler {
   export interface IOptions extends ChatModel.IOptions {
-    aiProvider: IAIProvider;
+    providerRegistry: IAIProviderRegistry;
   }
 }

package/src/completion-provider.ts

@@ -3,10 +3,9 @@ import {
   IInlineCompletionContext,
   IInlineCompletionProvider
 } from '@jupyterlab/completer';
-import { BaseLanguageModel } from '@langchain/core/language_models/base';
-import { ReadonlyPartialJSONObject } from '@lumino/coreutils';
 
-import { getCompleter, IBaseCompleter, BaseCompleter } from './llm-models';
+import { IBaseCompleter } from './llm-models';
+import { IAIProviderRegistry } from './tokens';
 
 /**
  * The generic completion provider to register to the completion provider manager.
@@ -15,67 +14,44 @@ export class CompletionProvider implements IInlineCompletionProvider {
   readonly identifier = '@jupyterlite/ai';
 
   constructor(options: CompletionProvider.IOptions) {
-    const { name, settings } = options;
+    this._providerRegistry = options.providerRegistry;
     this._requestCompletion = options.requestCompletion;
-    this.setCompleter(name, settings);
-  }
 
-  /**
-   * Set the completer.
-   *
-   * @param name - the name of the completer.
-   * @param settings - The settings associated to the completer.
-   */
-  setCompleter(name: string, settings: ReadonlyPartialJSONObject) {
-    try {
-      this._completer = getCompleter(name, settings);
-      if (this._completer) {
-        this._completer.requestCompletion = this._requestCompletion;
+    this._providerRegistry.providerChanged.connect(() => {
+      if (this.completer) {
+        this.completer.requestCompletion = this._requestCompletion;
       }
-      this._name = this._completer === null ? 'None' : name;
-    } catch (e: any) {
-      this._completer = null;
-      this._name = 'None';
-      throw e;
-    }
+    });
   }
 
   /**
    * Get the current completer name.
    */
   get name(): string {
-    return this._name;
+    return this._providerRegistry.currentName;
   }
 
   /**
    * Get the current completer.
    */
   get completer(): IBaseCompleter | null {
-    return this._completer;
-  }
-
-  /**
-   * Get the LLM completer.
-   */
-  get llmCompleter(): BaseLanguageModel | null {
-    return this._completer?.provider || null;
+    return this._providerRegistry.currentCompleter;
   }
 
   async fetch(
     request: CompletionHandler.IRequest,
     context: IInlineCompletionContext
   ) {
-    return this._completer?.fetch(request, context);
+    return this.completer?.fetch(request, context);
   }
 
-  private _name: string = 'None';
+  private _providerRegistry: IAIProviderRegistry;
   private _requestCompletion: () => void;
-  private _completer: IBaseCompleter | null = null;
 }
 
 export namespace CompletionProvider {
-  export interface IOptions extends BaseCompleter.IOptions {
-    name: string;
+  export interface IOptions {
+    providerRegistry: IAIProviderRegistry;
     requestCompletion: () => void;
   }
 }

package/src/index.ts

@@ -16,12 +16,16 @@ import { ICompletionProviderManager } from '@jupyterlab/completer';
 import { INotebookTracker } from '@jupyterlab/notebook';
 import { IRenderMimeRegistry } from '@jupyterlab/rendermime';
 import { ISettingRegistry } from '@jupyterlab/settingregistry';
+import { IFormRendererRegistry } from '@jupyterlab/ui-components';
+import { ReadonlyPartialJSONObject } from '@lumino/coreutils';
 
 import { ChatHandler } from './chat-handler';
-import { getSettings } from './llm-models';
-import { AIProvider } from './provider';
+import { CompletionProvider } from './completion-provider';
+import { AIProviders } from './llm-models';
+import { AIProviderRegistry } from './provider';
+import { aiSettingsRenderer } from './settings/panel';
 import { renderSlashCommandOption } from './slash-commands';
-import { IAIProvider } from './token';
+import { IAIProviderRegistry } from './tokens';
 
 const autocompletionRegistryPlugin: JupyterFrontEndPlugin<IAutocompletionRegistry> =
   {
@@ -54,11 +58,11 @@ const chatPlugin: JupyterFrontEndPlugin<void> = {
   id: '@jupyterlite/ai:chat',
   description: 'LLM chat extension',
   autoStart: true,
-  requires: [IAIProvider, IRenderMimeRegistry, IAutocompletionRegistry],
+  requires: [IAIProviderRegistry, IRenderMimeRegistry, IAutocompletionRegistry],
   optional: [INotebookTracker, ISettingRegistry, IThemeManager],
   activate: async (
     app: JupyterFrontEnd,
-    aiProvider: IAIProvider,
+    providerRegistry: IAIProviderRegistry,
     rmRegistry: IRenderMimeRegistry,
     autocompletionRegistry: IAutocompletionRegistry,
     notebookTracker: INotebookTracker | null,
@@ -74,8 +78,8 @@ const chatPlugin: JupyterFrontEndPlugin<void> = {
     }
 
     const chatHandler = new ChatHandler({
-      aiProvider: aiProvider,
-      activeCellManager: activeCellManager
+      providerRegistry,
+      activeCellManager
     });
 
     let sendWithShiftEnter = false;
@@ -129,50 +133,53 @@ const chatPlugin: JupyterFrontEndPlugin<void> = {
   }
 };
 
-const aiProviderPlugin: JupyterFrontEndPlugin<IAIProvider> = {
-  id: '@jupyterlite/ai:ai-provider',
+const completerPlugin: JupyterFrontEndPlugin<void> = {
+  id: '@jupyterlite/ai:completer',
   autoStart: true,
-  requires: [ICompletionProviderManager, ISettingRegistry],
-  provides: IAIProvider,
+  requires: [IAIProviderRegistry, ICompletionProviderManager],
   activate: (
     app: JupyterFrontEnd,
-    manager: ICompletionProviderManager,
-    settingRegistry: ISettingRegistry
-  ): IAIProvider => {
-    const aiProvider = new AIProvider({
-      completionProviderManager: manager,
+    providerRegistry: IAIProviderRegistry,
+    manager: ICompletionProviderManager
+  ): void => {
+    const completer = new CompletionProvider({
+      providerRegistry,
       requestCompletion: () => app.commands.execute('inline-completer:invoke')
     });
+    manager.registerInlineProvider(completer);
+  }
+};
 
-    let currentProvider = 'None';
+const providerRegistryPlugin: JupyterFrontEndPlugin<IAIProviderRegistry> = {
+  id: '@jupyterlite/ai:provider-registry',
+  autoStart: true,
+  requires: [IFormRendererRegistry, ISettingRegistry],
+  optional: [IRenderMimeRegistry],
+  provides: IAIProviderRegistry,
+  activate: (
+    app: JupyterFrontEnd,
+    editorRegistry: IFormRendererRegistry,
+    settingRegistry: ISettingRegistry,
+    rmRegistry?: IRenderMimeRegistry
+  ): IAIProviderRegistry => {
+    const providerRegistry = new AIProviderRegistry();
+
+    editorRegistry.addRenderer(
+      '@jupyterlite/ai:provider-registry.AIprovider',
+      aiSettingsRenderer({ providerRegistry, rmRegistry })
+    );
     settingRegistry
-      .load(aiProviderPlugin.id)
+      .load(providerRegistryPlugin.id)
       .then(settings => {
         const updateProvider = () => {
-          const provider = settings.get('provider').composite as string;
-          if (provider !== currentProvider) {
-            // Update the settings panel.
-            currentProvider = provider;
-            const settingsProperties = settings.schema.properties;
-            if (settingsProperties) {
-              const schemaKeys = Object.keys(settingsProperties);
-              schemaKeys.forEach(key => {
-                if (key !== 'provider') {
-                  delete settings.schema.properties?.[key];
-                }
-              });
-              const properties = getSettings(provider);
-              if (properties === null) {
-                return;
-              }
-              Object.entries(properties).forEach(([name, value], index) => {
-                settingsProperties[name] = value as ISettingRegistry.IProperty;
-              });
-            }
-          }
-
           // Update the settings to the AI providers.
-          aiProvider.setModels(provider, settings.composite);
+          const providerSettings = (settings.get('AIprovider').composite ?? {
+            provider: 'None'
+          }) as ReadonlyPartialJSONObject;
+          providerRegistry.setProvider(
+            providerSettings.provider as string,
+            providerSettings
+          );
         };
 
         settings.changed.connect(() => updateProvider());
@@ -180,13 +187,21 @@ const aiProviderPlugin: JupyterFrontEndPlugin<IAIProvider> = {
       })
       .catch(reason => {
         console.error(
-          `Failed to load settings for ${aiProviderPlugin.id}`,
+          `Failed to load settings for ${providerRegistryPlugin.id}`,
           reason
         );
       });
 
-    return aiProvider;
+    // Initialize the registry with the default providers
+    AIProviders.forEach(provider => providerRegistry.add(provider));
+
+    return providerRegistry;
   }
 };
 
-export default [chatPlugin, autocompletionRegistryPlugin, aiProviderPlugin];
+export default [
+  providerRegistryPlugin,
+  autocompletionRegistryPlugin,
+  chatPlugin,
+  completerPlugin
+];

package/src/llm-models/index.ts

@@ -1,3 +1,50 @@
+import { ChatAnthropic } from '@langchain/anthropic';
+import { ChromeAI } from '@langchain/community/experimental/llms/chrome_ai';
+import { ChatMistralAI } from '@langchain/mistralai';
+import { ChatOpenAI } from '@langchain/openai';
+
+import { AnthropicCompleter } from './anthropic-completer';
+import { CodestralCompleter } from './codestral-completer';
+import { ChromeCompleter } from './chrome-completer';
+import { OpenAICompleter } from './openai-completer';
+
+import { instructions } from '../settings/instructions';
+import { ProviderSettings } from '../settings/schemas';
+
+import { IAIProvider } from '../tokens';
+
 export * from './base-completer';
-export * from './codestral-completer';
-export * from './utils';
+
+const AIProviders: IAIProvider[] = [
+  {
+    name: 'Anthropic',
+    chatModel: ChatAnthropic,
+    completer: AnthropicCompleter,
+    settingsSchema: ProviderSettings.Anthropic,
+    errorMessage: (error: any) => error.error.error.message
+  },
+  {
+    name: 'ChromeAI',
+    // TODO: fix
+    // @ts-expect-error: missing properties
+    chatModel: ChromeAI,
+    completer: ChromeCompleter,
+    instructions: instructions.ChromeAI,
+    settingsSchema: ProviderSettings.ChromeAI
+  },
+  {
+    name: 'MistralAI',
+    chatModel: ChatMistralAI,
+    completer: CodestralCompleter,
+    instructions: instructions.MistralAI,
+    settingsSchema: ProviderSettings.MistralAI
+  },
+  {
+    name: 'OpenAI',
+    chatModel: ChatOpenAI,
+    completer: OpenAICompleter,
+    settingsSchema: ProviderSettings.OpenAI
+  }
+];
+
+export { AIProviders };

package/src/provider.ts

@@ -4,11 +4,13 @@ import { BaseChatModel } from '@langchain/core/language_models/chat_models';
 import { ISignal, Signal } from '@lumino/signaling';
 import { ReadonlyPartialJSONObject } from '@lumino/coreutils';
 
-import { CompletionProvider } from './completion-provider';
-import { getChatModel, IBaseCompleter } from './llm-models';
-import { IAIProvider } from './token';
+import { IBaseCompleter } from './llm-models';
+import { IAIProvider, IAIProviderRegistry } from './tokens';
+import { JSONSchema7 } from 'json-schema';
 
-export const chatSystemPrompt = (options: AIProvider.IPromptOptions) => `
+export const chatSystemPrompt = (
+  options: AIProviderRegistry.IPromptOptions
+) => `
 You are Jupyternaut, a conversational assistant living in JupyterLab to help users.
 You are not a language model, but rather an application built on a foundation model from ${options.provider_name}.
 You are talkative and you provide lots of specific details from the foundation model's context.
@@ -36,40 +38,79 @@ would write.
 Do not include the prompt in the output, only the string that should be appended to the current input.
 `;
 
-export class AIProvider implements IAIProvider {
-  constructor(options: AIProvider.IOptions) {
-    this._completionProvider = new CompletionProvider({
-      name: 'None',
-      settings: {},
-      requestCompletion: options.requestCompletion
-    });
-    options.completionProviderManager.registerInlineProvider(
-      this._completionProvider
-    );
+export class AIProviderRegistry implements IAIProviderRegistry {
+  /**
+   * Get the list of provider names.
+   */
+  get providers(): string[] {
+    return Array.from(this._providers.keys());
+  }
+
+  /**
+   * Add a new provider.
+   */
+  add(provider: IAIProvider): void {
+    if (this._providers.has(provider.name)) {
+      throw new Error(
+        `A AI provider named '${provider.name}' is already registered`
+      );
+    }
+    this._providers.set(provider.name, provider);
   }
 
-  get name(): string {
+  /**
+   * Get the current provider name.
+   */
+  get currentName(): string {
     return this._name;
   }
 
   /**
    * Get the current completer of the completion provider.
    */
-  get completer(): IBaseCompleter | null {
-    if (this._name === null) {
+  get currentCompleter(): IBaseCompleter | null {
+    if (this._name === 'None') {
       return null;
     }
-    return this._completionProvider.completer;
+    return this._completer;
   }
 
   /**
    * Get the current llm chat model.
    */
-  get chatModel(): BaseChatModel | null {
-    if (this._name === null) {
+  get currentChatModel(): BaseChatModel | null {
+    if (this._name === 'None') {
       return null;
     }
-    return this._llmChatModel;
+    return this._chatModel;
+  }
+
+  /**
+   * Get the settings schema of a given provider.
+   */
+  getSettingsSchema(provider: string): JSONSchema7 {
+    return (this._providers.get(provider)?.settingsSchema?.properties ||
+      {}) as JSONSchema7;
+  }
+
+  /**
+   * Get the instructions of a given provider.
+   */
+  getInstructions(provider: string): string | undefined {
+    return this._providers.get(provider)?.instructions;
+  }
+
+  /**
+   * Format an error message from the current provider.
+   */
+  formatErrorMessage(error: any): string {
+    if (this._currentProvider?.errorMessage) {
+      return this._currentProvider?.errorMessage(error);
+    }
+    if (error.message) {
+      return error.message;
+    }
+    return error;
   }
 
   /**
@@ -87,43 +128,59 @@ export class AIProvider implements IAIProvider {
   }
 
   /**
-   * Set the models (chat model and completer).
-   * Creates the models if the name has changed, otherwise only updates their config.
+   * Set the providers (chat model and completer).
+   * Creates the providers if the name has changed, otherwise only updates their config.
    *
-   * @param name - the name of the model to use.
+   * @param name - the name of the provider to use.
    * @param settings - the settings for the models.
    */
-  setModels(name: string, settings: ReadonlyPartialJSONObject) {
-    try {
-      this._completionProvider.setCompleter(name, settings);
-      this._completerError = '';
-    } catch (e: any) {
-      this._completerError = e.message;
+  setProvider(name: string, settings: ReadonlyPartialJSONObject): void {
+    this._currentProvider = this._providers.get(name) ?? null;
+
+    if (this._currentProvider?.completer !== undefined) {
+      try {
+        this._completer = new this._currentProvider.completer({ ...settings });
+        this._completerError = '';
+      } catch (e: any) {
+        this._completerError = e.message;
+      }
+    } else {
+      this._completer = null;
     }
-    try {
-      this._llmChatModel = getChatModel(name, settings);
-      this._chatError = '';
-    } catch (e: any) {
-      this._chatError = e.message;
-      this._llmChatModel = null;
+
+    if (this._currentProvider?.chatModel !== undefined) {
+      try {
+        this._chatModel = new this._currentProvider.chatModel({ ...settings });
+        this._chatError = '';
+      } catch (e: any) {
+        this._chatError = e.message;
+        this._chatModel = null;
+      }
+    } else {
+      this._chatModel = null;
     }
     this._name = name;
-    this._modelChange.emit();
+    this._providerChanged.emit();
   }
 
-  get modelChange(): ISignal<IAIProvider, void> {
-    return this._modelChange;
+  /**
+   * A signal emitting when the provider or its settings has changed.
+   */
+  get providerChanged(): ISignal<IAIProviderRegistry, void> {
+    return this._providerChanged;
   }
 
-  private _completionProvider: CompletionProvider;
-  private _llmChatModel: BaseChatModel | null = null;
+  private _currentProvider: IAIProvider | null = null;
+  private _completer: IBaseCompleter | null = null;
+  private _chatModel: BaseChatModel | null = null;
   private _name: string = 'None';
-  private _modelChange = new Signal<IAIProvider, void>(this);
+  private _providerChanged = new Signal<IAIProviderRegistry, void>(this);
   private _chatError: string = '';
   private _completerError: string = '';
+  private _providers = new Map<string, IAIProvider>();
 }
 
-export namespace AIProvider {
+export namespace AIProviderRegistry {
   /**
    * The options for the LLM provider.
    */

package/src/settings/instructions.ts

@@ -0,0 +1,48 @@
+import { IDict } from '../tokens';
+
+const chromeAiInstructions = `
+<i class="fas fa-exclamation-triangle"></i> Support for ChromeAI is still experimental and only available in Google Chrome.
+
+You can test ChromeAI is enabled in your browser by going to the following URL: https://chromeai.org/
+
+Enable the proper flags in Google Chrome.
+
+- 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
+
+<img src="https://github.com/user-attachments/assets/d48f46cc-52ee-4ce5-9eaf-c763cdbee04c" alt="A screenshot showing how to enable the ChromeAI flag in Google Chrome" width="500px">
+
+Then restart Chrome for these changes to take effect.
+
+<i class="fas fa-exclamation-triangle"></i> 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.
+
+<i class="fa fa-info-circle" aria-hidden="true"></i> For more information about Chrome Built-in AI: https://developer.chrome.com/docs/ai/get-started
+`;
+
+const mistralAIInstructions = `
+<i class="fas fa-exclamation-triangle"></i> This extension is still very much experimental. It is not an official MistralAI extension.
+
+1. Go to https://console.mistral.ai/api-keys/ and create an API key.
+
+<img src="https://raw.githubusercontent.com/jupyterlite/ai/refs/heads/main/img/1-api-key.png" alt="Screenshot showing how to create an API key" width="500px">
+
+2. Open the JupyterLab settings and go to the **Ai providers** section to select the \`MistralAI\`
+   provider and the API key (required).
+
+<img src="https://raw.githubusercontent.com/jupyterlite/ai/refs/heads/main/img/2-jupyterlab-settings.png" alt="Screenshot showing how to add the API key to the settings" width="500px">
+
+3. Open the chat, or use the inline completer
+
+<img src="https://raw.githubusercontent.com/jupyterlite/ai/refs/heads/main/img/3-usage.png" alt="Screenshot showing how to use the chat" width="500px">
+`;
+
+export const instructions: IDict = {
+  ChromeAI: chromeAiInstructions,
+  MistralAI: mistralAIInstructions
+};