@jupyterlite/ai 0.4.0 → 0.6.0

package/src/components/stop-button.tsx

@@ -0,0 +1,56 @@
+/*
+ * Copyright (c) Jupyter Development Team.
+ * Distributed under the terms of the Modified BSD License.
+ */
+
+import StopIcon from '@mui/icons-material/Stop';
+import React from 'react';
+
+import { InputToolbarRegistry, TooltippedButton } from '@jupyter/chat';
+
+/**
+ * Properties of the stop button.
+ */
+export interface IStopButtonProps
+  extends InputToolbarRegistry.IToolbarItemProps {
+  /**
+   * The function to stop streaming.
+   */
+  stopStreaming: () => void;
+}
+
+/**
+ * The stop button.
+ */
+export function StopButton(props: IStopButtonProps): JSX.Element {
+  const tooltip = 'Stop streaming';
+  return (
+    <TooltippedButton
+      onClick={props.stopStreaming}
+      tooltip={tooltip}
+      buttonProps={{
+        size: 'small',
+        variant: 'contained',
+        title: tooltip
+      }}
+    >
+      <StopIcon />
+    </TooltippedButton>
+  );
+}
+
+/**
+ * factory returning the toolbar item.
+ */
+export function stopItem(
+  stopStreaming: () => void
+): InputToolbarRegistry.IToolbarItem {
+  return {
+    element: (props: InputToolbarRegistry.IToolbarItemProps) => {
+      const stopProps: IStopButtonProps = { ...props, stopStreaming };
+      return StopButton(stopProps);
+    },
+    position: 50,
+    hidden: true /* hidden by default */
+  };
+}

package/src/{llm-models/anthropic-completer.ts → default-providers/Anthropic/completer.ts}

@@ -6,8 +6,8 @@ import { ChatAnthropic } from '@langchain/anthropic';
 import { BaseChatModel } from '@langchain/core/language_models/chat_models';
 import { AIMessage, SystemMessage } from '@langchain/core/messages';
 
