@jupyterlite/ai 0.2.0 → 0.4.0

package/src/completion-provider.ts

@@ -3,10 +3,9 @@ import {
   IInlineCompletionContext,
   IInlineCompletionProvider
 } from '@jupyterlab/completer';
-import { LLM } from '@langchain/core/language_models/llms';
 
-import { getCompleter, IBaseCompleter, BaseCompleter } from './llm-models';
-import { ReadonlyPartialJSONObject } from '@lumino/coreutils';
+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(): LLM | 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/icons.ts

@@ -0,0 +1,18 @@
+/*
+ * Copyright (c) Jupyter Development Team.
+ * Distributed under the terms of the Modified BSD License.
+ */
+
+import { LabIcon } from '@jupyterlab/ui-components';
+
+/**
+ * This icon is based on the jupyternaut icon from Jupyter AI:
+ * https://github.com/jupyterlab/jupyter-ai/blob/main/packages/jupyter-ai/style/icons/jupyternaut.svg
+ * With a small tweak for the colors to match the JupyterLite icon.
+ */
+import jupyternautLiteSvg from '../style/icons/jupyternaut-lite.svg';
+
+export const jupyternautLiteIcon = new LabIcon({
+  name: '@jupyterlite/ai:jupyternaut-lite',
+  svgstr: jupyternautLiteSvg
+});

package/src/index.ts

@@ -1,8 +1,11 @@
 import {
   ActiveCellManager,
+  AutocompletionRegistry,
   buildChatSidebar,
   buildErrorWidget,
-  IActiveCellManager
+  IActiveCellManager,
+  IAutocompletionCommandsProps,
+  IAutocompletionRegistry
 } from '@jupyter/chat';
 import {
   JupyterFrontEnd,
@@ -13,21 +16,55 @@ 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 { AIProvider } from './provider';
-import { IAIProvider } from './token';
+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 { IAIProviderRegistry } from './tokens';
+
+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 chatPlugin: JupyterFrontEndPlugin<void> = {
   id: '@jupyterlite/ai:chat',
   description: 'LLM chat extension',
   autoStart: true,
+  requires: [IAIProviderRegistry, IRenderMimeRegistry, IAutocompletionRegistry],
   optional: [INotebookTracker, ISettingRegistry, IThemeManager],
-  requires: [IAIProvider, IRenderMimeRegistry],
   activate: async (
     app: JupyterFrontEnd,
-    aiProvider: IAIProvider,
+    providerRegistry: IAIProviderRegistry,
     rmRegistry: IRenderMimeRegistry,
+    autocompletionRegistry: IAutocompletionRegistry,
     notebookTracker: INotebookTracker | null,
     settingsRegistry: ISettingRegistry | null,
     themeManager: IThemeManager | null
@@ -41,18 +78,23 @@ const chatPlugin: JupyterFrontEndPlugin<void> = {
     }
 
     const chatHandler = new ChatHandler({
-      aiProvider: aiProvider,
-      activeCellManager: activeCellManager
+      providerRegistry,
+      activeCellManager
     });
 
     let sendWithShiftEnter = false;
     let enableCodeToolbar = true;
+    let personaName = 'AI';
 
     function loadSetting(setting: ISettingRegistry.ISettings): void {
       sendWithShiftEnter = setting.get('sendWithShiftEnter')
         .composite as boolean;
       enableCodeToolbar = setting.get('enableCodeToolbar').composite as boolean;
+      personaName = setting.get('personaName').composite as string;
+
+      // set the properties
       chatHandler.config = { sendWithShiftEnter, enableCodeToolbar };
+      chatHandler.personaName = personaName;
     }
 
     Promise.all([app.restored, settingsRegistry?.load(chatPlugin.id)])
@@ -77,9 +119,10 @@ const chatPlugin: JupyterFrontEndPlugin<void> = {
       chatWidget = buildChatSidebar({
         model: chatHandler,
         themeManager,
-        rmRegistry
+        rmRegistry,
+        autocompletionRegistry
       });
-      chatWidget.title.caption = 'Codestral Chat';
+      chatWidget.title.caption = 'Jupyterlite AI Chat';
     } catch (e) {
       chatWidget = buildErrorWidget(themeManager);
     }
@@ -90,27 +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);
+  }
+};
 
+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;
-          aiProvider.setModels(provider, settings.composite);
+          // 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());
@@ -118,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, aiProviderPlugin];
+export default [
+  providerRegistryPlugin,
+  autocompletionRegistryPlugin,
+  chatPlugin,
+  completerPlugin
+];

