@jupyterlite/ai 0.2.0 → 0.4.0

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/llm-models/openai-completer.ts

@@ -0,0 +1,67 @@
+import {
+  CompletionHandler,
+  IInlineCompletionContext
+} from '@jupyterlab/completer';
+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';
+
+export class OpenAICompleter implements IBaseCompleter {
+  constructor(options: BaseCompleter.IOptions) {
+    this._openAIProvider = new ChatOpenAI({ ...options.settings });
+  }
+
+  get provider(): BaseChatModel {
+    return this._openAIProvider;
+  }
+
+  /**
+   * Getter and setter for the initial prompt.
+   */
+  get prompt(): string {
+    return this._prompt;
+  }
+  set prompt(value: string) {
+    this._prompt = value;
+  }
+
+  async fetch(
+    request: CompletionHandler.IRequest,
+    context: IInlineCompletionContext
+  ) {
+    const { text, offset: cursorOffset } = request;
+    const prompt = text.slice(0, cursorOffset);
+
+    const messages = [new SystemMessage(this._prompt), new AIMessage(prompt)];
+
+    try {
+      const response = await this._openAIProvider.invoke(messages);
+      const items = [];
+      if (typeof response.content === 'string') {
+        items.push({
+          insertText: response.content
+        });
+      } else {
+        response.content.forEach(content => {
+          if (content.type !== 'text') {
+            return;
+          }
+          items.push({
+            insertText: content.text,
+            filterText: prompt.substring(prompt.length)
+          });
+        });
+      }
+      return { items };
+    } catch (error) {
+      console.error('Error fetching completions', error);
+      return { items: [] };
+    }
+  }
+
+  private _openAIProvider: ChatOpenAI;
+  private _prompt: string = COMPLETION_SYSTEM_PROMPT;
+}

package/src/llm-models/svg.d.ts

@@ -0,0 +1,9 @@
+/*
+ * Copyright (c) Jupyter Development Team.
+ * Distributed under the terms of the Modified BSD License.
+ */
+
+declare module '*.svg' {
+  const value: string;
+  export default value;
+}

package/src/provider.ts

