@jupyterlite/ai 0.2.0 → 0.3.0

package/lib/chat-handler.d.ts

@@ -8,11 +8,23 @@ export type ConnectionMessage = {
 export declare class ChatHandler extends ChatModel {
     constructor(options: ChatHandler.IOptions);
     get provider(): BaseChatModel | null;
+    /**
+     * Getter and setter for the persona name.
+     */
+    get personaName(): string;
+    set personaName(value: string);
+    /**
+     * Getter and setter for the initial prompt.
+     */
+    get prompt(): string;
+    set prompt(value: string);
     sendMessage(message: INewMessage): Promise<boolean>;
     getHistory(): Promise<IChatHistory>;
     dispose(): void;
     messageAdded(message: IChatMessage): void;
     private _aiProvider;
+    private _personaName;
+    private _prompt;
     private _errorMessage;
     private _history;
     private _defaultErrorMessage;

package/lib/chat-handler.js

@@ -3,28 +3,72 @@
  * Distributed under the terms of the Modified BSD License.
  */
 import { ChatModel } from '@jupyter/chat';
-import { AIMessage, HumanMessage, mergeMessageRuns } from '@langchain/core/messages';
+import { AIMessage, HumanMessage, mergeMessageRuns, SystemMessage } from '@langchain/core/messages';
 import { UUID } from '@lumino/coreutils';
 import { getErrorMessage } from './llm-models';
+import { chatSystemPrompt } from './provider';
+import { jupyternautLiteIcon } from './icons';
+/**
+ * The base64 encoded SVG string of the jupyternaut lite icon.
+ * Encode so it can be passed as avatar_url to jupyter-chat.
+ */
+const AI_AVATAR_BASE64 = btoa(jupyternautLiteIcon.svgstr);
+const AI_AVATAR = `data:image/svg+xml;base64,${AI_AVATAR_BASE64}`;
 export class ChatHandler extends ChatModel {
     constructor(options) {
         super(options);
+        this._personaName = 'AI';
         this._errorMessage = '';
         this._history = { messages: [] };
         this._defaultErrorMessage = 'AI provider not configured';
         this._aiProvider = options.aiProvider;
+        this._prompt = chatSystemPrompt({ provider_name: this._aiProvider.name });
         this._aiProvider.modelChange.connect(() => {
             this._errorMessage = this._aiProvider.chatError;
+            this._prompt = chatSystemPrompt({ provider_name: this._aiProvider.name });
         });
     }
     get provider() {
         return this._aiProvider.chatModel;
     }
+    /**
+     * Getter and setter for the persona name.
+     */
+    get personaName() {
+        return this._personaName;
+    }
+    set personaName(value) {
+        this.messages.forEach(message => {
+            if (message.sender.username === this._personaName) {
+                const updated = { ...message };
+                updated.sender.username = value;
+                this.messageAdded(updated);
+            }
+        });
+        this._personaName = value;
+    }
+    /**
+     * Getter and setter for the initial prompt.
+     */
+    get prompt() {
+        return this._prompt;
+    }
+    set prompt(value) {
+        this._prompt = value;
+    }
     async sendMessage(message) {
+        var _a;
+        const body = message.body;
+        if (body.startsWith('/clear')) {
+            // TODO: do we need a clear method?
+            this.messagesDeleted(0, this.messages.length);
+            this._history.messages = [];
+            return false;
+        }
         message.id = UUID.uuid4();
         const msg = {
             id: message.id,
-            body: message.body,
+            body,
             sender: { username: 'User' },
             time: Date.now(),
             type: 'msg'
@@ -42,29 +86,34 @@ export class ChatHandler extends ChatModel {
             return false;
         }
         this._history.messages.push(msg);
-        const messages = mergeMessageRuns(this._history.messages.map(msg => {
+        const messages = mergeMessageRuns([new SystemMessage(this._prompt)]);
+        messages.push(...this._history.messages.map(msg => {
             if (msg.sender.username === 'User') {
                 return new HumanMessage(msg.body);
             }
             return new AIMessage(msg.body);
         }));
-        this.updateWriters([{ username: 'AI' }]);
-        return this._aiProvider.chatModel
-            .invoke(messages)
-            .then(response => {
-            const content = response.content;
-            const botMsg = {
-                id: UUID.uuid4(),
-                body: content.toString(),
-                sender: { username: 'AI' },
-                time: Date.now(),
-                type: 'msg'
-            };
-            this.messageAdded(botMsg);
+        const sender = { username: this._personaName, avatar_url: AI_AVATAR };
+        this.updateWriters([sender]);
+        // create an empty message to be filled by the AI provider
+        const botMsg = {
+            id: UUID.uuid4(),
+            body: '',
+            sender,
+            time: Date.now(),
+            type: 'msg'
+        };
+        let content = '';
+        try {
+            for await (const chunk of await this._aiProvider.chatModel.stream(messages)) {
+                content += (_a = chunk.content) !== null && _a !== void 0 ? _a : chunk;
+                botMsg.body = content;
+                this.messageAdded(botMsg);
+            }
             this._history.messages.push(botMsg);
             return true;
-        })
-            .catch(reason => {
+        }
+        catch (reason) {
             const error = getErrorMessage(this._aiProvider.name, reason);
             const errorMsg = {
                 id: UUID.uuid4(),
@@ -75,10 +124,10 @@ export class ChatHandler extends ChatModel {
             };
             this.messageAdded(errorMsg);
             return false;
-        })
-            .finally(() => {
+        }
+        finally {
             this.updateWriters([]);
-        });
+        }
     }
     async getHistory() {
         return this._history;

package/lib/completion-provider.d.ts

@@ -1,7 +1,7 @@
 import { CompletionHandler, IInlineCompletionContext, IInlineCompletionProvider } from '@jupyterlab/completer';
-import { LLM } from '@langchain/core/language_models/llms';
-import { IBaseCompleter, BaseCompleter } from './llm-models';
+import { BaseLanguageModel } from '@langchain/core/language_models/base';
 import { ReadonlyPartialJSONObject } from '@lumino/coreutils';
+import { IBaseCompleter, BaseCompleter } from './llm-models';
 /**
  * The generic completion provider to register to the completion provider manager.
  */
@@ -26,7 +26,7 @@ export declare class CompletionProvider implements IInlineCompletionProvider {
     /**
      * Get the LLM completer.
      */
-    get llmCompleter(): LLM | null;
+    get llmCompleter(): BaseLanguageModel | null;
     fetch(request: CompletionHandler.IRequest, context: IInlineCompletionContext): Promise<any>;
     private _name;
     private _requestCompletion;

package/lib/icons.d.ts

@@ -0,0 +1,2 @@
+import { LabIcon } from '@jupyterlab/ui-components';
+export declare const jupyternautLiteIcon: LabIcon;

package/lib/icons.js

@@ -0,0 +1,15 @@
+/*
+ * 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/lib/index.d.ts

@@ -1,4 +1,5 @@
+import { IAutocompletionRegistry } from '@jupyter/chat';
 import { JupyterFrontEndPlugin } from '@jupyterlab/application';
 import { IAIProvider } from './token';
-declare const _default: (JupyterFrontEndPlugin<void> | JupyterFrontEndPlugin<IAIProvider>)[];
+declare const _default: (JupyterFrontEndPlugin<void> | JupyterFrontEndPlugin<IAutocompletionRegistry> | JupyterFrontEndPlugin<IAIProvider>)[];
 export default _default;

package/lib/index.js

@@ -1,19 +1,46 @@
-import { ActiveCellManager, buildChatSidebar, buildErrorWidget } from '@jupyter/chat';
+import { ActiveCellManager, AutocompletionRegistry, buildChatSidebar, buildErrorWidget, IAutocompletionRegistry } from '@jupyter/chat';
 import { IThemeManager } from '@jupyterlab/apputils';
 import { ICompletionProviderManager } from '@jupyterlab/completer';
 import { INotebookTracker } from '@jupyterlab/notebook';
 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 = {
+    id: '@jupyterlite/ai:autocompletion-registry',
+    description: 'Autocompletion registry',
+    autoStart: true,
+    provides: IAutocompletionRegistry,
+    activate: () => {
+        const autocompletionRegistry = new AutocompletionRegistry();
+        const options = ['/clear'];
+        const autocompletionCommands = {
+            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 = {
     id: '@jupyterlite/ai:chat',
     description: 'LLM chat extension',
     autoStart: true,
+    requires: [IAIProvider, IRenderMimeRegistry, IAutocompletionRegistry],
     optional: [INotebookTracker, ISettingRegistry, IThemeManager],
-    requires: [IAIProvider, IRenderMimeRegistry],
-    activate: async (app, aiProvider, rmRegistry, notebookTracker, settingsRegistry, themeManager) => {
+    activate: async (app, aiProvider, rmRegistry, autocompletionRegistry, notebookTracker, settingsRegistry, themeManager) => {
         let activeCellManager = null;
         if (notebookTracker) {
             activeCellManager = new ActiveCellManager({
@@ -27,11 +54,15 @@ const chatPlugin = {
         });
         let sendWithShiftEnter = false;
         let enableCodeToolbar = true;
+        let personaName = 'AI';
         function loadSetting(setting) {
             sendWithShiftEnter = setting.get('sendWithShiftEnter')
                 .composite;
             enableCodeToolbar = setting.get('enableCodeToolbar').composite;
+            personaName = setting.get('personaName').composite;
+            // set the properties
             chatHandler.config = { sendWithShiftEnter, enableCodeToolbar };
+            chatHandler.personaName = personaName;
         }
         Promise.all([app.restored, settingsRegistry === null || settingsRegistry === void 0 ? void 0 : settingsRegistry.load(chatPlugin.id)])
             .then(([, settings]) => {
@@ -50,9 +81,10 @@ const chatPlugin = {
             chatWidget = buildChatSidebar({
                 model: chatHandler,
                 themeManager,
-                rmRegistry
+                rmRegistry,
+                autocompletionRegistry
             });
-            chatWidget.title.caption = 'Codestral Chat';
+            chatWidget.title.caption = 'Jupyterlite AI Chat';
         }
         catch (e) {
             chatWidget = buildErrorWidget(themeManager);
@@ -71,11 +103,34 @@ const aiProviderPlugin = {
             completionProviderManager: manager,
             requestCompletion: () => app.commands.execute('inline-completer:invoke')
         });
+        let currentProvider = 'None';
         settingRegistry
             .load(aiProviderPlugin.id)
             .then(settings => {
             const updateProvider = () => {
                 const provider = settings.get('provider').composite;
+                if (provider !== currentProvider) {
+                    // Update the settings panel.
+                    currentProvider = provider;
+                    const settingsProperties = settings.schema.properties;
+                    if (settingsProperties) {
+                        const schemaKeys = Object.keys(settingsProperties);
+                        schemaKeys.forEach(key => {
+                            var _a;
+                            if (key !== 'provider') {
+                                (_a = settings.schema.properties) === null || _a === void 0 ? true : delete _a[key];
+                            }
+                        });
+                        const properties = getSettings(provider);
+                        if (properties === null) {
+                            return;
+                        }
+                        Object.entries(properties).forEach(([name, value], index) => {
+                            settingsProperties[name] = value;
+                        });
+                    }
+                }
+                // Update the settings to the AI providers.
                 aiProvider.setModels(provider, settings.composite);
             };
             settings.changed.connect(() => updateProvider());
@@ -87,4 +142,4 @@ const aiProviderPlugin = {
         return aiProvider;
     }
 };
-export default [chatPlugin, aiProviderPlugin];
+export default [chatPlugin, autocompletionRegistryPlugin, aiProviderPlugin];

package/lib/llm-models/anthropic-completer.d.ts

@@ -0,0 +1,19 @@
+import { CompletionHandler, IInlineCompletionContext } from '@jupyterlab/completer';
+import { BaseChatModel } from '@langchain/core/language_models/chat_models';
+import { BaseCompleter, IBaseCompleter } from './base-completer';
+export declare class AnthropicCompleter implements IBaseCompleter {
+    constructor(options: BaseCompleter.IOptions);
+    get provider(): BaseChatModel;
+    /**
+     * Getter and setter for the initial prompt.
+     */
+    get prompt(): string;
+    set prompt(value: string);
+    fetch(request: CompletionHandler.IRequest, context: IInlineCompletionContext): Promise<{
+        items: {
+            insertText: string;
+        }[];
+    }>;
+    private _anthropicProvider;
+    private _prompt;
+}

package/lib/llm-models/anthropic-completer.js

@@ -0,0 +1,57 @@
+import { ChatAnthropic } from '@langchain/anthropic';
+import { AIMessage, SystemMessage } from '@langchain/core/messages';
+import { COMPLETION_SYSTEM_PROMPT } from '../provider';
+export class AnthropicCompleter {
+    constructor(options) {
+        this._prompt = COMPLETION_SYSTEM_PROMPT;
+        this._anthropicProvider = new ChatAnthropic({ ...options.settings });
+    }
+    get provider() {
+        return this._anthropicProvider;
+    }
+    /**
+     * Getter and setter for the initial prompt.
+     */
+    get prompt() {
+        return this._prompt;
+    }
+    set prompt(value) {
+        this._prompt = value;
+    }
+    async fetch(request, context) {
+        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: [] };
+        }
+    }
+}

package/lib/llm-models/base-completer.d.ts

@@ -1,11 +1,15 @@
 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/lib/llm-models/chrome-completer.d.ts

@@ -0,0 +1,19 @@
+import { CompletionHandler, IInlineCompletionContext } from '@jupyterlab/completer';
+import { LLM } from '@langchain/core/language_models/llms';
+import { BaseCompleter, IBaseCompleter } from './base-completer';
+export declare class ChromeCompleter implements IBaseCompleter {
+    constructor(options: BaseCompleter.IOptions);
+    /**
+     * Getter and setter for the initial prompt.
+     */
+    get prompt(): string;
+    set prompt(value: string);
+    get provider(): LLM;
+    fetch(request: CompletionHandler.IRequest, context: IInlineCompletionContext): Promise<{
+        items: {
+            insertText: string;
+        }[];
+    }>;
+    private _chromeProvider;
+    private _prompt;
+}

package/lib/llm-models/chrome-completer.js

@@ -0,0 +1,67 @@
+import { ChromeAI } from '@langchain/community/experimental/llms/chrome_ai';
+import { HumanMessage, SystemMessage } from '@langchain/core/messages';
+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 {
+    constructor(options) {
+        this._prompt = COMPLETION_SYSTEM_PROMPT;
+        this._chromeProvider = new ChromeAI({ ...options.settings });
+    }
+    /**
+     * Getter and setter for the initial prompt.
+     */
+    get prompt() {
+        return this._prompt;
+    }
+    set prompt(value) {
+        this._prompt = value;
+    }
+    get provider() {
+        return this._chromeProvider;
+    }
+    async fetch(request, context) {
+        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: [] };
+        }
+    }
+}

package/lib/llm-models/codestral-completer.d.ts

@@ -1,15 +1,16 @@
 import { CompletionHandler, IInlineCompletionContext } from '@jupyterlab/completer';
-import { LLM } from '@langchain/core/language_models/llms';
+import { BaseChatModel } from '@langchain/core/language_models/chat_models';
 import { BaseCompleter, IBaseCompleter } from './base-completer';
 export declare class CodestralCompleter implements IBaseCompleter {
     constructor(options: BaseCompleter.IOptions);
-    get provider(): LLM;
-    set requestCompletion(value: () => void);
-    fetch(request: CompletionHandler.IRequest, context: IInlineCompletionContext): Promise<{
-        items: any;
-    }>;
-    private _requestCompletion?;
+    get provider(): BaseChatModel;
+    /**
+     * Getter and setter for the initial prompt.
+     */
+    get prompt(): string;
+    set prompt(value: string);
+    fetch(request: CompletionHandler.IRequest, context: IInlineCompletionContext): Promise<any>;
     private _throttler;
     private _mistralProvider;
-    private _currentData;
+    private _prompt;
 }

package/lib/llm-models/codestral-completer.js

@@ -1,75 +1,58 @@
-import { MistralAI } from '@langchain/mistralai';
+import { HumanMessage, SystemMessage } from '@langchain/core/messages';
+import { ChatMistralAI } from '@langchain/mistralai';
 import { Throttler } from '@lumino/polling';
+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 {
     constructor(options) {
-        this._currentData = null;
-        // this._requestCompletion = options.requestCompletion;
-        this._mistralProvider = new MistralAI({ ...options.settings });
-        this._throttler = new Throttler(async (data) => {
-            var _a;
-            const invokedData = data;
-            // Request completion.
-            const request = this._mistralProvider.completionWithRetry(data, {}, false);
-            const timeoutPromise = new Promise(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 !== ((_a = this._currentData) === null || _a === void 0 ? void 0 : _a.prompt)) {
-                return {
-                    items: [],
-                    fetchAgain: true
-                };
-            }
+        this._prompt = COMPLETION_SYSTEM_PROMPT;
+        this._mistralProvider = new ChatMistralAI({ ...options.settings });
+        this._throttler = new Throttler(async (messages) => {
+            const response = await this._mistralProvider.invoke(messages);
             // Extract results of completion request.
-            const items = response.choices.map((choice) => {
-                return { insertText: choice.message.content };
-            });
-            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() {
         return this._mistralProvider;
     }
-    set requestCompletion(value) {
-        this._requestCompletion = value;
+    /**
+     * Getter and setter for the initial prompt.
+     */
+    get prompt() {
+        return this._prompt;
+    }
+    set prompt(value) {
+        this._prompt = value;
     }
     async fetch(request, context) {
         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 = [
+            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);