@jupyterlite/ai 0.14.0 → 0.16.0

package/src/index.ts

@@ -23,6 +23,7 @@ import {
 import {
   ICommandPalette,
   IThemeManager,
+  showDialog,
   showErrorMessage,
   WidgetTracker
 } from '@jupyterlab/apputils';
@@ -31,6 +32,8 @@ import { ICompletionProviderManager } from '@jupyterlab/completer';
 
 import { IDocumentManager } from '@jupyterlab/docmanager';
 
+import { FileDialog } from '@jupyterlab/filebrowser';
+
 import { INotebookTracker } from '@jupyterlab/notebook';
 
 import { IRenderMimeRegistry } from '@jupyterlab/rendermime';
@@ -48,28 +51,36 @@ import {
 } from '@jupyterlab/translation';
 
 import {
+  fileUploadIcon,
+  saveIcon,
   settingsIcon,
   Toolbar,
   ToolbarButton
 } from '@jupyterlab/ui-components';
 
-import { ISecretsManager, SecretsManager } from 'jupyter-secrets-manager';
-
 import { PromiseDelegate, UUID } from '@lumino/coreutils';
+
 import { DisposableSet } from '@lumino/disposable';
+
 import { CommandRegistry } from '@lumino/commands';
 
+import { IComponentsRendererFactory } from 'jupyter-chat-components';
+
+import { ISecretsManager, SecretsManager } from 'jupyter-secrets-manager';
+
 import { AgentManagerFactory } from './agent';
 
 import { AIChatModel } from './chat-model';
+
 import { RenderedMessageOutputAreaCompat } from './rendered-message-outputarea';
 
 import { ClearCommandProvider } from './chat-commands/clear';
+
 import { SkillsCommandProvider } from './chat-commands/skills';
 
 import { ProviderRegistry } from './providers/provider-registry';
 
-import { ApprovalButtons } from './approval-buttons';
+import { SaveComponentWidget } from './components/save-button';
 
 import { ChatModelHandler } from './chat-model-handler';
 
@@ -77,13 +88,14 @@ import {
   CommandIds,
   IAgentManagerFactory,
   type IAISecretsAccess,
+  IAISettingsModel,
+  IChatModelHandler,
+  IDiffManager,
+  type IProviderConfig,
   IProviderRegistry,
   IToolRegistry,
   ISkillRegistry,
-  SECRETS_NAMESPACE,
-  IAISettingsModel,
-  IChatModelHandler,
-  IDiffManager
+  SECRETS_NAMESPACE
 } from './tokens';
 
 import {
@@ -102,10 +114,10 @@ import {
   createToolSelectItem,
   stopItem,
   CompletionStatusWidget,
-  TokenUsageWidget
+  UsageWidget
 } from './components';
 
-import { AISettingsModel, IProviderConfig } from './models/settings-model';
+import { AISettingsModel } from './models/settings-model';
 
 import { loadSkillsFromPaths, SkillRegistry } from './skills';
 
@@ -117,7 +129,9 @@ import {
   createDiscoverCommandsTool,
   createExecuteCommandTool
 } from './tools/commands';
+
 import { createDiscoverSkillsTool, createLoadSkillTool } from './tools/skills';
+
 import { createBrowserFetchTool } from './tools/web';
 
 import { AISettingsWidget } from './widgets/ai-settings';
@@ -320,17 +334,16 @@ const chatModelHandler: JupyterFrontEndPlugin<IChatModelHandler> = {
   ],
   optional: [IProviderRegistry, IToolRegistry, ITranslator],
   provides: IChatModelHandler,
-  activate: (
+  activate: async (
     app: JupyterFrontEnd,
-    settingsModel: AISettingsModel,
-    agentManagerFactory: AgentManagerFactory,
+    settingsModel: IAISettingsModel,
+    agentManagerFactory: IAgentManagerFactory,
     docManager: IDocumentManager,
     rmRegistry: IRenderMimeRegistry,
     providerRegistry?: IProviderRegistry,
-    toolRegistry?: IToolRegistry,
-    translator?: ITranslator
-  ): IChatModelHandler => {
-    const trans = (translator ?? nullTranslator).load('jupyterlite_ai');
+    toolRegistry?: IToolRegistry
+  ): Promise<IChatModelHandler> => {
+    await app.serviceManager.ready;
 
     return new ChatModelHandler({
       settingsModel,
@@ -339,7 +352,7 @@ const chatModelHandler: JupyterFrontEndPlugin<IChatModelHandler> = {
       rmRegistry,
       providerRegistry,
       toolRegistry,
-      trans
+      contentsManager: app.serviceManager.contents
     });
   }
 };
@@ -364,22 +377,29 @@ const plugin: JupyterFrontEndPlugin<IChatTracker> = {
     ILayoutRestorer,
     ILabShell,
     INotebookTracker,
-    ITranslator
+    ITranslator,
+    IComponentsRendererFactory,
+    ICommandPalette,
+    IDocumentManager
   ],
   activate: (
     app: JupyterFrontEnd,
     rmRegistry: IRenderMimeRegistry,
     inputToolbarFactory: IInputToolbarRegistryFactory,
     modelHandler: IChatModelHandler,
-    settingsModel: AISettingsModel,
+    settingsModel: IAISettingsModel,
     chatCommandRegistry: IChatCommandRegistry,
     themeManager?: IThemeManager,
     restorer?: ILayoutRestorer,
     labShell?: ILabShell,
     notebookTracker?: INotebookTracker,
-    translator?: ITranslator
+    translator?: ITranslator,
+    chatComponentsFactory?: IComponentsRendererFactory,
+    palette?: ICommandPalette,
+    documentManager?: IDocumentManager
   ): IChatTracker => {
     const trans = (translator ?? nullTranslator).load('jupyterlite_ai');
+
     // Create attachment opener registry to handle file attachments
     const attachmentOpenerRegistry = new AttachmentOpenerRegistry();
     attachmentOpenerRegistry.set('file', attachment => {
@@ -436,7 +456,10 @@ const plugin: JupyterFrontEndPlugin<IChatTracker> = {
           name = `${modelName}-${i}`;
           i += 1;
         }
-        const model = modelHandler.createModel(name, provider);
+        const model = modelHandler.createModel({
+          name,
+          activeProvider: provider
+        });
         return { model };
       },
       getChatNames: async () => {
@@ -493,7 +516,7 @@ const plugin: JupyterFrontEndPlugin<IChatTracker> = {
       chatPanel.disposed.connect(disconnectSettingsButtonListener);
     }
 
-    let tokenUsageWidget: TokenUsageWidget | null = null;
+    let usageWidget: UsageWidget | null = null;
     chatPanel.chatOpened.connect((_, widget) => {
       const model = widget.model as AIChatModel;
 
@@ -511,19 +534,28 @@ const plugin: JupyterFrontEndPlugin<IChatTracker> = {
       model.agentManager.activeProviderChanged.connect(saveTracker);
 
       // Update the token usage widget.
-      tokenUsageWidget?.dispose();
+      usageWidget?.dispose();
 
-      tokenUsageWidget = new TokenUsageWidget({
+      usageWidget = new UsageWidget({
         tokenUsageChanged: model.tokenUsageChanged,
         settingsModel,
         initialTokenUsage: model.agentManager.tokenUsage,
         translator: trans
       });
-      chatPanel.current?.toolbar.insertBefore(
-        'markRead',
-        'token-usage',
-        tokenUsageWidget
-      );
+      chatPanel.current?.toolbar.insertBefore('markRead', 'usage', usageWidget);
+
+      if (model.saveAvailable) {
+        const saveChatButton = new SaveComponentWidget({
+          model,
+          translator: trans
+        });
+
+        chatPanel.current?.toolbar.insertAfter(
+          'markRead',
+          'saveChat',
+          saveChatButton
+        );
+      }
 
       // Listen for writers change to display the stop button.
       function writersChanged(_: IChatModel, writers: IChatModel.IWriter[]) {
@@ -543,11 +575,6 @@ const plugin: JupyterFrontEndPlugin<IChatTracker> = {
 
       model.writersChanged?.connect(writersChanged);
 
-      // Associate an approval buttons object to the chat.
-      const approvalButton = new ApprovalButtons({
-        chatPanel: widget,
-        agentManager: model.agentManager
-      });
       // Temporary compat: keep output-area CSS context for MIME renderers
       // until jupyter-chat provides it natively.
       const outputAreaCompat = new RenderedMessageOutputAreaCompat({
@@ -560,11 +587,17 @@ const plugin: JupyterFrontEndPlugin<IChatTracker> = {
         model.writersChanged?.disconnect(writersChanged);
 
         // Dispose of the approval buttons widget when the chat is disposed.
-        approvalButton.dispose();
         outputAreaCompat.dispose();
       });
     });
 
+    /**
+     * Update the available chat list when settings config changed.
+     */
+    settingsModel.stateChanged.connect(() => {
+      chatPanel.updateChatList();
+    });
+
     app.shell.add(chatPanel, 'left', { rank: 1000 });
 
     if (restorer) {
@@ -602,9 +635,32 @@ const plugin: JupyterFrontEndPlugin<IChatTracker> = {
       modelHandler,
       trans,
       themeManager,
-      labShell
+      labShell,
+      palette,
+      documentManager
     );
 
+    /**
+     * The callback to approve or reject a tool.
+     */
+    function toolCallApproval(
+      targetId: string,
+      approvalId: string,
+      isApproved: boolean
+    ) {
+      const model = tracker.find(chat => chat.model.name === targetId)?.model;
+      if (!model) {
+        return;
+      }
+      isApproved
+        ? (model as AIChatModel).agentManager.approveToolCall(approvalId)
+        : (model as AIChatModel).agentManager.rejectToolCall(approvalId);
+    }
+
+    if (chatComponentsFactory) {
+      chatComponentsFactory.toolCallApproval = toolCallApproval;
+    }
+
     return tracker;
   }
 };
@@ -615,13 +671,15 @@ function registerCommands(
   chatPanel: MultiChatPanel,
   attachmentOpenerRegistry: IAttachmentOpenerRegistry,
   inputToolbarFactory: IInputToolbarRegistryFactory,
-  settingsModel: AISettingsModel,
+  settingsModel: IAISettingsModel,
   chatCommandRegistry: IChatCommandRegistry,
   tracker: WidgetTracker<MainAreaChat | ChatWidget>,
   modelRegistry: IChatModelHandler,
   trans: TranslationBundle,
   themeManager?: IThemeManager,
-  labShell?: ILabShell
+  labShell?: ILabShell,
+  palette?: ICommandPalette,
+  documentManager?: IDocumentManager
 ) {
   const { commands } = app;
 
@@ -735,7 +793,10 @@ function registerCommands(
           name = providerConfig.name;
         }
 
-        const model = modelRegistry.createModel(name, provider);
+        const model = modelRegistry.createModel({
+          name,
+          activeProvider: provider
+        });
         if (!model) {
           return false;
         }
@@ -817,13 +878,12 @@ function registerCommands(
         previousModel.name = UUID.uuid4();
 
         // Create a new model by duplicating the previous model attributes.
-        const model = modelRegistry.createModel(
-          args.name as string,
-          previousModel.agentManager.activeProvider,
-          previousModel.agentManager.tokenUsage
-        );
-        previousModel?.messages.forEach(message => {
-          model?.messageAdded({ ...message.content });
+        const model = modelRegistry.createModel({
+          name: args.name as string,
+          activeProvider: previousModel.agentManager.activeProvider,
+          tokenUsage: previousModel.agentManager.tokenUsage,
+          messages: previousModel.messages,
+          autosave: previousModel.autosave
         });
 
         // Wait (with timeout) for the tracker to have updated the previous widget.
@@ -869,6 +929,135 @@ function registerCommands(
         }
       }
     });
+
+    commands.addCommand(CommandIds.saveChat, {
+      label: args => (args.isPalette ? trans.__('Save chat') : ''),
+      caption: trans.__('Save the chat as local file'),
+      icon: saveIcon,
+      execute: async (args): Promise<boolean> => {
+        let model: AIChatModel | null = null;
+        if (args.name) {
+          tracker.forEach(widget => {
+            if (widget.model.name === args.name) {
+              model = widget.model as AIChatModel;
+            }
+          });
+        } else {
+          model = (tracker.currentWidget?.model as AIChatModel) ?? null;
+        }
+        if (model === null) {
+          console.log('No chat to save');
+          return false;
+        }
+
+        model.save();
+        return true;
+      },
+      describedBy: {
+        args: {
+          type: 'object',
+          properties: {
+            isPalette: {
+              type: 'boolean',
+              description: trans.__('Whether the command is in palette')
+            },
+            name: {
+              type: 'string',
+              description: trans.__('The name of the chat to save')
+            }
+          }
+        }
+      }
+    });
+
+    commands.addCommand(CommandIds.restoreChat, {
+      label: args => (args.isPalette ? trans.__('Restore chat') : ''),
+      caption: trans.__('Restore the chat from a local file'),
+      icon: fileUploadIcon,
+      isVisible: () => !!documentManager,
+      execute: async (args): Promise<boolean> => {
+        if (!documentManager) {
+          console.warn('The restoration is not possible');
+          return false;
+        }
+        let model: AIChatModel | null = null;
+        if (args.name) {
+          tracker.forEach(widget => {
+            if (widget.model.name === args.name) {
+              model = widget.model as AIChatModel;
+            }
+          });
+        } else {
+          model = (tracker.currentWidget?.model as AIChatModel) ?? null;
+        }
+        if (model === null) {
+          console.warn('There is no chat to restore');
+          return false;
+        }
+
+        let backupDirExists = false;
+        await app.serviceManager.contents
+          .get(settingsModel.config.chatBackupDirectory, { content: false })
+          .then(() => (backupDirExists = true))
+          .catch(() => (backupDirExists = false));
+
+        const selection = await FileDialog.getOpenFiles({
+          title: trans.__('Select files to attach'),
+          manager: documentManager,
+          defaultPath: backupDirExists
+            ? settingsModel.config.chatBackupDirectory
+            : ''
+        });
+
+        const filepath = selection.value?.[0].path;
+
+        if (!filepath) {
+          return false;
+        }
+
+        if (model.messages.length) {
+          const result = await showDialog({
+            body: trans.__('All the message will be deleted')
+          });
+          if (!result.button.accept) {
+            return false;
+          }
+        }
+        return model.restore(filepath);
+      },
+      describedBy: {
+        args: {
+          type: 'object',
+          properties: {
+            isPalette: {
+              type: 'boolean',
+              description: trans.__('Whether the command is in palette')
+            },
+            name: {
+              type: 'string',
+              description: trans.__('The name of the chat to save')
+            }
+          }
+        }
+      }
+    });
+
+    if (palette) {
+      palette.addItem({
+        category: trans.__('AI Assistant'),
+        command: CommandIds.saveChat,
+        args: {
+          isPalette: true
+        }
+      });
+      palette.addItem({
+        category: trans.__('AI Assistant'),
+        command: CommandIds.restoreChat,
+        args: {
+          isPalette: true
+        }
+      });
+    }
   }
 }
 
@@ -876,7 +1065,7 @@ function registerCommands(
  * A plugin to provide the agent manager factory and completion provider.
  * These objects require the secrets manager token with the same namespace.
  */
-const agentManagerFactory: JupyterFrontEndPlugin<AgentManagerFactory> =
+const agentManagerFactory: JupyterFrontEndPlugin<IAgentManagerFactory> =
   SecretsManager.sign(SECRETS_NAMESPACE, token => {
     Private.setAISecretsToken(token);
 
@@ -889,12 +1078,12 @@ const agentManagerFactory: JupyterFrontEndPlugin<AgentManagerFactory> =
       optional: [ISkillRegistry, ICompletionProviderManager, ISecretsManager],
       activate: (
         app: JupyterFrontEnd,
-        settingsModel: AISettingsModel,
+        settingsModel: IAISettingsModel,
         providerRegistry: IProviderRegistry,
         skillRegistry?: ISkillRegistry,
         completionManager?: ICompletionProviderManager,
         secretsManager?: ISecretsManager
-      ): AgentManagerFactory => {
+      ): IAgentManagerFactory => {
         const agentManagerFactory = new AgentManagerFactory({
           settingsModel,
           skillRegistry,
@@ -940,8 +1129,8 @@ const settingsPanelPlugin: JupyterFrontEndPlugin<void> = {
   ],
   activate: (
     app: JupyterFrontEnd,
-    settingsModel: AISettingsModel,
-    agentManagerFactory: AgentManagerFactory,
+    settingsModel: IAISettingsModel,
+    agentManagerFactory: IAgentManagerFactory,
     providerRegistry: IProviderRegistry,
     palette?: ICommandPalette,
     restorer?: ILayoutRestorer,
@@ -1005,13 +1194,16 @@ const settingsPanelPlugin: JupyterFrontEndPlugin<void> = {
 /**
  * Built-in completion providers plugin
  */
-const settingsModel: JupyterFrontEndPlugin<AISettingsModel> = {
+const settingsModel: JupyterFrontEndPlugin<IAISettingsModel> = {
   id: '@jupyterlite/ai:settings-model',
   description: 'Provide the AI settings model',
   autoStart: true,
   provides: IAISettingsModel,
   requires: [ISettingRegistry],
-  activate: (app: JupyterFrontEnd, settingRegistry: ISettingRegistry) => {
+  activate: (
+    app: JupyterFrontEnd,
+    settingRegistry: ISettingRegistry
+  ): IAISettingsModel => {
     return new AISettingsModel({ settingRegistry });
   }
 };
@@ -1027,7 +1219,7 @@ const diffManager: JupyterFrontEndPlugin<IDiffManager> = {
   requires: [IAISettingsModel],
   activate: (
     app: JupyterFrontEnd,
-    settingsModel: AISettingsModel
+    settingsModel: IAISettingsModel
   ): IDiffManager => {
     return new DiffManager({
       commands: app.commands,
@@ -1058,7 +1250,7 @@ const toolRegistry: JupyterFrontEndPlugin<IToolRegistry> = {
   provides: IToolRegistry,
   activate: (
     app: JupyterFrontEnd,
-    settingsModel: AISettingsModel,
+    settingsModel: IAISettingsModel,
     skillRegistry?: ISkillRegistry
   ) => {
     const toolRegistry = new ToolRegistry();
@@ -1098,7 +1290,7 @@ const inputToolbarFactory: JupyterFrontEndPlugin<IInputToolbarRegistryFactory> =
     optional: [ITranslator],
     activate: (
       app: JupyterFrontEnd,
-      settingsModel: AISettingsModel,
+      settingsModel: IAISettingsModel,
       toolRegistry: IToolRegistry,
       providerRegistry: IProviderRegistry,
       translator?: ITranslator
@@ -1148,7 +1340,7 @@ const completionStatus: JupyterFrontEndPlugin<void> = {
   optional: [IStatusBar, ITranslator],
   activate: (
     app: JupyterFrontEnd,
-    settingsModel: AISettingsModel,
+    settingsModel: IAISettingsModel,
     statusBar: IStatusBar | null,
     translator?: ITranslator
   ) => {
@@ -1179,7 +1371,7 @@ const skillsPlugin: JupyterFrontEndPlugin<void> = {
   optional: [ICommandPalette, ITranslator],
   activate: async (
     app: JupyterFrontEnd,
-    settingsModel: AISettingsModel,
+    settingsModel: IAISettingsModel,
     docManager: IDocumentManager,
     skillRegistry: ISkillRegistry,
     palette?: ICommandPalette,