@@ -4,44 +4,113 @@ 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';
-
-export class AIProvider implements IAIProvider {
-  constructor(options: AIProvider.IOptions) {
-    this._completionProvider = new CompletionProvider({
-      name: 'None',
-      settings: {},
-      requestCompletion: options.requestCompletion
-    });
-    options.completionProviderManager.registerInlineProvider(
-      this._completionProvider
-    );
+import { IBaseCompleter } from './llm-models';
+import { IAIProvider, IAIProviderRegistry } from './tokens';
+import { JSONSchema7 } from 'json-schema';
+
+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.
+`;
+
+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;
   }
 
   /**
@@ -59,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.
    */
@@ -110,6 +195,16 @@ export namespace AIProvider {
     requestCompletion: () => void;
   }
 
+  /**
+   * The options for the Chat system prompt.
+   */
+  export interface IPromptOptions {
+    /**
+     * The provider name.
+     */
+    provider_name: string;
+  }
+
   /**
    * This function indicates whether a key is writable in an object.
    * https://stackoverflow.com/questions/54724875/can-we-check-whether-property-is-readonly-in-typescript

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
+};

package/src/settings/panel.tsx

@@ -0,0 +1,257 @@
+import { IRenderMimeRegistry } from '@jupyterlab/rendermime';
+import { ISettingRegistry } from '@jupyterlab/settingregistry';
+import { FormComponent, IFormRenderer } from '@jupyterlab/ui-components';
+import { JSONExt } from '@lumino/coreutils';
+import { IChangeEvent } from '@rjsf/core';
+import type { FieldProps } from '@rjsf/utils';
+import validator from '@rjsf/validator-ajv8';
+import { JSONSchema7 } from 'json-schema';
+import React from 'react';
+
+import baseSettings from './schemas/base.json';
+import { IAIProviderRegistry, IDict } from '../tokens';
+
+const MD_MIME_TYPE = 'text/markdown';
+const STORAGE_NAME = '@jupyterlite/ai:settings';
+const INSTRUCTION_CLASS = 'jp-AISettingsInstructions';
+
+export const aiSettingsRenderer = (options: {
+  providerRegistry: IAIProviderRegistry;
+  rmRegistry?: IRenderMimeRegistry;
+}): IFormRenderer => {
+  return {
+    fieldRenderer: (props: FieldProps) => {
+      props.formContext = { ...props.formContext, ...options };
+      return <AiSettings {...props} />;
+    }
+  };
+};
+
+export interface ISettingsFormStates {
+  schema: JSONSchema7;
+  instruction: HTMLElement | null;
+}
+
+const WrappedFormComponent = (props: any): JSX.Element => {
+  return <FormComponent {...props} validator={validator} />;
+};
+
+export class AiSettings extends React.Component<
+  FieldProps,
+  ISettingsFormStates
+> {
+  constructor(props: FieldProps) {
+    super(props);
+    if (!props.formContext.providerRegistry) {
+      throw new Error(
+        'The provider registry is needed to enable the jupyterlite-ai settings panel'
+      );
+    }
+    this._providerRegistry = props.formContext.providerRegistry;
+    this._rmRegistry = props.formContext.rmRegistry ?? null;
+    this._settings = props.formContext.settings;
+
+    // Initialize the providers schema.
+    const providerSchema = JSONExt.deepCopy(baseSettings) as any;
+    providerSchema.properties.provider = {
+      type: 'string',
+      title: 'Provider',
+      description: 'The AI provider to use for chat and completion',
+      default: 'None',
+      enum: ['None'].concat(this._providerRegistry.providers)
+    };
+    this._providerSchema = providerSchema as JSONSchema7;
+
+    // Check if there is saved values in local storage, otherwise use the settings from
+    // the setting registry (led to default if there are no user settings).
+    const storageSettings = localStorage.getItem(STORAGE_NAME);
+    if (storageSettings === null) {
+      const labSettings = this._settings.get('AIprovider').composite;
+      if (labSettings && Object.keys(labSettings).includes('provider')) {
+        // Get the provider name.
+        const provider = Object.entries(labSettings).find(
+          v => v[0] === 'provider'
+        )?.[1] as string;
+        // Save the settings.
+        const settings: any = {
+          _current: provider
+        };
+        settings[provider] = labSettings;
+        localStorage.setItem(STORAGE_NAME, JSON.stringify(settings));
+      }
+    }
+
+    // Initialize the settings from the saved ones.
+    this._provider = this.getCurrentProvider();
+    this._currentSettings = this.getSettings();
+
+    // Initialize the schema.
+    const schema = this._buildSchema();
+    this.state = { schema, instruction: null };
+
+    this._renderInstruction();
+
+    // Update the setting registry.
+    this._settings
+      .set('AIprovider', this._currentSettings)
+      .catch(console.error);
+  }
+
+  /**
+   * Get the current provider from the local storage.
+   */
+  getCurrentProvider(): string {
+    const settings = JSON.parse(localStorage.getItem(STORAGE_NAME) || '{}');
+    return settings['_current'] ?? 'None';
+  }
+
+  /**
+   * Save the current provider to the local storage.
+   */
+  saveCurrentProvider(): void {
+    const settings = JSON.parse(localStorage.getItem(STORAGE_NAME) || '{}');
+    settings['_current'] = this._provider;
+    localStorage.setItem(STORAGE_NAME, JSON.stringify(settings));
+  }
+
+  /**
+   * Get settings from local storage for a given provider.
+   */
+  getSettings(): IDict<any> {
+    const settings = JSON.parse(localStorage.getItem(STORAGE_NAME) || '{}');
+    return settings[this._provider] ?? { provider: this._provider };
+  }
+
+  /**
+   * Save settings in local storage for a given provider.
+   */
+  saveSettings(value: IDict<any>) {
+    const settings = JSON.parse(localStorage.getItem(STORAGE_NAME) ?? '{}');
+    settings[this._provider] = value;
+    localStorage.setItem(STORAGE_NAME, JSON.stringify(settings));
+  }
+
+  /**
+   * 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.
+   */
+  private _buildSchema(): JSONSchema7 {
+    const schema = JSONExt.deepCopy(baseSettings) as any;
+    this._uiSchema = {};
+    const settingsSchema = this._providerRegistry.getSettingsSchema(
+      this._provider
+    );
+
+    if (settingsSchema) {
+      Object.entries(settingsSchema).forEach(([key, value]) => {
+        schema.properties[key] = value;
+        this._updateUiSchema(key);
+      });
+    }
+    return schema as JSONSchema7;
+  }
+
+  /**
+   * Update the schema state for the given provider, that trigger the re-rendering of
+   * the component.
+   */
+  private _updateSchema() {
+    const schema = this._buildSchema();
+    this.setState({ schema });
+  }
+
+  /**
+   * Render the markdown instructions for the current provider.
+   */
+  private async _renderInstruction(): Promise<void> {
+    let instructions = this._providerRegistry.getInstructions(this._provider);
+    if (!this._rmRegistry || !instructions) {
+      this.setState({ instruction: null });
+      return;
+    }
+    instructions = `---\n\n${instructions}\n\n---`;
+    const renderer = this._rmRegistry.createRenderer(MD_MIME_TYPE);
+    const model = this._rmRegistry.createModel({
+      data: { [MD_MIME_TYPE]: instructions }
+    });
+    await renderer.renderModel(model);
+    this.setState({ instruction: renderer.node });
+  }
+
+  /**
+   * Triggered when the provider hes changed, to update the schema and values.
+   * Update the Jupyterlab settings accordingly.
+   */
+  private _onProviderChanged = (e: IChangeEvent) => {
+    const provider = e.formData.provider;
+    if (provider === this._currentSettings.provider) {
+      return;
+    }
+    this._provider = provider;
+    this.saveCurrentProvider();
+    this._currentSettings = this.getSettings();
+    this._updateSchema();
+    this._renderInstruction();
+    this._settings
+      .set('AIprovider', { provider: this._provider, ...this._currentSettings })
+      .catch(console.error);
+  };
+
+  /**
+   * Triggered when the form value has changed, to update the current settings and save
+   * it in local storage.
+   * Update the Jupyterlab settings accordingly.
+   */
+  private _onFormChange = (e: IChangeEvent) => {
+    this._currentSettings = JSONExt.deepCopy(e.formData);
+    this.saveSettings(this._currentSettings);
+    this._settings
+      .set('AIprovider', { provider: this._provider, ...this._currentSettings })
+      .catch(console.error);
+  };
+
+  render(): JSX.Element {
+    return (
+      <>
+        <WrappedFormComponent
+          formData={{ provider: this._provider }}
+          schema={this._providerSchema}
+          onChange={this._onProviderChanged}
+        />
+        {this.state.instruction !== null && (
+          <details>
+            <summary className={INSTRUCTION_CLASS}>Instructions</summary>
+            <span
+              ref={node =>
+                node && node.replaceChildren(this.state.instruction!)
+              }
+            />
+          </details>
+        )}
+        <WrappedFormComponent
+          formData={this._currentSettings}
+          schema={this.state.schema}
+          onChange={this._onFormChange}
+          uiSchema={this._uiSchema}
+        />
+      </>
+    );
+  }
+
+  private _providerRegistry: IAIProviderRegistry;
+  private _provider: string;
+  private _providerSchema: JSONSchema7;
+  private _rmRegistry: IRenderMimeRegistry | null;
+  private _currentSettings: IDict<any> = { provider: 'None' };
+  private _uiSchema: IDict<any> = {};
+  private _settings: ISettingRegistry.ISettings;
+}

package/src/settings/schemas/index.ts

@@ -0,0 +1,15 @@
+import { IDict } from '../../tokens';
+
+import ChromeAI from './_generated/ChromeAI.json';
+import MistralAI from './_generated/MistralAI.json';
+import Anthropic from './_generated/Anthropic.json';
+import OpenAI from './_generated/OpenAI.json';
+
+const ProviderSettings: IDict<any> = {
+  ChromeAI,
+  MistralAI,
+  Anthropic,
+  OpenAI
+};
+
+export { ProviderSettings };