@jupyterlite/ai 0.5.0 → 0.6.1

package/src/index.ts

@@ -4,7 +4,8 @@ import {
   buildErrorWidget,
   ChatCommandRegistry,
   IActiveCellManager,
-  IChatCommandRegistry
+  IChatCommandRegistry,
+  InputToolbarRegistry
 } from '@jupyter/chat';
 import {
   JupyterFrontEnd,
@@ -14,20 +15,24 @@ 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 { ISecretsManager, SecretsManager } from 'jupyter-secrets-manager';
 
 import { ChatHandler } from './chat-handler';
 import { CompletionProvider } from './completion-provider';
 import { defaultProviderPlugins } from './default-providers';
 import { AIProviderRegistry } from './provider';
-import { aiSettingsRenderer } from './settings/panel';
-import { IAIProviderRegistry } from './tokens';
+import { aiSettingsRenderer, SettingConnector } from './settings';
+import { IAIProviderRegistry, PLUGIN_IDS } from './tokens';
+import { stopItem } from './components/stop-button';
 
 const chatCommandRegistryPlugin: JupyterFrontEndPlugin<IChatCommandRegistry> = {
-  id: '@jupyterlite/ai:autocompletion-registry',
+  id: PLUGIN_IDS.chatCommandRegistry,
   description: 'Autocompletion registry',
   autoStart: true,
   provides: IChatCommandRegistry,
@@ -39,7 +44,7 @@ const chatCommandRegistryPlugin: JupyterFrontEndPlugin<IChatCommandRegistry> = {
 };
 
 const chatPlugin: JupyterFrontEndPlugin<void> = {
-  id: '@jupyterlite/ai:chat',
+  id: PLUGIN_IDS.chat,
   description: 'LLM chat extension',
   autoStart: true,
   requires: [IAIProviderRegistry, IRenderMimeRegistry, IChatCommandRegistry],
@@ -99,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,
-        chatCommandRegistry
+        chatCommandRegistry,
+        inputToolbarRegistry
       });
       chatWidget.title.caption = 'Jupyterlite AI Chat';
     } catch (e) {
@@ -118,7 +141,7 @@ const chatPlugin: JupyterFrontEndPlugin<void> = {
 };
 
 const completerPlugin: JupyterFrontEndPlugin<void> = {
-  id: '@jupyterlite/ai:completer',
+  id: PLUGIN_IDS.completer,
   autoStart: true,
   requires: [IAIProviderRegistry, ICompletionProviderManager],
   activate: (
@@ -134,51 +157,77 @@ const completerPlugin: JupyterFrontEndPlugin<void> = {
   }
 };
 
-const providerRegistryPlugin: JupyterFrontEndPlugin<IAIProviderRegistry> = {
-  id: '@jupyterlite/ai:provider-registry',
-  autoStart: true,
-  requires: [IFormRendererRegistry, ISettingRegistry],
-  optional: [IRenderMimeRegistry, ISecretsManager],
-  provides: IAIProviderRegistry,
-  activate: (
-    app: JupyterFrontEnd,
-    editorRegistry: IFormRendererRegistry,
-    settingRegistry: ISettingRegistry,
-    rmRegistry?: IRenderMimeRegistry,
-    secretsManager?: ISecretsManager
-  ): IAIProviderRegistry => {
-    const providerRegistry = new AIProviderRegistry();
-
-    editorRegistry.addRenderer(
-      '@jupyterlite/ai:provider-registry.AIprovider',
-      aiSettingsRenderer({ providerRegistry, rmRegistry, secretsManager })
-    );
-    settingRegistry
-      .load(providerRegistryPlugin.id)
-      .then(settings => {
-        const updateProvider = () => {
-          // Update the settings to the AI providers.
-          const providerSettings = (settings.get('AIprovider').composite ?? {
-            provider: 'None'
-          }) as ReadonlyPartialJSONObject;
-          providerRegistry.setProvider(
-            providerSettings.provider as string,
-            providerSettings
-          );
-        };
-
-        settings.changed.connect(() => updateProvider());
-        updateProvider();
-      })
-      .catch(reason => {
-        console.error(
-          `Failed to load settings for ${providerRegistryPlugin.id}`,
-          reason
-        );
+const providerRegistryPlugin: JupyterFrontEndPlugin<IAIProviderRegistry> =
+  SecretsManager.sign(PLUGIN_IDS.providerRegistry, token => ({
+    id: PLUGIN_IDS.providerRegistry,
+    autoStart: true,
+    requires: [IFormRendererRegistry, ISettingRegistry],
+    optional: [IRenderMimeRegistry, ISecretsManager, ISettingConnector],
+    provides: IAIProviderRegistry,
+    activate: (
+      app: JupyterFrontEnd,
+      editorRegistry: IFormRendererRegistry,
+      settingRegistry: ISettingRegistry,
+      rmRegistry?: IRenderMimeRegistry,
+      secretsManager?: ISecretsManager,
+      settingConnector?: ISettingConnector
+    ): IAIProviderRegistry => {
+      const providerRegistry = new AIProviderRegistry({
+        token,
+        secretsManager
       });
 
-    return providerRegistry;
-  }
+      editorRegistry.addRenderer(
+        `${PLUGIN_IDS.providerRegistry}.AIprovider`,
+        aiSettingsRenderer({
+          providerRegistry,
+          secretsToken: token,
+          rmRegistry,
+          secretsManager,
+          settingConnector
+        })
+      );
+
+      settingRegistry
+        .load(providerRegistryPlugin.id)
+        .then(settings => {
+          const updateProvider = () => {
+            // Update the settings to the AI providers.
+            const providerSettings = (settings.get('AIprovider').composite ?? {
+              provider: 'None'
+            }) as ReadonlyPartialJSONObject;
+            providerRegistry.setProvider({
+              name: providerSettings.provider as string,
+              settings: providerSettings
+            });
+          };
+
+          settings.changed.connect(() => updateProvider());
+          updateProvider();
+        })
+        .catch(reason => {
+          console.error(
+            `Failed to load settings for ${providerRegistryPlugin.id}`,
+            reason
+          );
+        });
+
+      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: 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 [
@@ -186,5 +235,6 @@ export default [
   chatCommandRegistryPlugin,
   chatPlugin,
   completerPlugin,
+  settingsConnector,
   ...defaultProviderPlugins
 ];

package/src/provider.ts

@@ -1,12 +1,21 @@
-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 { JSONSchema7 } from 'json-schema';
+import { ISecretsManager } from 'jupyter-secrets-manager';
 
 import { IBaseCompleter } from './base-completer';
-import { IAIProvider, IAIProviderRegistry } from './tokens';
-import { JSONSchema7 } from 'json-schema';
+import { getSecretId, SECRETS_REPLACEMENT } from './settings';
+import {
+  IAIProvider,
+  IAIProviderRegistry,
+  IDict,
+  ISetProviderOptions,
+  PLUGIN_IDS
+} from './tokens';
+
+const SECRETS_NAMESPACE = PLUGIN_IDS.providerRegistry;
 
 export const chatSystemPrompt = (
   options: AIProviderRegistry.IPromptOptions
@@ -39,6 +48,14 @@ 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;
+    Private.setToken(options.token);
+  }
+
   /**
    * 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,40 @@ 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(
+          Private.getToken(),
+          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 +197,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 +219,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 +228,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 +237,13 @@ export namespace AIProviderRegistry {
    */
   export interface IOptions {
     /**
-     * The completion provider manager in which register the LLM completer.
+     * The secrets manager used in the application.
      */
-    completionProviderManager: ICompletionProviderManager;
+    secretsManager?: ISecretsManager;
     /**
-     * The application commands registry.
+     * The token used to request the secrets manager.
      */
-    requestCompletion: () => void;
+    token: symbol;
   }
 
   /**
@@ -247,3 +298,24 @@ export namespace AIProviderRegistry {
     });
   }
 }
+
+namespace Private {
+  /**
+   * The token to use with the secrets manager.
+   */
+  let secretsToken: symbol;
+
+  /**
+   * Set of the token.
+   */
+  export function setToken(value: symbol): void {
+    secretsToken = value;
+  }
+
+  /**
+   * get the token.
+   */
+  export function getToken(): symbol {
+    return secretsToken;
+  }
+}

package/src/settings/index.ts

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

package/src/settings/panel.tsx

@@ -1,5 +1,8 @@
 import { IRenderMimeRegistry } from '@jupyterlab/rendermime';
-import { ISettingRegistry } from '@jupyterlab/settingregistry';
+import {
+  ISettingConnector,
+  ISettingRegistry
+} from '@jupyterlab/settingregistry';
 import { FormComponent, IFormRenderer } from '@jupyterlab/ui-components';
 import { ArrayExt } from '@lumino/algorithm';
 import { JSONExt } from '@lumino/coreutils';
@@ -10,19 +13,27 @@ import { JSONSchema7 } from 'json-schema';
 import { ISecretsManager } from 'jupyter-secrets-manager';
 import React from 'react';
 
+import { getSecretId, SettingConnector } from '.';
 import baseSettings from './base.json';
-import { IAIProviderRegistry, IDict } from '../tokens';
+import { IAIProviderRegistry, IDict, PLUGIN_IDS } from '../tokens';
 
-const SECRETS_NAMESPACE = '@jupyterlite/ai';
 const MD_MIME_TYPE = 'text/markdown';
 const STORAGE_NAME = '@jupyterlite/ai:settings';
 const INSTRUCTION_CLASS = 'jp-AISettingsInstructions';
+const SECRETS_NAMESPACE = PLUGIN_IDS.providerRegistry;
 
 export const aiSettingsRenderer = (options: {
   providerRegistry: IAIProviderRegistry;
+  secretsToken?: symbol;
   rmRegistry?: IRenderMimeRegistry;
   secretsManager?: ISecretsManager;
+  settingConnector?: ISettingConnector;
 }): IFormRenderer => {
+  const { secretsToken } = options;
+  delete options.secretsToken;
+  if (secretsToken) {
+    Private.setToken(secretsToken);
+  }
   return {
     fieldRenderer: (props: FieldProps) => {
       props.formContext = { ...props.formContext, ...options };
@@ -54,10 +65,13 @@ export class AiSettings extends React.Component<
     this._providerRegistry = props.formContext.providerRegistry;
     this._rmRegistry = props.formContext.rmRegistry ?? null;
     this._secretsManager = props.formContext.secretsManager ?? null;
+    this._settingConnector = props.formContext.settingConnector ?? null;
     this._settings = props.formContext.settings;
 
     this._useSecretsManager =
       (this._settings.get('UseSecretsManager').composite as boolean) ?? true;
+    this._hideSecretFields =
+      (this._settings.get('HideSecretFields').composite as boolean) ?? true;
 
     // Initialize the providers schema.
     const providerSchema = JSONExt.deepCopy(baseSettings) as any;
@@ -110,6 +124,12 @@ export class AiSettings extends React.Component<
       if (useSecretsManager !== this._useSecretsManager) {
         this.updateUseSecretsManager(useSecretsManager);
       }
+      const hideSecretFields =
+        (this._settings.get('HideSecretFields').composite as boolean) ?? true;
+      if (hideSecretFields !== this._hideSecretFields) {
+        this._hideSecretFields = hideSecretFields;
+        this._updateSchema();
+      }
     });
   }
 
@@ -123,15 +143,16 @@ export class AiSettings extends React.Component<
       return;
     }
 
-    await this._secretsManager.detachAll(SECRETS_NAMESPACE);
+    await this._secretsManager.detachAll(Private.getToken(), SECRETS_NAMESPACE);
     this._formInputs = [...inputs];
     this._unsavedFields = [];
     for (let i = 0; i < inputs.length; i++) {
       if (inputs[i].type.toLowerCase() === 'password') {
         const label = inputs[i].getAttribute('label');
         if (label) {
-          const id = `${this._provider}-${label}`;
+          const id = getSecretId(this._provider, label);
           this._secretsManager.attach(
+            Private.getToken(),
             SECRETS_NAMESPACE,
             id,
             inputs[i],
@@ -141,6 +162,16 @@ export class AiSettings extends React.Component<
         }
       }
     }
+    if (this._settingConnector instanceof SettingConnector) {
+      this._settingConnector.doNotSave = this._unsavedFields;
+    }
+  }
+
+  componentWillUnmount(): void {
+    if (!this._secretsManager || !this._useSecretsManager) {
+      return;
+    }
+    this._secretsManager.detachAll(Private.getToken(), SECRETS_NAMESPACE);
   }
 
   /**
@@ -184,9 +215,12 @@ export class AiSettings extends React.Component<
     if (!value) {
       // Detach all the password inputs attached to the secrets manager, and save the
       // current settings to the local storage to save the password.
-      this._secretsManager?.detachAll(SECRETS_NAMESPACE);
+      this._secretsManager?.detachAll(Private.getToken(), SECRETS_NAMESPACE);
       this._formInputs = [];
       this._unsavedFields = [];
+      if (this._settingConnector instanceof SettingConnector) {
+        this._settingConnector.doNotSave = [];
+      }
       this.saveSettings(this._currentSettings);
     } else {
       // Remove all the keys stored locally and attach the password inputs to the
@@ -202,18 +236,11 @@ export class AiSettings extends React.Component<
       localStorage.setItem(STORAGE_NAME, JSON.stringify(settings));
       this.componentDidUpdate();
     }
+    this._settings
+      .set('AIprovider', { provider: this._provider, ...this._currentSettings })
+      .catch(console.error);
   };
 
-  /**
-   * Update the UI schema of the form.
-   * Currently use to hide API keys.
-   */
-  private _updateUiSchema(key: string) {
-    if (key.toLowerCase().includes('key')) {
-      this._uiSchema[key] = { 'ui:widget': 'password' };
-    }
-  }
-
   /**
    * Build the schema for a given provider.
    */
@@ -226,8 +253,13 @@ export class AiSettings extends React.Component<
 
     if (settingsSchema) {
       Object.entries(settingsSchema).forEach(([key, value]) => {
+        if (key.toLowerCase().includes('key')) {
+          if (this._hideSecretFields) {
+            return;
+          }
+          this._uiSchema[key] = { 'ui:widget': 'password' };
+        }
         schema.properties[key] = value;
-        this._updateUiSchema(key);
       });
     }
     return schema as JSONSchema7;
@@ -335,8 +367,10 @@ export class AiSettings extends React.Component<
   private _provider: string;
   private _providerSchema: JSONSchema7;
   private _useSecretsManager: boolean;
+  private _hideSecretFields: boolean;
   private _rmRegistry: IRenderMimeRegistry | null;
   private _secretsManager: ISecretsManager | null;
+  private _settingConnector: ISettingConnector | null;
   private _currentSettings: IDict<any> = { provider: 'None' };
   private _uiSchema: IDict<any> = {};
   private _settings: ISettingRegistry.ISettings;
@@ -344,3 +378,24 @@ export class AiSettings extends React.Component<
   private _unsavedFields: string[] = [];
   private _formInputs: HTMLInputElement[] = [];
 }
+
+namespace Private {
+  /**
+   * The token to use with the secrets manager.
+   */
+  let secretsToken: symbol;
+
+  /**
+   * Set of the token.
+   */
+  export function setToken(value: symbol): void {
+    secretsToken = value;
+  }
+
+  /**
+   * get the token.
+   */
+  export function getToken(): symbol {
+    return secretsToken;
+  }
+}

package/src/settings/settings-connector.ts

@@ -0,0 +1,89 @@
+import { PageConfig } from '@jupyterlab/coreutils';
+import {
+  ISettingConnector,
+  ISettingRegistry
+} from '@jupyterlab/settingregistry';
+import { DataConnector, IDataConnector } from '@jupyterlab/statedb';
+import { Throttler } from '@lumino/polling';
+import * as json5 from 'json5';
+
+import { SECRETS_REPLACEMENT } from '.';
+
+/**
+ * A data connector for fetching settings.
+ *
+ * #### Notes
+ * This connector adds a query parameter to the base services setting manager.
+ */
+export class SettingConnector
+  extends DataConnector<ISettingRegistry.IPlugin, string>
+  implements ISettingConnector
+{
+  constructor(connector: IDataConnector<ISettingRegistry.IPlugin, string>) {
+    super();
+    this._connector = connector;
+  }
+
+  set doNotSave(fields: string[]) {
+    this._doNotSave = [...fields];
+  }
+
+  /**
+   * Fetch settings for a plugin.
+   * @param id - The plugin ID
+   *
+   * #### Notes
+   * The REST API requests are throttled at one request per plugin per 100ms.
+   */
+  fetch(id: string): Promise<ISettingRegistry.IPlugin | undefined> {
+    const throttlers = this._throttlers;
+    if (!(id in throttlers)) {
+      throttlers[id] = new Throttler(() => this._connector.fetch(id), 100);
+    }
+    return throttlers[id].invoke();
+  }
+
+  async list(query: 'ids'): Promise<{ ids: string[] }>;
+  async list(
+    query: 'active' | 'all'
+  ): Promise<{ ids: string[]; values: ISettingRegistry.IPlugin[] }>;
+  async list(
+    query: 'active' | 'all' | 'ids' = 'all'
+  ): Promise<{ ids: string[]; values?: ISettingRegistry.IPlugin[] }> {
+    const { isDisabled } = PageConfig.Extension;
+    const { ids, values } = await this._connector.list(
+      query === 'ids' ? 'ids' : undefined
+    );
+
+    if (query === 'all') {
+      return { ids, values };
+    }
+
+    if (query === 'ids') {
+      return { ids };
+    }
+
+    return {
+      ids: ids.filter(id => !isDisabled(id)),
+      values: values.filter(({ id }) => !isDisabled(id))
+    };
+  }
+
+  async save(id: string, raw: string): Promise<void> {
+    const settings = json5.parse(raw);
+    this._doNotSave.forEach(field => {
+      if (
+        settings['AIprovider'] !== undefined &&
+        settings['AIprovider'][field] !== undefined &&
+        settings['AIprovider'][field] !== ''
+      ) {
+        settings['AIprovider'][field] = SECRETS_REPLACEMENT;
+      }
+    });
+    await this._connector.save(id, json5.stringify(settings, null, 2));
+  }
+
+  private _connector: IDataConnector<ISettingRegistry.IPlugin, string>;
+  private _doNotSave: string[] = [];
+  private _throttlers: { [key: string]: Throttler } = Object.create(null);
+}

package/src/settings/utils.ts

@@ -0,0 +1,5 @@
+export const SECRETS_REPLACEMENT = '***';
+
+export function getSecretId(provider: string, label: string) {
+  return `${provider}-${label}`;
+}