@jupyterlite/ai 0.6.2 → 0.8.0

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

@@ -3,31 +3,16 @@ import {
   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';
+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 });
   }
 
-  get completer(): BaseChatModel {
-    return this._completer;
-  }
-
-  /**
-   * 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
@@ -39,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)
     ];
 
@@ -70,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

@@ -3,11 +3,9 @@ import {
   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';
+import { BaseCompleter } from '../../base-completer';
 
 /**
  * Regular expression to match the '```' string at the start of a string.
@@ -30,25 +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;
-  }
-
-  get completer(): LLM {
-    return this._completer;
-  }
-
   async fetch(
     request: CompletionHandler.IRequest,
     context: IInlineCompletionContext
@@ -59,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)
     ];
 
@@ -84,6 +69,5 @@ export class ChromeCompleter implements IBaseCompleter {
     }
   }
 
-  private _completer: ChromeAI;
-  private _prompt: string = COMPLETION_SYSTEM_PROMPT;
+  protected _completer: ChromeAI;
 }

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

@@ -22,3 +22,24 @@ 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>
 `;
+
+/**
+ * Check if the browser supports ChromeAI and the model is available.
+ */
+export async function compatibilityCheck(): Promise<string | null> {
+  // Check if the browser supports the ChromeAI model
+  if (
+    typeof window === 'undefined' ||
+    !('LanguageModel' in window) ||
+    window.LanguageModel === undefined ||
+    (window.LanguageModel as any).availability === undefined
+  ) {
+    return 'Your browser does not support ChromeAI. Please use an updated chrome based browser like Google Chrome, and follow the instructions in settings to enable it.';
+  }
+  const languageModel = window.LanguageModel as any;
+  if (!(await languageModel.availability())) {
+    return 'The ChromeAI model is not available in your browser. Please ensure you have enabled the necessary flags in Google Chrome as described in the instructions in settings.';
+  }
+  // If the model is available, return null to indicate compatibility
+  return null;
+}

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,86 +2,68 @@ import {
   CompletionHandler,
   IInlineCompletionContext
 } from '@jupyterlab/completer';