-import { BaseCompleter, IBaseCompleter } from './base-completer';
-import { COMPLETION_SYSTEM_PROMPT } from '../provider';
+import { BaseCompleter, IBaseCompleter } from '../../base-completer';
+import { COMPLETION_SYSTEM_PROMPT } from '../../provider';
 
 export class AnthropicCompleter implements IBaseCompleter {
   constructor(options: BaseCompleter.IOptions) {

package/src/{llm-models/chrome-completer.ts → default-providers/ChromeAI/completer.ts}

@@ -5,8 +5,9 @@ import {
 import { ChromeAI } from '@langchain/community/experimental/llms/chrome_ai';
 import { LLM } from '@langchain/core/language_models/llms';
 import { HumanMessage, SystemMessage } from '@langchain/core/messages';
-import { BaseCompleter, IBaseCompleter } from './base-completer';
-import { COMPLETION_SYSTEM_PROMPT } from '../provider';
+
+import { BaseCompleter, IBaseCompleter } from '../../base-completer';
+import { COMPLETION_SYSTEM_PROMPT } from '../../provider';
 
 /**
  * Regular expression to match the '```' string at the start of a string.

package/src/default-providers/ChromeAI/instructions.ts

@@ -0,0 +1,24 @@
+export default `
+<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>
+`;

package/src/{llm-models/codestral-completer.ts → default-providers/MistralAI/completer.ts}

@@ -11,8 +11,8 @@ import {
 import { ChatMistralAI } from '@langchain/mistralai';
 import { Throttler } from '@lumino/polling';
 
-import { BaseCompleter, IBaseCompleter } from './base-completer';
-import { COMPLETION_SYSTEM_PROMPT } from '../provider';
+import { BaseCompleter, IBaseCompleter } from '../../base-completer';
+import { COMPLETION_SYSTEM_PROMPT } from '../../provider';
 
 /**
  * The Mistral API has a rate limit of 1 request per second

package/src/default-providers/MistralAI/instructions.ts

@@ -0,0 +1,16 @@
+export default `
+<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">
+`;

package/src/{llm-models/openai-completer.ts → default-providers/OpenAI/completer.ts}

@@ -6,8 +6,8 @@ import { BaseChatModel } from '@langchain/core/language_models/chat_models';
 import { AIMessage, SystemMessage } from '@langchain/core/messages';
 import { ChatOpenAI } from '@langchain/openai';
 
-import { BaseCompleter, IBaseCompleter } from './base-completer';
-import { COMPLETION_SYSTEM_PROMPT } from '../provider';
+import { BaseCompleter, IBaseCompleter } from '../../base-completer';
+import { COMPLETION_SYSTEM_PROMPT } from '../../provider';
 
 export class OpenAICompleter implements IBaseCompleter {
   constructor(options: BaseCompleter.IOptions) {

package/src/default-providers/index.ts

@@ -0,0 +1,71 @@
+import {
+  JupyterFrontEnd,
+  JupyterFrontEndPlugin
+} from '@jupyterlab/application';
+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 { IAIProvider, IAIProviderRegistry } from '../tokens';
+
+// Import completers
+import { AnthropicCompleter } from './Anthropic/completer';
+import { ChromeCompleter } from './ChromeAI/completer';
+import { CodestralCompleter } from './MistralAI/completer';
+import { OpenAICompleter } from './OpenAI/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 OpenAISettings from './OpenAI/settings-schema.json';
+
+// Import instructions
+import ChromeAIInstructions from './ChromeAI/instructions';
+import MistralAIInstructions from './MistralAI/instructions';
+
+// Build the AIProvider list
+const AIProviders: IAIProvider[] = [
+  {
+    name: 'Anthropic',
+    chatModel: ChatAnthropic,
+    completer: AnthropicCompleter,
+    settingsSchema: AnthropicSettings,
+    errorMessage: (error: any) => error.error.error.message
+  },
+  {
+    name: 'ChromeAI',
+    // TODO: fix
+    // @ts-expect-error: missing properties
+    chatModel: ChromeAI,
+    completer: ChromeCompleter,
+    instructions: ChromeAIInstructions,
+    settingsSchema: ChromeAISettings
+  },
+  {
+    name: 'MistralAI',
+    chatModel: ChatMistralAI,
+    completer: CodestralCompleter,
+    instructions: MistralAIInstructions,
+    settingsSchema: MistralAISettings
+  },
+  {
+    name: 'OpenAI',
+    chatModel: ChatOpenAI,
+    completer: OpenAICompleter,
+    settingsSchema: OpenAISettings
+  }
+];
+
+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);
+      }
+    };
+  });

package/src/index.ts

@@ -1,11 +1,11 @@
 import {
   ActiveCellManager,
-  AutocompletionRegistry,
   buildChatSidebar,
   buildErrorWidget,
+  ChatCommandRegistry,
   IActiveCellManager,
-  IAutocompletionCommandsProps,
-  IAutocompletionRegistry
+  IChatCommandRegistry,
+  InputToolbarRegistry
 } from '@jupyter/chat';
 import {
   JupyterFrontEnd,
@@ -15,56 +15,45 @@ import { ReactWidget, IThemeManager } from '@jupyterlab/apputils';
 import { ICompletionProviderManager } from '@jupyterlab/completer';
 import { INotebookTracker } from '@jupyterlab/notebook';
 import { IRenderMimeRegistry } from '@jupyterlab/rendermime';
-import { ISettingRegistry } from '@jupyterlab/settingregistry';
+import {
+  ISettingConnector,
+  ISettingRegistry
+} from '@jupyterlab/settingregistry';
 import { IFormRendererRegistry } from '@jupyterlab/ui-components';
 import { ReadonlyPartialJSONObject } from '@lumino/coreutils';
+import { ISecretsManager } from 'jupyter-secrets-manager';
 
 import { ChatHandler } from './chat-handler';
 import { CompletionProvider } from './completion-provider';
-import { AIProviders } from './llm-models';
+import { defaultProviderPlugins } from './default-providers';
 import { AIProviderRegistry } from './provider';
-import { aiSettingsRenderer } from './settings/panel';
-import { renderSlashCommandOption } from './slash-commands';
+import { aiSettingsRenderer, SettingConnector } from './settings';
 import { IAIProviderRegistry } from './tokens';
+import { stopItem } from './components/stop-button';
 
-const autocompletionRegistryPlugin: JupyterFrontEndPlugin<IAutocompletionRegistry> =
-  {
-    id: '@jupyterlite/ai:autocompletion-registry',
-    description: 'Autocompletion registry',
-    autoStart: true,
-    provides: IAutocompletionRegistry,
-    activate: () => {
-      const autocompletionRegistry = new AutocompletionRegistry();
-      const options = ['/clear'];
-      const autocompletionCommands: IAutocompletionCommandsProps = {
-        opener: '/',
-        commands: options.map(option => {
-          return {
-            id: option.slice(1),
-            label: option,
-            description: 'Clear the chat window'
-          };
-        }),
-        props: {
-          renderOption: renderSlashCommandOption
-        }
-      };
-      autocompletionRegistry.add('jupyterlite-ai', autocompletionCommands);
-      return autocompletionRegistry;
-    }
-  };
+const chatCommandRegistryPlugin: JupyterFrontEndPlugin<IChatCommandRegistry> = {
+  id: '@jupyterlite/ai:autocompletion-registry',
+  description: 'Autocompletion registry',
+  autoStart: true,
+  provides: IChatCommandRegistry,
+  activate: () => {
+    const registry = new ChatCommandRegistry();
+    registry.addProvider(new ChatHandler.ClearCommandProvider());
+    return registry;
+  }
+};
 
 const chatPlugin: JupyterFrontEndPlugin<void> = {
   id: '@jupyterlite/ai:chat',
   description: 'LLM chat extension',
   autoStart: true,
-  requires: [IAIProviderRegistry, IRenderMimeRegistry, IAutocompletionRegistry],
+  requires: [IAIProviderRegistry, IRenderMimeRegistry, IChatCommandRegistry],
   optional: [INotebookTracker, ISettingRegistry, IThemeManager],
   activate: async (
     app: JupyterFrontEnd,
     providerRegistry: IAIProviderRegistry,
     rmRegistry: IRenderMimeRegistry,
-    autocompletionRegistry: IAutocompletionRegistry,
+    chatCommandRegistry: IChatCommandRegistry,
     notebookTracker: INotebookTracker | null,
     settingsRegistry: ISettingRegistry | null,
     themeManager: IThemeManager | null
@@ -115,12 +104,30 @@ const chatPlugin: JupyterFrontEndPlugin<void> = {
       });
 
     let chatWidget: ReactWidget | null = null;
+
+    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
+      ) {
+        inputToolbarRegistry.hide('send');
+        inputToolbarRegistry.show('stop');
+      } else {
+        inputToolbarRegistry.hide('stop');
+        inputToolbarRegistry.show('send');
+      }
+    });
+
     try {
       chatWidget = buildChatSidebar({
         model: chatHandler,
         themeManager,
         rmRegistry,
-        autocompletionRegistry
+        chatCommandRegistry,
+        inputToolbarRegistry
       });
       chatWidget.title.caption = 'Jupyterlite AI Chat';
     } catch (e) {
@@ -154,20 +161,28 @@ const providerRegistryPlugin: JupyterFrontEndPlugin<IAIProviderRegistry> = {
   id: '@jupyterlite/ai:provider-registry',
   autoStart: true,
   requires: [IFormRendererRegistry, ISettingRegistry],
-  optional: [IRenderMimeRegistry],
+  optional: [IRenderMimeRegistry, ISecretsManager, ISettingConnector],
   provides: IAIProviderRegistry,
   activate: (
     app: JupyterFrontEnd,
     editorRegistry: IFormRendererRegistry,
     settingRegistry: ISettingRegistry,
-    rmRegistry?: IRenderMimeRegistry
+    rmRegistry?: IRenderMimeRegistry,
+    secretsManager?: ISecretsManager,
+    settingConnector?: ISettingConnector
   ): IAIProviderRegistry => {
-    const providerRegistry = new AIProviderRegistry();
+    const providerRegistry = new AIProviderRegistry({ secretsManager });
 
     editorRegistry.addRenderer(
       '@jupyterlite/ai:provider-registry.AIprovider',
-      aiSettingsRenderer({ providerRegistry, rmRegistry })
+      aiSettingsRenderer({
+        providerRegistry,
+        rmRegistry,
+        secretsManager,
+        settingConnector
+      })
     );
+
     settingRegistry
       .load(providerRegistryPlugin.id)
       .then(settings => {
@@ -176,10 +191,10 @@ const providerRegistryPlugin: JupyterFrontEndPlugin<IAIProviderRegistry> = {
           const providerSettings = (settings.get('AIprovider').composite ?? {
             provider: 'None'
           }) as ReadonlyPartialJSONObject;
-          providerRegistry.setProvider(
-            providerSettings.provider as string,
-            providerSettings
-          );
+          providerRegistry.setProvider({
+            name: providerSettings.provider as string,
+            settings: providerSettings
+          });
         };
 
         settings.changed.connect(() => updateProvider());
@@ -192,16 +207,29 @@ const providerRegistryPlugin: JupyterFrontEndPlugin<IAIProviderRegistry> = {
         );
       });
 
-    // Initialize the registry with the default providers
-    AIProviders.forEach(provider => providerRegistry.add(provider));
-
     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: JupyterFrontEndPlugin<ISettingConnector> = {
+  id: '@jupyterlite/ai:settings-connector',
+  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,
-  autocompletionRegistryPlugin,
+  chatCommandRegistryPlugin,
   chatPlugin,
-  completerPlugin
+  completerPlugin,
+  settingsConnector,
+  ...defaultProviderPlugins
 ];

package/src/provider.ts

@@ -1,12 +1,22 @@
-import { ICompletionProviderManager } 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';
 import { ReadonlyPartialJSONObject } from '@lumino/coreutils';
-
-import { IBaseCompleter } from './llm-models';
-import { IAIProvider, IAIProviderRegistry } from './tokens';
 import { JSONSchema7 } from 'json-schema';
+import { ISecretsManager } from 'jupyter-secrets-manager';
+
+import { IBaseCompleter } from './base-completer';
+import {
+  getSecretId,
+  SECRETS_NAMESPACE,
+  SECRETS_REPLACEMENT
+} from './settings';
+import {
+  IAIProvider,
+  IAIProviderRegistry,
+  IDict,
+  ISetProviderOptions
+} from './tokens';
 
 export const chatSystemPrompt = (
   options: AIProviderRegistry.IPromptOptions
@@ -39,6 +49,13 @@ Do not include the prompt in the output, only the string that should be appended
 `;
 
 export class AIProviderRegistry implements IAIProviderRegistry {
+  /**
+   * The constructor of the provider registry.
+   */
+  constructor(options: AIProviderRegistry.IOptions) {
+    this._secretsManager = options.secretsManager || null;
+  }
+
   /**
    * Get the list of provider names.
    */
@@ -56,6 +73,11 @@ export class AIProviderRegistry implements IAIProviderRegistry {
       );
     }
     this._providers.set(provider.name, provider);
+
+    // Set the provider if the loading has been deferred.
+    if (provider.name === this._deferredProvider?.name) {
+      this.setProvider(this._deferredProvider);
+    }
   }
 
   /**
@@ -131,15 +153,36 @@ export class AIProviderRegistry implements IAIProviderRegistry {
    * 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 provider to use.
-   * @param settings - the settings for the models.
+   * @param options - An object with the name and the settings of the provider to use.
    */
-  setProvider(name: string, settings: ReadonlyPartialJSONObject): void {
+  async setProvider(options: ISetProviderOptions): Promise<void> {
+    const { name, settings } = options;
     this._currentProvider = this._providers.get(name) ?? null;
+    if (this._currentProvider === null) {
+      // The current provider may not be loaded when the settings are first loaded.
+      // Let's defer the provider loading.
+      this._deferredProvider = options;
+    } else {
+      this._deferredProvider = null;
+    }
+
+    // Build a new settings object containing the secrets.
+    const fullSettings: IDict = {};
+    for (const key of Object.keys(settings)) {
+      if (settings[key] === SECRETS_REPLACEMENT) {
+        const id = getSecretId(name, key);
+        const secrets = await this._secretsManager?.get(SECRETS_NAMESPACE, id);
+        fullSettings[key] = secrets?.value || settings[key];
+        continue;
+      }
+      fullSettings[key] = settings[key];
+    }
 
     if (this._currentProvider?.completer !== undefined) {
       try {
-        this._completer = new this._currentProvider.completer({ ...settings });
+        this._completer = new this._currentProvider.completer({
+          ...fullSettings
+        });
         this._completerError = '';
       } catch (e: any) {
         this._completerError = e.message;
@@ -150,7 +193,9 @@ export class AIProviderRegistry implements IAIProviderRegistry {
 
     if (this._currentProvider?.chatModel !== undefined) {
       try {
-        this._chatModel = new this._currentProvider.chatModel({ ...settings });
+        this._chatModel = new this._currentProvider.chatModel({
+          ...fullSettings
+        });
         this._chatError = '';
       } catch (e: any) {
         this._chatError = e.message;
@@ -170,6 +215,7 @@ export class AIProviderRegistry implements IAIProviderRegistry {
     return this._providerChanged;
   }
 
+  private _secretsManager: ISecretsManager | null;
   private _currentProvider: IAIProvider | null = null;
   private _completer: IBaseCompleter | null = null;
   private _chatModel: BaseChatModel | null = null;
@@ -178,6 +224,7 @@ export class AIProviderRegistry implements IAIProviderRegistry {
   private _chatError: string = '';
   private _completerError: string = '';
   private _providers = new Map<string, IAIProvider>();
+  private _deferredProvider: ISetProviderOptions | null = null;
 }
 
 export namespace AIProviderRegistry {
@@ -186,13 +233,9 @@ export namespace AIProviderRegistry {
    */
   export interface IOptions {
     /**
-     * The completion provider manager in which register the LLM completer.
-     */
-    completionProviderManager: ICompletionProviderManager;
-    /**
-     * The application commands registry.
+     * The secrets manager used in the application.
      */
-    requestCompletion: () => void;
+    secretsManager?: ISecretsManager;
   }
 
   /**

package/src/settings/index.ts

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