package/src/llm-models/anthropic-completer.ts

@@ -0,0 +1,75 @@
+import {
+  CompletionHandler,
+  IInlineCompletionContext
+} from '@jupyterlab/completer';
+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';
+
+export class AnthropicCompleter implements IBaseCompleter {
+  constructor(options: BaseCompleter.IOptions) {
+    this._anthropicProvider = new ChatAnthropic({ ...options.settings });
+  }
+
+  get provider(): BaseChatModel {
+    return this._anthropicProvider;
+  }
+
+  /**
+   * 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);
+
+    // Anthropic does not allow whitespace at the end of the AIMessage
+    const trimmedPrompt = prompt.trim();
+
+    const messages = [
+      new SystemMessage(this._prompt),
+      new AIMessage(trimmedPrompt)
+    ];
+
+    try {
+      const response = await this._anthropicProvider.invoke(messages);
+      const items = [];
+
+      // Anthropic can return string or complex content, a list of string/images/other.
+      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(trimmedPrompt.length)
+          });
+        });
+      }
+      return { items };
+    } catch (error) {
+      console.error('Error fetching completions', error);
+      return { items: [] };
+    }
+  }
+
+  private _anthropicProvider: ChatAnthropic;
+  private _prompt: string = COMPLETION_SYSTEM_PROMPT;
+}

package/src/llm-models/base-completer.ts

@@ -2,14 +2,19 @@ import {
   CompletionHandler,
   IInlineCompletionContext
 } from '@jupyterlab/completer';
-import { LLM } from '@langchain/core/language_models/llms';
+import { BaseLanguageModel } from '@langchain/core/language_models/base';
 import { ReadonlyPartialJSONObject } from '@lumino/coreutils';
 
 export interface IBaseCompleter {
   /**
    * The LLM completer.
    */
-  provider: LLM;
+  provider: BaseLanguageModel;
+
+  /**
+   * The completion prompt.
+   */
+  prompt: string;
 
   /**
    * The function to fetch a new completion.

package/src/llm-models/chrome-completer.ts

@@ -0,0 +1,88 @@
+import {
+  CompletionHandler,
+  IInlineCompletionContext
+} from '@jupyterlab/completer';
+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';
+
+/**
+ * Regular expression to match the '```' string at the start of a string.
+ * So the completions returned by the LLM can still be kept after removing the code block formatting.
+ *
+ * For example, if the response contains the following content after typing `import pandas`:
+ *
+ * ```python
+ * as pd
+ * ```
+ *
+ * The formatting string after removing the code block delimiters will be:
+ *
+ * as pd
+ */
+const CODE_BLOCK_START_REGEX = /^```(?:[a-zA-Z]+)?\n?/;
+
+/**
+ * Regular expression to match the '```' string at the end of a string.
+ */
+const CODE_BLOCK_END_REGEX = /```$/;
+
+export class ChromeCompleter implements IBaseCompleter {
+  constructor(options: BaseCompleter.IOptions) {
+    this._chromeProvider = new ChromeAI({ ...options.settings });
+  }
+
+  /**
+   * Getter and setter for the initial prompt.
+   */
+  get prompt(): string {
+    return this._prompt;
+  }
+  set prompt(value: string) {
+    this._prompt = value;
+  }
+
+  get provider(): LLM {
+    return this._chromeProvider;
+  }
+
+  async fetch(
+    request: CompletionHandler.IRequest,
+    context: IInlineCompletionContext
+  ) {
+    const { text, offset: cursorOffset } = request;
+    const prompt = text.slice(0, cursorOffset);
+
+    const trimmedPrompt = prompt.trim();
+
+    const messages = [
+      new SystemMessage(this._prompt),
+      new HumanMessage(trimmedPrompt)
+    ];
+
+    try {
+      let response = await this._chromeProvider.invoke(messages);
+
+      // ChromeAI sometimes returns a string starting with '```',
+      // so process the response to remove the code block delimiters
+      if (CODE_BLOCK_START_REGEX.test(response)) {
+        response = response
+          .replace(CODE_BLOCK_START_REGEX, '')
+          .replace(CODE_BLOCK_END_REGEX, '');
+      }
+
+      const items = [{ insertText: response }];
+      return {
+        items
+      };
+    } catch (error) {
+      console.error('Error fetching completion:', error);
+      return { items: [] };
+    }
+  }
+
+  private _chromeProvider: ChromeAI;
+  private _prompt: string = COMPLETION_SYSTEM_PROMPT;
+}

