@jupyterlite/ai 0.2.0

package/LICENSE

@@ -0,0 +1,30 @@
+BSD 3-Clause License
+
+Copyright (c) 2024, JupyterLite Contributors
+Copyright (c) 2024, Jeremy Tuloup
+All rights reserved.
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are met:
+
+1. Redistributions of source code must retain the above copyright notice, this
+   list of conditions and the following disclaimer.
+
+2. Redistributions in binary form must reproduce the above copyright notice,
+   this list of conditions and the following disclaimer in the documentation
+   and/or other materials provided with the distribution.
+
+3. Neither the name of the copyright holder nor the names of its
+   contributors may be used to endorse or promote products derived from
+   this software without specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

package/README.md

@@ -0,0 +1,112 @@
+# jupyterlite-ai
+
+[![Github Actions Status](https://github.com/jupyterlite/ai/workflows/Build/badge.svg)](https://github.com/jupyterlite/ai/actions/workflows/build.yml)
+[![lite-badge](https://jupyterlite.rtfd.io/en/latest/_static/badge.svg)](https://jupyterlite.github.io/ai/lab/index.html)
+
+AI code completions and chat for JupyterLab, Notebook 7 and JupyterLite, powered by MistralAI ✨
+
+[a screencast showing the Codestral extension in JupyterLite](https://github.com/jupyterlite/ai/assets/591645/855c4e3e-3a63-4868-8052-5c9909922c21)
+
+## Requirements
+
+> [!NOTE]
+> This extension is meant to be used in JupyterLite to enable AI code completions and chat in the browser, with a specific provider.
+> To enable more AI providers in JupyterLab and Jupyter Notebook, we recommend using the [Jupyter AI](https://github.com/jupyterlab/jupyter-ai) extension directly.
+> At the moment Jupyter AI is not compatible with JupyterLite, but might be to some extent in the future.
+
+- JupyterLab >= 4.1.0 or Notebook >= 7.1.0
+
+> [!WARNING]
+> This extension is still very much experimental. It is not an official MistralAI extension.
+> It is exploring the integration of the MistralAI API with JupyterLab, which can also be used in [JupyterLite](https://jupyterlite.readthedocs.io/).
+> For a more complete AI extension for JupyterLab, see [Jupyter AI](https://github.com/jupyterlab/jupyter-ai).
+
+## ✨ Try it in your browser ✨
+
+You can try the extension in your browser using JupyterLite:
+
+[![lite-badge](https://jupyterlite.rtfd.io/en/latest/_static/badge.svg)](https://jupyterlite.github.io/ai/lab/index.html)
+
+See the [Usage](#usage) section below for more information on how to provide your API key.
+
+## Install
+
+To install the extension, execute:
+
+```bash
+pip install jupyterlite-ai
+```
+
+# Usage
+
+1. Go to https://console.mistral.ai/api-keys/ and create an API key.
+
+![Screenshot showing how to create an API key](./img/1-api-key.png)
+
+2. Open the JupyterLab settings and go to the Codestral section to enter the API key
+
+![Screenshot showing how to add the API key to the settings](./img/2-jupyterlab-settings.png)
+
+3. Open the chat, or use the inline completer
+
+![Screenshot showing how to use the chat](./img/3-usage.png)
+
+## Uninstall
+
+To remove the extension, execute:
+
+```bash
+pip uninstall jupyterlite-ai
+```
+
+## Contributing
+
+### Development install
+
+Note: You will need NodeJS to build the extension package.
+
+The `jlpm` command is JupyterLab's pinned version of
+[yarn](https://yarnpkg.com/) that is installed with JupyterLab. You may use
+`yarn` or `npm` in lieu of `jlpm` below.
+
+```bash
+# Clone the repo to your local environment
+# Change directory to the jupyterlite_ai directory
+# Install package in development mode
+pip install -e "."
+# Link your development version of the extension with JupyterLab
+jupyter labextension develop . --overwrite
+# Rebuild extension Typescript source after making changes
+jlpm build
+```
+
+You can watch the source directory and run JupyterLab at the same time in different terminals to watch for changes in the extension's source and automatically rebuild the extension.
+
+```bash
+# Watch the source directory in one terminal, automatically rebuilding when needed
+jlpm watch
+# Run JupyterLab in another terminal
+jupyter lab
+```
+
+With the watch command running, every saved change will immediately be built locally and available in your running JupyterLab. Refresh JupyterLab to load the change in your browser (you may need to wait several seconds for the extension to be rebuilt).
+
+By default, the `jlpm build` command generates the source maps for this extension to make it easier to debug using the browser dev tools. To also generate source maps for the JupyterLab core extensions, you can run the following command:
+
+```bash
+jupyter lab build --minimize=False
+```
+
+### Development uninstall
+
+```bash
+pip uninstall jupyterlite-ai
+```
+
+In development mode, you will also need to remove the symlink created by `jupyter labextension develop`
+command. To find its location, you can run `jupyter labextension list` to figure out where the `labextensions`
+folder is located. Then you can remove the symlink named `@jupyterlite/ai` within that folder.
+
+### Packaging the extension
+
+See [RELEASE](RELEASE.md)

package/lib/chat-handler.d.ts

@@ -0,0 +1,24 @@
+import { ChatModel, IChatHistory, IChatMessage, INewMessage } from '@jupyter/chat';
+import type { BaseChatModel } from '@langchain/core/language_models/chat_models';
+import { IAIProvider } from './token';
+export type ConnectionMessage = {
+    type: 'connection';
+    client_id: string;
+};
+export declare class ChatHandler extends ChatModel {
+    constructor(options: ChatHandler.IOptions);
+    get provider(): BaseChatModel | null;
+    sendMessage(message: INewMessage): Promise<boolean>;
+    getHistory(): Promise<IChatHistory>;
+    dispose(): void;
+    messageAdded(message: IChatMessage): void;
+    private _aiProvider;
+    private _errorMessage;
+    private _history;
+    private _defaultErrorMessage;
+}
+export declare namespace ChatHandler {
+    interface IOptions extends ChatModel.IOptions {
+        aiProvider: IAIProvider;
+    }
+}

package/lib/chat-handler.js

@@ -0,0 +1,92 @@
+/*
+ * Copyright (c) Jupyter Development Team.
+ * Distributed under the terms of the Modified BSD License.
+ */
+import { ChatModel } from '@jupyter/chat';
+import { AIMessage, HumanMessage, mergeMessageRuns } from '@langchain/core/messages';
+import { UUID } from '@lumino/coreutils';
+import { getErrorMessage } from './llm-models';
+export class ChatHandler extends ChatModel {
+    constructor(options) {
+        super(options);
+        this._errorMessage = '';
+        this._history = { messages: [] };
+        this._defaultErrorMessage = 'AI provider not configured';
+        this._aiProvider = options.aiProvider;
+        this._aiProvider.modelChange.connect(() => {
+            this._errorMessage = this._aiProvider.chatError;
+        });
+    }
+    get provider() {
+        return this._aiProvider.chatModel;
+    }
+    async sendMessage(message) {
+        message.id = UUID.uuid4();
+        const msg = {
+            id: message.id,
+            body: message.body,
+            sender: { username: 'User' },
+            time: Date.now(),
+            type: 'msg'
+        };
+        this.messageAdded(msg);
+        if (this._aiProvider.chatModel === null) {
+            const errorMsg = {
+                id: UUID.uuid4(),
+                body: `**${this._errorMessage ? this._errorMessage : this._defaultErrorMessage}**`,
+                sender: { username: 'ERROR' },
+                time: Date.now(),
+                type: 'msg'
+            };
+            this.messageAdded(errorMsg);
+            return false;
+        }
+        this._history.messages.push(msg);
+        const messages = mergeMessageRuns(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);
+            this._history.messages.push(botMsg);
+            return true;
+        })
+            .catch(reason => {
+            const error = getErrorMessage(this._aiProvider.name, reason);
+            const errorMsg = {
+                id: UUID.uuid4(),
+                body: `**${error}**`,
+                sender: { username: 'ERROR' },
+                time: Date.now(),
+                type: 'msg'
+            };
+            this.messageAdded(errorMsg);
+            return false;
+        })
+            .finally(() => {
+            this.updateWriters([]);
+        });
+    }
+    async getHistory() {
+        return this._history;
+    }
+    dispose() {
+        super.dispose();
+    }
+    messageAdded(message) {
+        super.messageAdded(message);
+    }
+}

package/lib/completion-provider.d.ts

@@ -0,0 +1,40 @@
+import { CompletionHandler, IInlineCompletionContext, IInlineCompletionProvider } from '@jupyterlab/completer';
+import { LLM } from '@langchain/core/language_models/llms';
+import { IBaseCompleter, BaseCompleter } from './llm-models';
+import { ReadonlyPartialJSONObject } from '@lumino/coreutils';
+/**
+ * The generic completion provider to register to the completion provider manager.
+ */
+export declare class CompletionProvider implements IInlineCompletionProvider {
+    readonly identifier = "@jupyterlite/ai";
+    constructor(options: CompletionProvider.IOptions);
+    /**
+     * Set the completer.
+     *
+     * @param name - the name of the completer.
+     * @param settings - The settings associated to the completer.
+     */
+    setCompleter(name: string, settings: ReadonlyPartialJSONObject): void;
+    /**
+     * Get the current completer name.
+     */
+    get name(): string;
+    /**
+     * Get the current completer.
+     */
+    get completer(): IBaseCompleter | null;
+    /**
+     * Get the LLM completer.
+     */
+    get llmCompleter(): LLM | null;
+    fetch(request: CompletionHandler.IRequest, context: IInlineCompletionContext): Promise<any>;
+    private _name;
+    private _requestCompletion;
+    private _completer;
+}
+export declare namespace CompletionProvider {
+    interface IOptions extends BaseCompleter.IOptions {
+        name: string;
+        requestCompletion: () => void;
+    }
+}

package/lib/completion-provider.js

@@ -0,0 +1,57 @@
+import { getCompleter } from './llm-models';
+/**
+ * The generic completion provider to register to the completion provider manager.
+ */
+export class CompletionProvider {
+    constructor(options) {
+        this.identifier = '@jupyterlite/ai';
+        this._name = 'None';
+        this._completer = null;
+        const { name, settings } = options;
+        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, settings) {
+        try {
+            this._completer = getCompleter(name, settings);
+            if (this._completer) {
+                this._completer.requestCompletion = this._requestCompletion;
+            }
+            this._name = this._completer === null ? 'None' : name;
+        }
+        catch (e) {
+            this._completer = null;
+            this._name = 'None';
+            throw e;
+        }
+    }
+    /**
+     * Get the current completer name.
+     */
+    get name() {
+        return this._name;
+    }
+    /**
+     * Get the current completer.
+     */
+    get completer() {
+        return this._completer;
+    }
+    /**
+     * Get the LLM completer.
+     */
+    get llmCompleter() {
+        var _a;
+        return ((_a = this._completer) === null || _a === void 0 ? void 0 : _a.provider) || null;
+    }
+    async fetch(request, context) {
+        var _a;
+        return (_a = this._completer) === null || _a === void 0 ? void 0 : _a.fetch(request, context);
+    }
+}

package/lib/index.d.ts

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

package/lib/index.js

@@ -0,0 +1,90 @@
+import { ActiveCellManager, buildChatSidebar, buildErrorWidget } 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 { AIProvider } from './provider';
+import { IAIProvider } from './token';
+const chatPlugin = {
+    id: '@jupyterlite/ai:chat',
+    description: 'LLM chat extension',
+    autoStart: true,
+    optional: [INotebookTracker, ISettingRegistry, IThemeManager],
+    requires: [IAIProvider, IRenderMimeRegistry],
+    activate: async (app, aiProvider, rmRegistry, notebookTracker, settingsRegistry, themeManager) => {
+        let activeCellManager = null;
+        if (notebookTracker) {
+            activeCellManager = new ActiveCellManager({
+                tracker: notebookTracker,
+                shell: app.shell
+            });
+        }
+        const chatHandler = new ChatHandler({
+            aiProvider: aiProvider,
+            activeCellManager: activeCellManager
+        });
+        let sendWithShiftEnter = false;
+        let enableCodeToolbar = true;
+        function loadSetting(setting) {
+            sendWithShiftEnter = setting.get('sendWithShiftEnter')
+                .composite;
+            enableCodeToolbar = setting.get('enableCodeToolbar').composite;
+            chatHandler.config = { sendWithShiftEnter, enableCodeToolbar };
+        }
+        Promise.all([app.restored, settingsRegistry === null || settingsRegistry === void 0 ? void 0 : settingsRegistry.load(chatPlugin.id)])
+            .then(([, settings]) => {
+            if (!settings) {
+                console.warn('The SettingsRegistry is not loaded for the chat extension');
+                return;
+            }
+            loadSetting(settings);
+            settings.changed.connect(loadSetting);
+        })
+            .catch(reason => {
+            console.error(`Something went wrong when reading the settings.\n${reason}`);
+        });
+        let chatWidget = null;
+        try {
+            chatWidget = buildChatSidebar({
+                model: chatHandler,
+                themeManager,
+                rmRegistry
+            });
+            chatWidget.title.caption = 'Codestral Chat';
+        }
+        catch (e) {
+            chatWidget = buildErrorWidget(themeManager);
+        }
+        app.shell.add(chatWidget, 'left', { rank: 2000 });
+        console.log('Chat extension initialized');
+    }
+};
+const aiProviderPlugin = {
+    id: '@jupyterlite/ai:ai-provider',
+    autoStart: true,
+    requires: [ICompletionProviderManager, ISettingRegistry],
+    provides: IAIProvider,
+    activate: (app, manager, settingRegistry) => {
+        const aiProvider = new AIProvider({
+            completionProviderManager: manager,
+            requestCompletion: () => app.commands.execute('inline-completer:invoke')
+        });
+        settingRegistry
+            .load(aiProviderPlugin.id)
+            .then(settings => {
+            const updateProvider = () => {
+                const provider = settings.get('provider').composite;
+                aiProvider.setModels(provider, settings.composite);
+            };
+            settings.changed.connect(() => updateProvider());
+            updateProvider();
+        })
+            .catch(reason => {
+            console.error(`Failed to load settings for ${aiProviderPlugin.id}`, reason);
+        });
+        return aiProvider;
+    }
+};
+export default [chatPlugin, aiProviderPlugin];

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

@@ -0,0 +1,28 @@
+import { CompletionHandler, IInlineCompletionContext } from '@jupyterlab/completer';
+import { LLM } from '@langchain/core/language_models/llms';
+import { ReadonlyPartialJSONObject } from '@lumino/coreutils';
+export interface IBaseCompleter {
+    /**
+     * The LLM completer.
+     */
+    provider: LLM;
+    /**
+     * The function to fetch a new completion.
+     */
+    requestCompletion?: () => void;
+    /**
+     * The fetch request for the LLM completer.
+     */
+    fetch(request: CompletionHandler.IRequest, context: IInlineCompletionContext): Promise<any>;
+}
+/**
+ * The namespace for the base completer.
+ */
+export declare namespace BaseCompleter {
+    /**
+     * The options for the constructor of a completer.
+     */
+    interface IOptions {
+        settings: ReadonlyPartialJSONObject;
+    }
+}

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

@@ -0,0 +1 @@
+export {};

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

@@ -0,0 +1,15 @@
+import { CompletionHandler, IInlineCompletionContext } from '@jupyterlab/completer';
+import { LLM } from '@langchain/core/language_models/llms';
+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?;
+    private _throttler;
+    private _mistralProvider;
+    private _currentData;
+}

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

@@ -0,0 +1,79 @@
+import { MistralAI } from '@langchain/mistralai';
+import { Throttler } from '@lumino/polling';
+/**
+ * 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
+                };
+            }
+            // Extract results of completion request.
+            const items = response.choices.map((choice) => {
+                return { insertText: choice.message.content };
+            });
+            return {
+                items
+            };
+        }, { limit: INTERVAL });
+    }
+    get provider() {
+        return this._mistralProvider;
+    }
+    set requestCompletion(value) {
+        this._requestCompletion = 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: []
+        };
+        try {
+            this._currentData = data;
+            const completionResult = await this._throttler.invoke(data);
+            if (completionResult.fetchAgain) {
+                if (this._requestCompletion) {
+                    this._requestCompletion();
+                }
+            }
+            return { items: completionResult.items };
+        }
+        catch (error) {
+            console.error('Error fetching completions', error);
+            return { items: [] };
+        }
+    }
+}

package/lib/llm-models/index.d.ts

@@ -0,0 +1,3 @@
+export * from './base-completer';
+export * from './codestral-completer';
+export * from './utils';

package/lib/llm-models/index.js

@@ -0,0 +1,3 @@
+export * from './base-completer';
+export * from './codestral-completer';
+export * from './utils';

package/lib/llm-models/utils.d.ts

@@ -0,0 +1,15 @@
+import { BaseChatModel } from '@langchain/core/language_models/chat_models';
+import { IBaseCompleter } from './base-completer';
+import { ReadonlyPartialJSONObject } from '@lumino/coreutils';
+/**
+ * Get an LLM completer from the name.
+ */
+export declare function getCompleter(name: string, settings: ReadonlyPartialJSONObject): IBaseCompleter | null;
+/**
+ * Get an LLM chat model from the name.
+ */
+export declare function getChatModel(name: string, settings: ReadonlyPartialJSONObject): BaseChatModel | null;
+/**
+ * Get the error message from provider.
+ */
+export declare function getErrorMessage(name: string, error: any): string;

package/lib/llm-models/utils.js

@@ -0,0 +1,29 @@
+import { ChatMistralAI } from '@langchain/mistralai';
+import { CodestralCompleter } from './codestral-completer';
+/**
+ * Get an LLM completer from the name.
+ */
+export function getCompleter(name, settings) {
+    if (name === 'MistralAI') {
+        return new CodestralCompleter({ settings });
+    }
+    return null;
+}
+/**
+ * Get an LLM chat model from the name.
+ */
+export function getChatModel(name, settings) {
+    if (name === 'MistralAI') {
+        return new ChatMistralAI({ ...settings });
+    }
+    return null;
+}
+/**
+ * Get the error message from provider.
+ */
+export function getErrorMessage(name, error) {
+    if (name === 'MistralAI') {
+        return error.message;
+    }
+    return 'Unknown provider';
+}