-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 { 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 }
-    );
-  }
-
-  get completer(): BaseChatModel {
-    return this._completer;
-  }
-
-  /**
-   * 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

@@ -0,0 +1,54 @@
+import {
+  CompletionHandler,
+  IInlineCompletionContext
+} from '@jupyterlab/completer';
+import { AIMessage, SystemMessage } from '@langchain/core/messages';
+import { ChatOllama } from '@langchain/ollama';
+
+import { BaseCompleter } from '../../base-completer';
+
+export class OllamaCompleter extends BaseCompleter {
+  constructor(options: BaseCompleter.IOptions) {
+    super(options);
+    this._completer = new ChatOllama({ ...options.settings });
+  }
+
+  async fetch(
+    request: CompletionHandler.IRequest,
+    context: IInlineCompletionContext
+  ) {
+    const { text, offset: cursorOffset } = request;
+    const prompt = text.slice(0, cursorOffset);
+
+    const messages = [
+      new SystemMessage(this.systemPrompt),
+      new AIMessage(prompt)
+    ];
+
+    try {
+      const response = await this._completer.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: [] };
+    }
+  }
+
+  protected _completer: ChatOllama;
+}

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

@@ -0,0 +1,70 @@
+export default `
+Ollama allows to run large language models locally on your machine.
+To use it you need to install the Ollama CLI and pull the model you want to use.
+
+1. Install the Ollama CLI by following the instructions at <https://ollama.com/download>
+
+2. Pull the model you want to use by running the following command in your terminal:
+
+   \`\`\`bash
+   ollama pull <model-name>
+   \`\`\`
+
+   For example, to pull the Llama 2 model, run:
+
+   \`\`\`bash
+   ollama pull llama2
+   \`\`\`
+
+3. Once the model is pulled, you can use it in your application by running the following command:
+
+    \`\`\`bash
+    ollama serve
+    \`\`\`
+
+4. This model will be available in the extension, using the model name you used in the command above.
+
+<details>
+<summary>Deploying Lite/Lab on external server</summary>
+
+See https://objectgraph.com/blog/ollama-cors/ for more details.
+
+On Linux, you can run the following commands:
+
+1. Check if CORS is enabled on the server. You can do this by running the following command in your terminal:
+
+   \`\`\`bash
+   curl -X OPTIONS http://localhost:11434 -H "Origin: http://example.com" -H "Access-Control-Request-Method: GET" -I
+   \`\`\`
+
+   If CORS is disabled, you will see a response like this:
+
+   \`\`\`bash
+   HTTP/1.1 403 Forbidden
+   Date: Wed, 09 Oct 2024 10:12:15 GMT
+   Content-Length: 0
+    \`\`\`
+
+2. If CORS is not enabled, update _/etc/systemd/system/ollama.service_ with:
+
+   \`\`\`bash
+   [Service]
+   Environment="OLLAMA_HOST=0.0.0.0"
+   Environment="OLLAMA_ORIGINS=*"
+   \`\`\`
+
+3. Restart the service:
+
+   \`\`\`bash
+   sudo systemctl daemon-reload
+   sudo systemctl restart ollama
+   \`\`\`
+
+4. Check if CORS is enabled on the server again by running the following command in your terminal:
+
+    \`\`\`bash
+    curl -X OPTIONS http://localhost:11434 -H "Origin: http://example.com" -H "Access-Control-Request-Method: GET" -I
+    \`\`\`
+
+</details>
+`;

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

@@ -2,32 +2,17 @@ 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';
+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 });
   }
 
-  get completer(): BaseChatModel {
-    return this._completer;
-  }
-
-  /**
-   * 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
@@ -35,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);
@@ -62,6 +50,5 @@ export class OpenAICompleter implements IBaseCompleter {
     }
   }
 
-  private _completer: ChatOpenAI;
-  private _prompt: string = COMPLETION_SYSTEM_PROMPT;
+  protected _completer: ChatOpenAI;
 }

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

@@ -0,0 +1,151 @@
+import {
+  CompletionHandler,
+  IInlineCompletionContext
+} from '@jupyterlab/completer';
+import { HumanMessage, SystemMessage } from '@langchain/core/messages';
+import { ChatWebLLM } from '@langchain/community/chat_models/webllm';
+
+import { BaseCompleter } from '../../base-completer';
+
+/**
+ * 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 WebLLMCompleter extends BaseCompleter {
+  constructor(options: BaseCompleter.IOptions) {
+    super(options);
+    const model = options.settings.model as string;
+    // provide model separately since ChatWebLLM expects it
+    this._completer = new ChatWebLLM({
+      ...options.settings,
+      model
+    });
+
+    // Initialize the model and track its status
+    this._isInitialized = false;
+    this._isInitializing = false;
+    this._initError = null;
+    void this._initializeModel();
+  }
+
+  /**
+   * Initialize the WebLLM model
+   */
+  private async _initializeModel(): Promise<void> {
+    if (this._isInitialized || this._isInitializing) {
+      return;
+    }
+
+    this._isInitializing = true;
+    try {
+      await this._completer.initialize((progress: any) => {
+        console.log('WebLLM initialization progress:', progress);
+      });
+      this._isInitialized = true;
+      this._isInitializing = false;
+      console.log('WebLLM model successfully initialized');
+    } catch (error) {
+      this._initError =
+        error instanceof Error ? error : new Error(String(error));
+      this._isInitializing = false;
+      console.error('Failed to initialize WebLLM model:', error);
+    }
+  }
+
+  get provider(): ChatWebLLM {
+    return this._completer;
+  }
+
+  async fetch(
+    request: CompletionHandler.IRequest,
+    context: IInlineCompletionContext
+  ) {
+    // Abort any pending request
+    if (this._abortController) {
+      this._abortController.abort();
+    }
+
+    // Create a new abort controller for this request
+    this._abortController = new AbortController();
+    const signal = this._abortController.signal;
+
+    if (!this._isInitialized) {
+      if (this._initError) {
+        console.error('WebLLM model failed to initialize:', this._initError);
+        return { items: [] };
+      }
+
+      if (!this._isInitializing) {
+        // Try to initialize again if it's not currently initializing
+        await this._initializeModel();
+      } else {
+        console.log(
+          'WebLLM model is still initializing, please try again later'
+        );
+        return { items: [] };
+      }
+
+      // Return empty if still not initialized
+      if (!this._isInitialized) {
+        return { items: [] };
+      }
+    }
+
+    const { text, offset: cursorOffset } = request;
+    const prompt = text.slice(0, cursorOffset);
+    const trimmedPrompt = prompt.trim();
+
+    const messages = [
+      new SystemMessage(this.systemPrompt),
+      new HumanMessage(trimmedPrompt)
+    ];
+
+    try {
+      console.log('Trigger invoke');
+      const response = await this._completer.invoke(messages, { signal });
+      let content = response.content as string;
+      console.log('Response content:', content);
+
+      if (CODE_BLOCK_START_REGEX.test(content)) {
+        content = content
+          .replace(CODE_BLOCK_START_REGEX, '')
+          .replace(CODE_BLOCK_END_REGEX, '');
+      }
+
+      const items = [{ insertText: content }];
+      return {
+        items
+      };
+    } catch (error) {
+      if (error instanceof Error) {
+        console.error('Error fetching completion from WebLLM:', error.message);
+      } else {
+        console.error('Unknown error fetching completion from WebLLM:', error);
+      }
+      return { items: [] };
+    }
+  }
+
+  protected _completer: ChatWebLLM;
+  private _isInitialized: boolean = false;
+  private _isInitializing: boolean = false;
+  private _initError: Error | null = null;
+  private _abortController: AbortController | null = null;
+}