@jupyterlite/ai 0.7.0 → 0.8.1

package/src/base-completer.ts

@@ -2,13 +2,17 @@ import {
   CompletionHandler,
   IInlineCompletionContext
 } from '@jupyterlab/completer';
+import { BaseLanguageModel } from '@langchain/core/language_models/base';
 import { ReadonlyPartialJSONObject } from '@lumino/coreutils';
 
+import { DEFAULT_COMPLETION_SYSTEM_PROMPT } from './default-prompts';
+import { IAIProviderRegistry } from './tokens';
+
 export interface IBaseCompleter {
   /**
    * The completion prompt.
    */
-  prompt: string;
+  readonly systemPrompt: string;
 
   /**
    * The function to fetch a new completion.
@@ -24,6 +28,33 @@ export interface IBaseCompleter {
   ): Promise<any>;
 }
 
+export abstract class BaseCompleter implements IBaseCompleter {
+  constructor(options: BaseCompleter.IOptions) {
+    this._providerRegistry = options.providerRegistry;
+  }
+
+  /**
+   * Get the system prompt for the completion.
+   */
+  get systemPrompt(): string {
+    return (
+      this._providerRegistry.completerSystemPrompt ??
+      DEFAULT_COMPLETION_SYSTEM_PROMPT
+    );
+  }
+
+  /**
+   * The fetch request for the LLM completer.
+   */
+  abstract fetch(
+    request: CompletionHandler.IRequest,
+    context: IInlineCompletionContext
+  ): Promise<any>;
+
+  protected _providerRegistry: IAIProviderRegistry;
+  protected abstract _completer: BaseLanguageModel<any, any>;
+}
+
 /**
  * The namespace for the base completer.
  */
@@ -32,6 +63,13 @@ export namespace BaseCompleter {
    * The options for the constructor of a completer.
    */
   export interface IOptions {
+    /**
+     * The provider registry.
+     */
+    providerRegistry: IAIProviderRegistry;
+    /**
+     * The settings of the provider.
+     */
     settings: ReadonlyPartialJSONObject;
   }
 }

package/src/chat-handler.ts

@@ -23,8 +23,8 @@ import {
 } from '@langchain/core/messages';
 import { UUID } from '@lumino/coreutils';
 
+import { DEFAULT_CHAT_SYSTEM_PROMPT } from './default-prompts';
 import { jupyternautLiteIcon } from './icons';
-import { chatSystemPrompt } from './provider';
 import { IAIProviderRegistry } from './tokens';
 import { AIChatModel } from './types/ai-model';
 
