@jupyterlite/ai 0.6.2 → 0.8.0

package/src/provider.ts

@@ -1,7 +1,13 @@
+import { Notification } from '@jupyterlab/apputils';
+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';
 import { ReadonlyPartialJSONObject } from '@lumino/coreutils';
+import { Debouncer } from '@lumino/polling';
+import { ISignal, Signal } from '@lumino/signaling';
 import { JSONSchema7 } from 'json-schema';
 import { ISecretsManager } from 'jupyter-secrets-manager';
 
@@ -11,41 +17,13 @@ import {
   IAIProvider,
   IAIProviderRegistry,
   IDict,
-  ISetProviderOptions,
+  ModelRole,
   PLUGIN_IDS
 } from './tokens';
+import { AIChatModel, AICompleter } from './types/ai-model';
 
 const SECRETS_NAMESPACE = PLUGIN_IDS.providerRegistry;
-
-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.
-You may use Markdown to format your response.
-If your response includes code, they must be enclosed in Markdown fenced code blocks (with triple backticks before and after).
-If your response includes mathematical notation, they must be expressed in LaTeX markup and enclosed in LaTeX delimiters.
-All dollar quantities (of USD) must be formatted in LaTeX, with the \`$\` symbol escaped by a single backslash \`\\\`.
-- Example prompt: \`If I have \\\\$100 and spend \\\\$20, how much money do I have left?\`
-- **Correct** response: \`You have \\(\\$80\\) remaining.\`
-- **Incorrect** response: \`You have $80 remaining.\`
-If you do not know the answer to a question, answer truthfully by responding that you do not know.
-The following is a friendly conversation between you and a human.
-`;
-
-export const COMPLETION_SYSTEM_PROMPT = `
-You are an application built to provide helpful code completion suggestions.
-You should only produce code. Keep comments to minimum, use the
-programming language comment syntax. Produce clean code.
-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.
-Do not include the prompt in the output, only the string that should be appended to the current input.
-`;
+const NOTIFICATION_DELAY = 2000;
 
 export class AIProviderRegistry implements IAIProviderRegistry {
   /**
@@ -54,64 +32,122 @@ export class AIProviderRegistry implements IAIProviderRegistry {
   constructor(options: AIProviderRegistry.IOptions) {
     this._secretsManager = options.secretsManager || null;
     Private.setToken(options.token);
+
+    this._notifications = {
+      chat: new Debouncer(this._emitErrorNotification, NOTIFICATION_DELAY),
+      completer: new Debouncer(this._emitErrorNotification, NOTIFICATION_DELAY)
+    };
   }
 
   /**
    * 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) {
-      this.setProvider(this._deferredProvider);
+    // Set the providers if the loading has been deferred.
+    if (provider.name === this._deferredProvider.completer?.name) {
+      this.setCompleterProvider(this._deferredProvider.completer);
+    }
+    if (provider.name === this._deferredProvider.chat?.name) {
+      this.setChatProvider(this._deferredProvider.chat);
     }
   }
 
   /**
    * Get the current provider name.
    */
-  get currentName(): string {
-    return this._name;
+  currentName(role: ModelRole): string {
+    return Private.getName(role);
   }
 
   /**
-   * 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('completer') === '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.
+   * Getter/setter for the completer system prompt.
    */
-  get currentChatModel(): BaseChatModel | null {
-    if (this._name === 'None') {
+  get completerSystemPrompt(): string {
+    return this._completerPrompt.replaceAll(
+      '$provider_name$',
+      this.currentName('completer')
+    );
+  }
+  set completerSystemPrompt(value: string) {
+    this._completerPrompt = value;
+  }
+
+  /**
+   * Get the current AIChatModel.
+   */
+  get currentChatModel(): AIChatModel | null {
+    if (Private.getName('chat') === 'None') {
+      return null;
+    }
+    const currentProvider =
+      Private.providers.get(Private.getName('chat')) ?? 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)
+    };
+  }
+
+  /**
+   * Getter/setter for the chat system prompt.
+   */
+  get chatSystemPrompt(): string {
+    return this._chatPrompt.replaceAll(
+      '$provider_name$',
+      this.currentName('chat')
+    );
+  }
+  set chatSystemPrompt(value: string) {
+    this._chatPrompt = value;
   }
 
   /**
    * 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 +155,26 @@ 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('chat')) ?? null;
+    if (currentProvider?.errorMessage) {
+      return currentProvider?.errorMessage(error);
     }
     if (error.message) {
       return error.message;
@@ -136,101 +183,210 @@ export class AIProviderRegistry implements IAIProviderRegistry {
   }
 
   /**
-   * Get the current chat error;
+   * Get/set the current chat error;
    */
   get chatError(): string {
     return this._chatError;
   }
+  private set chatError(error: string) {
+    this._chatError = error;
+    if (error !== '') {
+      this._notifications.chat.invoke(`Chat: ${error}`);
+    }
+  }
 
   /**
-   * get the current completer error.
+   * Get/set the current completer error.
    */
   get completerError(): string {
     return this._completerError;
   }
+  private set completerError(error: string) {
+    this._completerError = error;
+    if (error !== '') {
+      this._notifications.completer.invoke(`Completer: ${error}`);
+    }
+  }
+
+  /**
+   * A function to emit a notification error.
+   */
+  private _emitErrorNotification(error: string) {
+    Notification.emit(error, 'error', {
+      autoClose: NOTIFICATION_DELAY
+    });
+  }
 
   /**
-   * Set the providers (chat model and completer).
-   * Creates the providers if the name has changed, otherwise only updates their config.
+   * 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.
    */
-  async setProvider(options: ISetProviderOptions): Promise<void> {
-    const { name, settings } = options;
-    this._currentProvider = this._providers.get(name) ?? null;
-    if (this._currentProvider === null) {
+  async setCompleterProvider(
+    settings: ReadonlyPartialJSONObject
+  ): Promise<void> {
+    this.completerError = '';
+    if (!Object.keys(settings).includes('provider')) {
+      Private.setName('completer', 'None');
+      Private.setCompleter(null);
+      this.completerError =
+        'The provider is missing from the completer settings';
+      return;
+    }
+    const provider = settings['provider'] as string;
+    const currentProvider = Private.providers.get(provider) ?? 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;
+      this._deferredProvider.completer = settings;
+      Private.setName('completer', provider);
+      Private.setCompleter(null);
+      return;
     } else {
-      this._deferredProvider = null;
+      this._deferredProvider.completer = 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
-        );
-        if (secrets !== undefined) {
-          fullSettings[key] = secrets.value;
-        }
-        continue;
+    const compatibilityCheck = this.getCompatibilityCheck(provider);
+    if (compatibilityCheck !== undefined) {
+      const error = await compatibilityCheck();
+      if (error !== null) {
+        this.completerError = error.trim();
+        Private.setName('completer', 'None');
+        this._providerChanged.emit('completer');
+        return;
       }
-      fullSettings[key] = settings[key];
     }
 
-    if (this._currentProvider?.completer !== undefined) {
+    // Get the settings including the secrets.
+    const fullSettings = await this._buildFullSettings(provider, settings);
+
+    if (currentProvider?.completer !== undefined) {
       try {
-        this._completer = new this._currentProvider.completer({
-          settings: fullSettings
-        });
-        this._completerError = '';
+        Private.setCompleter(
+          new currentProvider.completer({
+            providerRegistry: this,
+            settings: fullSettings
+          })
+        );
       } catch (e: any) {
-        this._completerError = e.message;
+        this.completerError = e.message;
       }
     } else {
-      this._completer = null;
+      Private.setCompleter(null);
+    }
+    Private.setName('completer', provider);
+    this._providerChanged.emit('completer');
+  }
+
+  /**
+   * 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.
+   */
+  async setChatProvider(settings: ReadonlyPartialJSONObject): Promise<void> {
+    this.chatError = '';
+    if (!Object.keys(settings).includes('provider')) {
+      Private.setName('chat', 'None');
+      Private.setChatModel(null);
+      this.chatError = 'The provider is missing from the chat settings';
+      return;
+    }
+    const provider = settings['provider'] as string;
+    const currentProvider = Private.providers.get(provider) ?? 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.chat = settings;
+      Private.setName('chat', provider);
+      Private.setChatModel(null);
+      return;
+    } else {
+      this._deferredProvider.chat = null;
     }
 
-    if (this._currentProvider?.chatModel !== undefined) {
+    const compatibilityCheck = this.getCompatibilityCheck(provider);
+    if (compatibilityCheck !== undefined) {
+      const error = await compatibilityCheck();
+      if (error !== null) {
+        this.chatError = error.trim();
+        Private.setName('chat', 'None');
+        this._providerChanged.emit('chat');
+        return;
+      }
+    }
+
+    // Get the settings including the secrets.
+    const fullSettings = await this._buildFullSettings(provider, settings);
+
+    if (currentProvider?.chat !== undefined) {
       try {
-        this._chatModel = new this._currentProvider.chatModel({
-          ...fullSettings
-        });
-        this._chatError = '';
+        Private.setChatModel(
+          new currentProvider.chat({
+            ...fullSettings
+          })
+        );
       } catch (e: any) {
-        this._chatError = e.message;
-        this._chatModel = null;
+        this.chatError = e.message;
+        Private.setChatModel(null);
       }
     } else {
-      this._chatModel = null;
+      Private.setChatModel(null);
     }
-    this._name = name;
-    this._providerChanged.emit();
+    Private.setName('chat', provider);
+    this._providerChanged.emit('chat');
   }
 
   /**
    * A signal emitting when the provider or its settings has changed.
    */
-  get providerChanged(): ISignal<IAIProviderRegistry, void> {
+  get providerChanged(): ISignal<IAIProviderRegistry, ModelRole> {
     return this._providerChanged;
   }
 
+  /**
+   * Build a new settings object containing the secrets.
+   */
+  private async _buildFullSettings(
+    name: string,
+    settings: IDict<any>
+  ): Promise<IDict<any>> {
+    // 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
+        );
+        if (secrets !== undefined) {
+          fullSettings[key] = secrets.value;
+        }
+        continue;
+      }
+      fullSettings[key] = settings[key];
+    }
+    return fullSettings;
+  }
+
   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 _providerChanged = new Signal<IAIProviderRegistry, ModelRole>(this);
   private _chatError: string = '';
   private _completerError: string = '';
-  private _providers = new Map<string, IAIProvider>();
-  private _deferredProvider: ISetProviderOptions | null = null;
+  private _notifications: {
+    [key in ModelRole]: Debouncer;
+  };
+  private _deferredProvider: {
+    [key in ModelRole]: ReadonlyPartialJSONObject | null;
+  } = {
+    chat: null,
+    completer: null
+  };
+  private _chatPrompt: string = '';
+  private _completerPrompt: string = '';
 }
 
 export namespace AIProviderRegistry {
@@ -303,21 +459,56 @@ 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.
+   */
+  const names: { [key in ModelRole]: string } = {
+    chat: 'None',
+    completer: 'None'
+  };
+  export function setName(role: ModelRole, value: string): void {
+    names[role] = value;
+  }
+  export function getName(role: ModelRole): string {
+    return names[role];
   }
 
   /**
-   * 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,3 @@
 export * from './panel';
-export * from './settings-connector';
+export * from './textarea';
 export * from './utils';