@jupyterlite/ai 0.2.0 → 0.3.0

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,
@@ -15,19 +18,49 @@ import { IRenderMimeRegistry } from '@jupyterlab/rendermime';
 import { ISettingRegistry } from '@jupyterlab/settingregistry';
 
 import { ChatHandler } from './chat-handler';
+import { getSettings } from './llm-models';
 import { AIProvider } from './provider';
+import { renderSlashCommandOption } from './slash-commands';
 import { IAIProvider } from './token';
 
+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: [IAIProvider, IRenderMimeRegistry, IAutocompletionRegistry],
   optional: [INotebookTracker, ISettingRegistry, IThemeManager],
-  requires: [IAIProvider, IRenderMimeRegistry],
   activate: async (
     app: JupyterFrontEnd,
     aiProvider: IAIProvider,
     rmRegistry: IRenderMimeRegistry,
+    autocompletionRegistry: IAutocompletionRegistry,
     notebookTracker: INotebookTracker | null,
     settingsRegistry: ISettingRegistry | null,
     themeManager: IThemeManager | null
@@ -47,12 +80,17 @@ const chatPlugin: JupyterFrontEndPlugin<void> = {
 
     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 +115,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);
     }
@@ -105,11 +144,34 @@ const aiProviderPlugin: JupyterFrontEndPlugin<IAIProvider> = {
       requestCompletion: () => app.commands.execute('inline-completer:invoke')
     });
 
+    let currentProvider = 'None';
     settingRegistry
       .load(aiProviderPlugin.id)
       .then(settings => {
         const updateProvider = () => {
           const provider = settings.get('provider').composite as string;
+          if (provider !== currentProvider) {
+            // Update the settings panel.
+            currentProvider = provider;
+            const settingsProperties = settings.schema.properties;
+            if (settingsProperties) {
+              const schemaKeys = Object.keys(settingsProperties);
+              schemaKeys.forEach(key => {
+                if (key !== 'provider') {
+                  delete settings.schema.properties?.[key];
+                }
+              });
+              const properties = getSettings(provider);
+              if (properties === null) {
+                return;
+              }
+              Object.entries(properties).forEach(([name, value], index) => {
+                settingsProperties[name] = value as ISettingRegistry.IProperty;
+              });
+            }
+          }
+
+          // Update the settings to the AI providers.
           aiProvider.setModels(provider, settings.composite);
         };
 
@@ -127,4 +189,4 @@ const aiProviderPlugin: JupyterFrontEndPlugin<IAIProvider> = {
   }
 };
 
-export default [chatPlugin, aiProviderPlugin];
+export default [chatPlugin, autocompletionRegistryPlugin, aiProviderPlugin];

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

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

@@ -1,8 +1,20 @@
+import { ChatAnthropic } from '@langchain/anthropic';
+import { ChromeAI } from '@langchain/community/experimental/llms/chrome_ai';
 import { BaseChatModel } from '@langchain/core/language_models/chat_models';
 import { ChatMistralAI } from '@langchain/mistralai';
+import { ChatOpenAI } from '@langchain/openai';
+
 import { IBaseCompleter } from './base-completer';
+import { AnthropicCompleter } from './anthropic-completer';
 import { CodestralCompleter } from './codestral-completer';
 import { ReadonlyPartialJSONObject } from '@lumino/coreutils';
+import { ChromeCompleter } from './chrome-completer';
+import { OpenAICompleter } from './openai-completer';
+
+import chromeAI from '../_provider-settings/chromeAI.json';
+import mistralAI from '../_provider-settings/mistralAI.json';
+import anthropic from '../_provider-settings/anthropic.json';
+import openAI from '../_provider-settings/openAI.json';
 
 /**
  * Get an LLM completer from the name.
@@ -13,6 +25,12 @@ export function getCompleter(
 ): IBaseCompleter | null {
   if (name === 'MistralAI') {
     return new CodestralCompleter({ settings });
+  } else if (name === 'Anthropic') {
+    return new AnthropicCompleter({ settings });
+  } else if (name === 'ChromeAI') {
+    return new ChromeCompleter({ settings });
+  } else if (name === 'OpenAI') {
+    return new OpenAICompleter({ settings });
   }
   return null;
 }
@@ -26,6 +44,14 @@ export function getChatModel(
 ): BaseChatModel | null {
   if (name === 'MistralAI') {
     return new ChatMistralAI({ ...settings });
+  } else if (name === 'Anthropic') {
+    return new ChatAnthropic({ ...settings });
+  } else if (name === 'ChromeAI') {
+    // TODO: fix
+    // @ts-expect-error: missing properties
+    return new ChromeAI({ ...settings });
+  } else if (name === 'OpenAI') {
+    return new ChatOpenAI({ ...settings });
   }
   return null;
 }
@@ -36,6 +62,29 @@ export function getChatModel(
 export function getErrorMessage(name: string, error: any): string {
   if (name === 'MistralAI') {
     return error.message;
+  } else if (name === 'Anthropic') {
+    return error.error.error.message;
+  } else if (name === 'ChromeAI') {
+    return error.message;
+  } else if (name === 'OpenAI') {
+    return error.message;
   }
   return 'Unknown provider';
 }
+
+/*
+ * Get an LLM completer from the name.
+ */
+export function getSettings(name: string): any {
+  if (name === 'MistralAI') {
+    return mistralAI.properties;
+  } else if (name === 'Anthropic') {
+    return anthropic.properties;
+  } else if (name === 'ChromeAI') {
+    return chromeAI.properties;
+  } else if (name === 'OpenAI') {
+    return openAI.properties;
+  }
+
+  return null;
+}