package/src/llm-models/codestral-completer.ts

@@ -2,72 +2,63 @@ import {
   CompletionHandler,
   IInlineCompletionContext
 } from '@jupyterlab/completer';
-import { LLM } from '@langchain/core/language_models/llms';
-import { MistralAI } from '@langchain/mistralai';
+import { BaseChatModel } from '@langchain/core/language_models/chat_models';
+import {
+  BaseMessage,
+  HumanMessage,
+  SystemMessage
+} from '@langchain/core/messages';
+import { ChatMistralAI } from '@langchain/mistralai';
 import { Throttler } from '@lumino/polling';
-import { CompletionRequest } from '@mistralai/mistralai';
 
 import { BaseCompleter, IBaseCompleter } from './base-completer';
+import { COMPLETION_SYSTEM_PROMPT } from '../provider';
 
 /**
  * The Mistral API has a rate limit of 1 request per second
  */
 const INTERVAL = 1000;
 
-/**
- * Timeout to avoid endless requests
- */
-const REQUEST_TIMEOUT = 3000;
-
 export class CodestralCompleter implements IBaseCompleter {
   constructor(options: BaseCompleter.IOptions) {
-    // this._requestCompletion = options.requestCompletion;
-    this._mistralProvider = new MistralAI({ ...options.settings });
+    this._mistralProvider = new ChatMistralAI({ ...options.settings });
     this._throttler = new Throttler(
-      async (data: CompletionRequest) => {
-        const invokedData = data;
-
-        // Request completion.
-        const request = this._mistralProvider.completionWithRetry(
-          data,
-          {},
-          false
-        );
-        const timeoutPromise = new Promise<null>(resolve => {
-          return setTimeout(() => resolve(null), REQUEST_TIMEOUT);
-        });
-
-        // Fetch again if the request is too long or if the prompt has changed.
-        const response = await Promise.race([request, timeoutPromise]);
-        if (
-          response === null ||
-          invokedData.prompt !== this._currentData?.prompt
-        ) {
-          return {
-            items: [],
-            fetchAgain: true
-          };
-        }
-
+      async (messages: BaseMessage[]) => {
+        const response = await this._mistralProvider.invoke(messages);
         // Extract results of completion request.
-        const items = response.choices.map((choice: any) => {
-          return { insertText: choice.message.content as string };
-        });
-
-        return {
-          items
-        };
+        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
+            });
+          });
+        }
+        return { items };
       },
       { limit: INTERVAL }
     );
   }
 
-  get provider(): LLM {
+  get provider(): BaseChatModel {
     return this._mistralProvider;
   }
 
-  set requestCompletion(value: () => void) {
-    this._requestCompletion = value;
+  /**
+   * Getter and setter for the initial prompt.
+   */
+  get prompt(): string {
+    return this._prompt;
+  }
+  set prompt(value: string) {
+    this._prompt = value;
   }
 
   async fetch(
@@ -76,38 +67,21 @@ export class CodestralCompleter implements IBaseCompleter {
   ) {
     const { text, offset: cursorOffset } = request;
     const prompt = text.slice(0, cursorOffset);
-    const suffix = text.slice(cursorOffset);
 
-    const data = {
-      prompt,
-      suffix,
-      model: this._mistralProvider.model,
-      // temperature: 0,
-      // top_p: 1,
-      // max_tokens: 1024,
-      // min_tokens: 0,
-      stream: false,
-      // random_seed: 1337,
-      stop: []
-    };
+    const messages: BaseMessage[] = [
+      new SystemMessage(this._prompt),
+      new HumanMessage(prompt)
+    ];
 
     try {
-      this._currentData = data;
-      const completionResult = await this._throttler.invoke(data);
-      if (completionResult.fetchAgain) {
-        if (this._requestCompletion) {
-          this._requestCompletion();
-        }
-      }
-      return { items: completionResult.items };
+      return await this._throttler.invoke(messages);
     } catch (error) {
       console.error('Error fetching completions', error);
       return { items: [] };
     }
   }
 
-  private _requestCompletion?: () => void;
   private _throttler: Throttler;
-  private _mistralProvider: MistralAI;
-  private _currentData: CompletionRequest | null = null;
+  private _mistralProvider: ChatMistralAI;
+  private _prompt: string = COMPLETION_SYSTEM_PROMPT;
 }