@@ -39,7 +39,7 @@ export const welcomeMessage = (providers: string[]) => `
 #### Ask JupyterLite AI
 
 
-The provider to use can be set in the settings editor, by selecting it from
+The provider to use can be set in the <button data-commandLinker-command="settingeditor:open" data-commandLinker-args='{"query": "AI provider"}' href="#">settings editor</button>, by selecting it from
 the <img src="${AI_AVATAR}" width="16" height="16"> _AI provider_ settings.
 
 The current providers that are available are _${providers.sort().join('_, _')}_.
@@ -56,15 +56,9 @@ export class ChatHandler extends AbstractChatModel {
   constructor(options: ChatHandler.IOptions) {
     super(options);
     this._providerRegistry = options.providerRegistry;
-    this._prompt = chatSystemPrompt({
-      provider_name: this._providerRegistry.currentName
-    });
 
     this._providerRegistry.providerChanged.connect(() => {
       this._errorMessage = this._providerRegistry.chatError;
-      this._prompt = chatSystemPrompt({
-        provider_name: this._providerRegistry.currentName
-      });
     });
   }
 
@@ -90,13 +84,12 @@ export class ChatHandler extends AbstractChatModel {
   }
 
   /**
-   * Getter and setter for the initial prompt.
+   * Get/set the system prompt for the chat.
    */
-  get prompt(): string {
-    return this._prompt;
-  }
-  set prompt(value: string) {
-    this._prompt = value;
+  get systemPrompt(): string {
+    return (
+      this._providerRegistry.chatSystemPrompt ?? DEFAULT_CHAT_SYSTEM_PROMPT
+    );
   }
 
   async sendMessage(message: INewMessage): Promise<boolean> {
@@ -112,7 +105,7 @@ export class ChatHandler extends AbstractChatModel {
       id: message.id,
       body,
       sender: { username: 'User' },
-      time: Date.now(),
+      time: Private.getTimestampMs(),
       type: 'msg'
     };
     this.messageAdded(msg);
@@ -122,7 +115,7 @@ export class ChatHandler extends AbstractChatModel {
         id: UUID.uuid4(),
         body: `**${this._errorMessage ? this._errorMessage : this._defaultErrorMessage}**`,
         sender: { username: 'ERROR' },
-        time: Date.now(),
+        time: Private.getTimestampMs(),
         type: 'msg'
       };
       this.messageAdded(errorMsg);
@@ -131,7 +124,7 @@ export class ChatHandler extends AbstractChatModel {
 
     this._history.messages.push(msg);
 
-    const messages = mergeMessageRuns([new SystemMessage(this._prompt)]);
+    const messages = mergeMessageRuns([new SystemMessage(this.systemPrompt)]);
     messages.push(
       ...this._history.messages.map(msg => {
         if (msg.sender.username === 'User') {
@@ -149,7 +142,7 @@ export class ChatHandler extends AbstractChatModel {
       id: UUID.uuid4(),
       body: '',
       sender,
-      time: Date.now(),
+      time: Private.getTimestampMs(),
       type: 'msg'
     };
 
@@ -173,7 +166,7 @@ export class ChatHandler extends AbstractChatModel {
         id: UUID.uuid4(),
         body: `**${error}**`,
         sender: { username: 'ERROR' },
-        time: Date.now(),
+        time: Private.getTimestampMs(),
         type: 'msg'
       };
       this.messageAdded(errorMsg);
@@ -206,7 +199,6 @@ export class ChatHandler extends AbstractChatModel {
 
   private _providerRegistry: IAIProviderRegistry;
   private _personaName = 'AI';
-  private _prompt: string;
   private _errorMessage: string = '';
   private _history: IChatHistory = { messages: [] };
   private _defaultErrorMessage = 'AI provider not configured';
@@ -241,7 +233,7 @@ export namespace ChatHandler {
         description: 'Clear the chat'
       }
     ];
-    async getChatCommands(inputModel: IInputModel) {
+    async listCommandCompletions(inputModel: IInputModel) {
       const match = inputModel.currentWord?.match(/^\/\w*/)?.[0];
       if (!match) {
         return [];
@@ -253,12 +245,18 @@ export namespace ChatHandler {
       return commands;
     }
 
-    async handleChatCommand(
-      command: ChatCommand,
-      inputModel: IInputModel
-    ): Promise<void> {
+    async onSubmit(inputModel: IInputModel): Promise<void> {
       // no handling needed because `replaceWith` is set in each command.
       return;
     }
   }
 }
+
+namespace Private {
+  /**
+   * Return the current timestamp in milliseconds.
+   */
+  export function getTimestampMs(): number {
+    return Date.now() / 1000;
+  }
+}

package/src/completion-provider.ts

@@ -28,7 +28,7 @@ export class CompletionProvider implements IInlineCompletionProvider {
    * Get the current completer name.
    */
   get name(): string {
-    return this._providerRegistry.currentName;
+    return this._providerRegistry.currentName('completer');
   }
 
   /**

package/src/default-prompts.ts

@@ -0,0 +1,33 @@
+/*
+ * Copyright (c) Jupyter Development Team.
+ * Distributed under the terms of the Modified BSD License.
+ */
+
+export const DEFAULT_CHAT_SYSTEM_PROMPT = `
+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 $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 DEFAULT_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 only contain the code that will complete the
+give code passed as input, no explanation whatsoever.
+Do not include the prompt in the output, only the string that should be appended to the current input.
+Here is the code to complete:
+`;

package/src/default-providers/Anthropic/completer.ts

@@ -5,24 +5,14 @@ import {
 import { ChatAnthropic } from '@langchain/anthropic';
 import { AIMessage, SystemMessage } from '@langchain/core/messages';
 
-import { BaseCompleter, IBaseCompleter } from '../../base-completer';
-import { COMPLETION_SYSTEM_PROMPT } from '../../provider';
+import { BaseCompleter } from '../../base-completer';
 
-export class AnthropicCompleter implements IBaseCompleter {
+export class AnthropicCompleter extends BaseCompleter {
   constructor(options: BaseCompleter.IOptions) {
+    super(options);
     this._completer = new ChatAnthropic({ ...options.settings });
   }
 
-  /**
-   * 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
@@ -34,7 +24,7 @@ export class AnthropicCompleter implements IBaseCompleter {
     const trimmedPrompt = prompt.trim();
 
     const messages = [
-      new SystemMessage(this._prompt),
+      new SystemMessage(this.systemPrompt),
       new AIMessage(trimmedPrompt)
     ];
 
@@ -65,6 +55,5 @@ export class AnthropicCompleter implements IBaseCompleter {
     }
   }
 
-  private _completer: ChatAnthropic;
-  private _prompt: string = COMPLETION_SYSTEM_PROMPT;
+  protected _completer: ChatAnthropic;
 }

package/src/default-providers/ChromeAI/completer.ts

@@ -5,8 +5,7 @@ import {
 import { ChromeAI } from '@langchain/community/experimental/llms/chrome_ai';
 import { HumanMessage, SystemMessage } from '@langchain/core/messages';
 
-import { BaseCompleter, IBaseCompleter } from '../../base-completer';
-import { COMPLETION_SYSTEM_PROMPT } from '../../provider';
+import { BaseCompleter } from '../../base-completer';
 
 /**
  * Regular expression to match the '```' string at the start of a string.
