@jupyterlite/ai 0.2.0

package/src/index.ts

@@ -0,0 +1,130 @@
+import {
+  ActiveCellManager,
+  buildChatSidebar,
+  buildErrorWidget,
+  IActiveCellManager
+} from '@jupyter/chat';
+import {
+  JupyterFrontEnd,
+  JupyterFrontEndPlugin
+} from '@jupyterlab/application';
+import { ReactWidget, 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: JupyterFrontEndPlugin<void> = {
+  id: '@jupyterlite/ai:chat',
+  description: 'LLM chat extension',
+  autoStart: true,
+  optional: [INotebookTracker, ISettingRegistry, IThemeManager],
+  requires: [IAIProvider, IRenderMimeRegistry],
+  activate: async (
+    app: JupyterFrontEnd,
+    aiProvider: IAIProvider,
+    rmRegistry: IRenderMimeRegistry,
+    notebookTracker: INotebookTracker | null,
+    settingsRegistry: ISettingRegistry | null,
+    themeManager: IThemeManager | null
+  ) => {
+    let activeCellManager: IActiveCellManager | null = 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: ISettingRegistry.ISettings): void {
+      sendWithShiftEnter = setting.get('sendWithShiftEnter')
+        .composite as boolean;
+      enableCodeToolbar = setting.get('enableCodeToolbar').composite as boolean;
+      chatHandler.config = { sendWithShiftEnter, enableCodeToolbar };
+    }
+
+    Promise.all([app.restored, 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: ReactWidget | null = null;
+    try {
+      chatWidget = buildChatSidebar({
+        model: chatHandler,
+        themeManager,
+        rmRegistry
+      });
+      chatWidget.title.caption = 'Codestral Chat';
+    } catch (e) {
+      chatWidget = buildErrorWidget(themeManager);
+    }
+
+    app.shell.add(chatWidget as ReactWidget, 'left', { rank: 2000 });
+
+    console.log('Chat extension initialized');
+  }
+};
+
+const aiProviderPlugin: JupyterFrontEndPlugin<IAIProvider> = {
+  id: '@jupyterlite/ai:ai-provider',
+  autoStart: true,
+  requires: [ICompletionProviderManager, ISettingRegistry],
+  provides: IAIProvider,
+  activate: (
+    app: JupyterFrontEnd,
+    manager: ICompletionProviderManager,
+    settingRegistry: ISettingRegistry
+  ): IAIProvider => {
+    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 as string;
+          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/src/llm-models/base-completer.ts

@@ -0,0 +1,38 @@
+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 namespace BaseCompleter {
+  /**
+   * The options for the constructor of a completer.
+   */
+  export interface IOptions {
+    settings: ReadonlyPartialJSONObject;
+  }
+}

package/src/llm-models/codestral-completer.ts

@@ -0,0 +1,113 @@
+import {
+  CompletionHandler,
+  IInlineCompletionContext
+} from '@jupyterlab/completer';
+import { LLM } from '@langchain/core/language_models/llms';
+import { MistralAI } from '@langchain/mistralai';
+import { Throttler } from '@lumino/polling';
+import { CompletionRequest } from '@mistralai/mistralai';
+
+import { BaseCompleter, IBaseCompleter } from './base-completer';
+
+/**
+ * The Mistral API has a rate limit of 1 request per second
+ */
+const INTERVAL = 1000;
+
+/**
+ * Timeout to avoid endless requests
+ */
+const REQUEST_TIMEOUT = 3000;
+
+export class CodestralCompleter implements IBaseCompleter {
+  constructor(options: BaseCompleter.IOptions) {
+    // this._requestCompletion = options.requestCompletion;
+    this._mistralProvider = new MistralAI({ ...options.settings });
+    this._throttler = new Throttler(
+      async (data: CompletionRequest) => {
+        const invokedData = data;
+
+        // Request completion.
+        const request = this._mistralProvider.completionWithRetry(
+          data,
+          {},
+          false
+        );
+        const timeoutPromise = new Promise<null>(resolve => {
+          return setTimeout(() => resolve(null), REQUEST_TIMEOUT);
+        });
+
+        // Fetch again if the request is too long or if the prompt has changed.
+        const response = await Promise.race([request, timeoutPromise]);
+        if (
+          response === null ||
+          invokedData.prompt !== this._currentData?.prompt
+        ) {
+          return {
+            items: [],
+            fetchAgain: true
+          };
+        }
+
+        // Extract results of completion request.
+        const items = response.choices.map((choice: any) => {
+          return { insertText: choice.message.content as string };
+        });
+
+        return {
+          items
+        };
+      },
+      { limit: INTERVAL }
+    );
+  }
+
+  get provider(): LLM {
+    return this._mistralProvider;
+  }
+
+  set requestCompletion(value: () => void) {
+    this._requestCompletion = value;
+  }
+
+  async fetch(
+    request: CompletionHandler.IRequest,
+    context: IInlineCompletionContext
+  ) {
+    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: [] };
+    }
+  }
+
+  private _requestCompletion?: () => void;
+  private _throttler: Throttler;
+  private _mistralProvider: MistralAI;
+  private _currentData: CompletionRequest | null = null;
+}

package/src/llm-models/index.ts

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

package/src/llm-models/utils.ts

@@ -0,0 +1,41 @@
+import { BaseChatModel } from '@langchain/core/language_models/chat_models';
+import { ChatMistralAI } from '@langchain/mistralai';
+import { IBaseCompleter } from './base-completer';
+import { CodestralCompleter } from './codestral-completer';
+import { ReadonlyPartialJSONObject } from '@lumino/coreutils';
+
+/**
+ * Get an LLM completer from the name.
+ */
+export function getCompleter(
+  name: string,
+  settings: ReadonlyPartialJSONObject
+): IBaseCompleter | null {
+  if (name === 'MistralAI') {
+    return new CodestralCompleter({ settings });
+  }
+  return null;
+}
+
+/**
+ * Get an LLM chat model from the name.
+ */
+export function getChatModel(
+  name: string,
+  settings: ReadonlyPartialJSONObject
+): BaseChatModel | null {
+  if (name === 'MistralAI') {
+    return new ChatMistralAI({ ...settings });
+  }
+  return null;
+}
+
+/**
+ * Get the error message from provider.
+ */
+export function getErrorMessage(name: string, error: any): string {
+  if (name === 'MistralAI') {
+    return error.message;
+  }
+  return 'Unknown provider';
+}

package/src/provider.ts

@@ -0,0 +1,154 @@
+import { ICompletionProviderManager } from '@jupyterlab/completer';
+import { BaseLanguageModel } from '@langchain/core/language_models/base';
+import { BaseChatModel } from '@langchain/core/language_models/chat_models';
+import { ISignal, Signal } from '@lumino/signaling';
+import { ReadonlyPartialJSONObject } from '@lumino/coreutils';
+
+import { CompletionProvider } from './completion-provider';
+import { getChatModel, IBaseCompleter } from './llm-models';
+import { IAIProvider } from './token';
+
+export class AIProvider implements IAIProvider {
+  constructor(options: AIProvider.IOptions) {
+    this._completionProvider = new CompletionProvider({
+      name: 'None',
+      settings: {},
+      requestCompletion: options.requestCompletion
+    });
+    options.completionProviderManager.registerInlineProvider(
+      this._completionProvider
+    );
+  }
+
+  get name(): string {
+    return this._name;
+  }
+
+  /**
+   * Get the current completer of the completion provider.
+   */
+  get completer(): IBaseCompleter | null {
+    if (this._name === null) {
+      return null;
+    }
+    return this._completionProvider.completer;
+  }
+
+  /**
+   * Get the current llm chat model.
+   */
+  get chatModel(): BaseChatModel | null {
+    if (this._name === null) {
+      return null;
+    }
+    return this._llmChatModel;
+  }
+
+  /**
+   * Get the current chat error;
+   */
+  get chatError(): string {
+    return this._chatError;
+  }
+
+  /**
+   * get the current completer error.
+   */
+  get completerError(): string {
+    return this._completerError;
+  }
+
+  /**
+   * Set the models (chat model and completer).
+   * Creates the models if the name has changed, otherwise only updates their config.
+   *
+   * @param name - the name of the model to use.
+   * @param settings - the settings for the models.
+   */
+  setModels(name: string, settings: ReadonlyPartialJSONObject) {
+    try {
+      this._completionProvider.setCompleter(name, settings);
+      this._completerError = '';
+    } catch (e: any) {
+      this._completerError = e.message;
+    }
+    try {
+      this._llmChatModel = getChatModel(name, settings);
+      this._chatError = '';
+    } catch (e: any) {
+      this._chatError = e.message;
+      this._llmChatModel = null;
+    }
+    this._name = name;
+    this._modelChange.emit();
+  }
+
+  get modelChange(): ISignal<IAIProvider, void> {
+    return this._modelChange;
+  }
+
+  private _completionProvider: CompletionProvider;
+  private _llmChatModel: BaseChatModel | null = null;
+  private _name: string = 'None';
+  private _modelChange = new Signal<IAIProvider, void>(this);
+  private _chatError: string = '';
+  private _completerError: string = '';
+}
+
+export namespace AIProvider {
+  /**
+   * The options for the LLM provider.
+   */
+  export interface IOptions {
+    /**
+     * The completion provider manager in which register the LLM completer.
+     */
+    completionProviderManager: ICompletionProviderManager;
+    /**
+     * The application commands registry.
+     */
+    requestCompletion: () => void;
+  }
+
+  /**
+   * This function indicates whether a key is writable in an object.
+   * https://stackoverflow.com/questions/54724875/can-we-check-whether-property-is-readonly-in-typescript
+   *
+   * @param obj - An object extending the BaseLanguageModel interface.
+   * @param key - A string as a key of the object.
+   * @returns a boolean whether the key is writable or not.
+   */
+  export function isWritable<T extends BaseLanguageModel>(
+    obj: T,
+    key: keyof T
+  ) {
+    const desc =
+      Object.getOwnPropertyDescriptor(obj, key) ||
+      Object.getOwnPropertyDescriptor(Object.getPrototypeOf(obj), key) ||
+      {};
+    return Boolean(desc.writable);
+  }
+
+  /**
+   * Update the config of a language model.
+   * It only updates the writable attributes of the model.
+   *
+   * @param model - the model to update.
+   * @param settings - the configuration s a JSON object.
+   */
+  export function updateConfig<T extends BaseLanguageModel>(
+    model: T,
+    settings: ReadonlyPartialJSONObject
+  ) {
+    Object.entries(settings).forEach(([key, value], index) => {
+      if (key in model) {
+        const modelKey = key as keyof typeof model;
+        if (isWritable(model, modelKey)) {
+          // eslint-disable-next-line @typescript-eslint/ban-ts-comment
+          // @ts-ignore
+          model[modelKey] = value;
+        }
+      }
+    });
+  }
+}

package/src/token.ts

@@ -0,0 +1,19 @@
+import { BaseChatModel } from '@langchain/core/language_models/chat_models';
+import { Token } from '@lumino/coreutils';
+import { ISignal } from '@lumino/signaling';
+
+import { IBaseCompleter } from './llm-models';
+
+export interface IAIProvider {
+  name: string;
+  completer: IBaseCompleter | null;
+  chatModel: BaseChatModel | null;
+  modelChange: ISignal<IAIProvider, void>;
+  chatError: string;
+  completerError: string;
+}
+
+export const IAIProvider = new Token<IAIProvider>(
+  '@jupyterlite/ai:AIProvider',
+  'Provider for chat and completion LLM provider'
+);

package/style/base.css

@@ -0,0 +1,7 @@
+/*
+    See the JupyterLab Developer Guide for useful CSS Patterns:
+
+    https://jupyterlab.readthedocs.io/en/stable/developer/css.html
+*/
+
+@import url('@jupyter/chat/style/index.css');

package/style/index.css

@@ -0,0 +1 @@
+@import url('base.css');

package/style/index.js

@@ -0,0 +1 @@
+import './base.css';