@@ -29,21 +28,12 @@ const CODE_BLOCK_START_REGEX = /^```(?:[a-zA-Z]+)?\n?/;
  */
 const CODE_BLOCK_END_REGEX = /```$/;
 
-export class ChromeCompleter implements IBaseCompleter {
+export class ChromeCompleter extends BaseCompleter {
   constructor(options: BaseCompleter.IOptions) {
+    super(options);
     this._completer = new ChromeAI({ ...options.settings });
   }
 
-  /**
-   * 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
@@ -54,7 +44,7 @@ export class ChromeCompleter implements IBaseCompleter {
     const trimmedPrompt = prompt.trim();
 
     const messages = [
-      new SystemMessage(this._prompt),
+      new SystemMessage(this.systemPrompt),
       new HumanMessage(trimmedPrompt)
     ];
 
@@ -79,6 +69,5 @@ export class ChromeCompleter implements IBaseCompleter {
     }
   }
 
-  private _completer: ChromeAI;
-  private _prompt: string = COMPLETION_SYSTEM_PROMPT;
+  protected _completer: ChromeAI;
 }

package/src/default-providers/Gemini/completer.ts

@@ -0,0 +1,61 @@
+import {
+  CompletionHandler,
+  IInlineCompletionContext
+} from '@jupyterlab/completer';
+import { ChatGoogleGenerativeAI } from '@langchain/google-genai';
+import { AIMessage, SystemMessage } from '@langchain/core/messages';
+
+import { BaseCompleter } from '../../base-completer';
+
+export class GeminiCompleter extends BaseCompleter {
+  constructor(options: BaseCompleter.IOptions) {
+    super(options);
+    this._completer = new ChatGoogleGenerativeAI({
+      model: 'gemini-pro',
+      ...options.settings
+    });
+  }
+
+  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.systemPrompt),
+      new AIMessage(trimmedPrompt)
+    ];
+
+    try {
+      const response = await this._completer.invoke(messages);
+      const items = [];
+
+      // Gemini 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: [] };
+    }
+  }
+
+  protected _completer: ChatGoogleGenerativeAI;
+}

package/src/default-providers/Gemini/instructions.ts

@@ -0,0 +1,9 @@
+export default `
+<i class="fas fa-exclamation-triangle"></i> This extension is still very much experimental. It is not an official Google extension.
+
+1. Go to <https://aistudio.google.com> and create an API key.
+
+2. Open the JupyterLab settings and go to the **Ai providers** section to select the \`Gemini\`
+   provider and add your API key (required).
+3. Open the chat, or use the inline completer.
+`;

package/src/default-providers/MistralAI/completer.ts

@@ -2,81 +2,68 @@ import {
   CompletionHandler,
   IInlineCompletionContext
 } from '@jupyterlab/completer';
-import {
-  BaseMessage,
-  HumanMessage,
-  SystemMessage
-} from '@langchain/core/messages';
-import { ChatMistralAI } from '@langchain/mistralai';
-import { Throttler } from '@lumino/polling';
+import { MistralAI } from '@langchain/mistralai';
 
-import { BaseCompleter, IBaseCompleter } from '../../base-completer';
-import { COMPLETION_SYSTEM_PROMPT } from '../../provider';
+import { BaseCompleter } from '../../base-completer';
 
-/**
- * The Mistral API has a rate limit of 1 request per second
- */
-const INTERVAL = 1000;
+const CODE_BLOCK_START_REGEX = /^```(?:[a-zA-Z]+)?\n?/;
+const CODE_BLOCK_END_REGEX = /```$/;
 
-export class CodestralCompleter implements IBaseCompleter {
+export class CodestralCompleter extends BaseCompleter {
   constructor(options: BaseCompleter.IOptions) {
-    this._completer = new ChatMistralAI({ ...options.settings });
-    this._throttler = new Throttler(
-      async (messages: BaseMessage[]) => {
-        const response = await this._completer.invoke(messages);
-        // Extract results of completion request.
-        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 }
-    );
-  }
-
-  /**
-   * Getter and setter for the initial prompt.
-   */
-  get prompt(): string {
-    return this._prompt;
-  }
-  set prompt(value: string) {
-    this._prompt = value;
+    super(options);
+    this._completer = new MistralAI({ ...options.settings });
   }
 
   async fetch(
     request: CompletionHandler.IRequest,
     context: IInlineCompletionContext
   ) {
-    const { text, offset: cursorOffset } = request;
-    const prompt = text.slice(0, cursorOffset);
+    try {
+      const { text, offset: cursorOffset } = request;
+      const prompt = this.systemPrompt.concat(text.slice(0, cursorOffset));
+      const suffix = text.slice(cursorOffset);
+      this._controller.abort();
+      this._controller = new AbortController();
 
-    const messages: BaseMessage[] = [
-      new SystemMessage(this._prompt),
-      new HumanMessage(prompt)
-    ];
+      const response = await this._completer.completionWithRetry(
+        {
+          prompt,
+          model: this._completer.model,
+          suffix
+        },
+        { signal: this._controller.signal },
+        false
+      );
+      const items = response.choices.map(choice => {
+        const messageContent = choice.message.content;
+        let content = '';
 
-    try {
-      return await this._throttler.invoke(messages);
+        if (typeof messageContent === 'string') {
+          content = messageContent
+            .replace(CODE_BLOCK_START_REGEX, '')
+            .replace(CODE_BLOCK_END_REGEX, '');
+        } else if (Array.isArray(messageContent)) {
+          // Handle ContentChunk[] case - extract text content
+          content = messageContent
+            .filter(chunk => chunk.type === 'text')
+            .map(chunk => chunk.text || '')
+            .join('')
+            .replace(CODE_BLOCK_START_REGEX, '')
+            .replace(CODE_BLOCK_END_REGEX, '');
+        }
+
+        return {
+          insertText: content
+        };
+      });
+      return { items };
     } catch (error) {
-      console.error('Error fetching completions', error);
+      // the request may be aborted
       return { items: [] };
     }
   }
 
-  private _throttler: Throttler;
-  private _completer: ChatMistralAI;
-  private _prompt: string = COMPLETION_SYSTEM_PROMPT;
+  private _controller = new AbortController();
+  protected _completer: MistralAI;
 }

package/src/default-providers/MistralAI/instructions.ts

@@ -10,6 +10,8 @@ export default `
 
     <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">
 
+**Note:** When using MistralAI for completions, only a subset of models are available. Please check [this resource](https://docs.mistral.ai/api/#tag/fim) to see the list of supported models for completions.
+
 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">

package/src/default-providers/Ollama/completer.ts

@@ -5,24 +5,14 @@ import {
 import { AIMessage, SystemMessage } from '@langchain/core/messages';
 import { ChatOllama } from '@langchain/ollama';
 
-import { BaseCompleter, IBaseCompleter } from '../../base-completer';
-import { COMPLETION_SYSTEM_PROMPT } from '../../provider';
+import { BaseCompleter } from '../../base-completer';
 
-export class OllamaCompleter implements IBaseCompleter {
+export class OllamaCompleter extends BaseCompleter {
   constructor(options: BaseCompleter.IOptions) {
+    super(options);
     this._completer = new ChatOllama({ ...options.settings });
   }
 
-  /**
-   * 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
@@ -30,7 +20,10 @@ export class OllamaCompleter implements IBaseCompleter {
     const { text, offset: cursorOffset } = request;
     const prompt = text.slice(0, cursorOffset);
 
-    const messages = [new SystemMessage(this._prompt), new AIMessage(prompt)];
+    const messages = [
+      new SystemMessage(this.systemPrompt),
+      new AIMessage(prompt)
+    ];
 
     try {
       const response = await this._completer.invoke(messages);
@@ -57,6 +50,5 @@ export class OllamaCompleter implements IBaseCompleter {
     }
   }
 
-  private _completer: ChatOllama;
-  private _prompt: string = COMPLETION_SYSTEM_PROMPT;
+  protected _completer: ChatOllama;
 }

package/src/default-providers/OpenAI/completer.ts

@@ -5,24 +5,14 @@ import {
 import { AIMessage, SystemMessage } from '@langchain/core/messages';
 import { ChatOpenAI } from '@langchain/openai';
 
-import { BaseCompleter, IBaseCompleter } from '../../base-completer';
-import { COMPLETION_SYSTEM_PROMPT } from '../../provider';
+import { BaseCompleter } from '../../base-completer';
 
-export class OpenAICompleter implements IBaseCompleter {
+export class OpenAICompleter extends BaseCompleter {
   constructor(options: BaseCompleter.IOptions) {
+    super(options);
     this._completer = new ChatOpenAI({ ...options.settings });
   }
 
-  /**
-   * 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
@@ -30,7 +20,10 @@ export class OpenAICompleter implements IBaseCompleter {
     const { text, offset: cursorOffset } = request;
     const prompt = text.slice(0, cursorOffset);
 
-    const messages = [new SystemMessage(this._prompt), new AIMessage(prompt)];
+    const messages = [
+      new SystemMessage(this.systemPrompt),
+      new AIMessage(prompt)
+    ];
 
     try {
       const response = await this._completer.invoke(messages);
@@ -57,6 +50,5 @@ export class OpenAICompleter implements IBaseCompleter {
     }
   }
 
-  private _completer: ChatOpenAI;
-  private _prompt: string = COMPLETION_SYSTEM_PROMPT;
+  protected _completer: ChatOpenAI;
 }