Enable/disable AI agent in chats

packages/jupyter-ai/jupyter_ai/extension.py

@@ -18,7 +18,7 @@
 from jupyter_server.extension.application import ExtensionApp
 from jupyter_server.utils import url_path_join
 from jupyterlab_collaborative_chat.ychat import YChat
-from pycrdt import ArrayEvent
+from pycrdt import ArrayEvent, MapEvent
 from tornado.web import StaticFileHandler
 from traitlets import Dict, Integer, List, Unicode
 
@@ -239,6 +239,9 @@ def initialize(self):
         # updating it.
         self.messages_indexes = {}
 
+        # The subscriptions to chats messages.
+        self.subscriptions = {}
+
     async def connect_chat(
         self, logger: EventLogger, schema_id: str, data: dict
     ) -> None:
@@ -257,8 +260,14 @@ async def connect_chat(
             )
             chat.awareness.set_local_state_field("user", BOT)
 
-            callback = partial(self.on_change, chat)
-            chat.ymessages.observe(callback)
+            # Check if the AI agent should be connected to the chat or not.
+            chat_meta = chat.get_metadata()
+            if "agents" in chat_meta and BOT["username"] in chat.get_metadata()["agents"]:
+                message_callback = partial(self.on_message_change, chat)
+                self.subscriptions[chat.get_id()] = chat.ymessages.observe(message_callback)
+
+            metadata_callback = partial(self.on_metadata_change, chat)
+            chat.ymetadata.observe(metadata_callback)
 
     async def get_chat(self, room_id: str) -> YChat:
         if COLLAB_VERSION == 3:
@@ -272,7 +281,7 @@ async def get_chat(self, room_id: str) -> YChat:
             document = room._document
         return document
 
-    def on_change(self, chat: YChat, events: ArrayEvent) -> None:
+    def on_message_change(self, chat: YChat, events: ArrayEvent) -> None:
         for change in events.delta:
             if not "insert" in change.keys():
                 continue
@@ -296,6 +305,22 @@ def on_change(self, chat: YChat, events: ArrayEvent) -> None:
                     self._route(chat_message, chat)
                 )
 
+    def on_metadata_change(self, chat: YChat, events: MapEvent) -> None:
+        """
+        Triggered when a chat metadata has changed.
+        It will connect or disconnect the AI agent from the chat.
+        """
+        if "agents" in events.keys:
+            if BOT["username"] in events.keys["agents"]["newValue"]:
+                message_callback = partial(self.on_message_change, chat)
+                self.subscriptions[chat.get_id()] = chat.ymessages.observe(message_callback)
+            else:
+                chat.ymessages.unobserve(self.subscriptions[chat.get_id()])
+                del self.subscriptions[chat.get_id()]
+
+            # Update the bot user avatar
+            chat.awareness.set_local_state_field("user", BOT)
+
     async def _route(self, message: HumanChatMessage, chat: YChat):
         """Method that routes an incoming message to the appropriate handler."""
         chat_handlers = self.settings["jai_chat_handlers"]

packages/jupyter-ai/package.json

@@ -62,7 +62,8 @@
     "@emotion/react": "^11.10.5",
     "@emotion/styled": "^11.10.5",
     "@jupyter/chat": "^0.5.0",
-    "@jupyter/collaboration": "^1",
+    "@jupyter/collaboration": "^2.1.3",
+    "@jupyter/docprovider": "^2.1.3",
     "@jupyterlab/application": "^4.2.0",
     "@jupyterlab/apputils": "^4.2.0",
     "@jupyterlab/cells": "^4.2.0",
@@ -79,6 +80,7 @@
     "@jupyterlab/ui-components": "^4.2.0",
     "@mui/icons-material": "^5.11.0",
     "@mui/material": "^5.11.0",
+    "jupyterlab-collaborative-chat": "0.5.0",
     "react": "^18.2.0",
     "react-dom": "^18.2.0"
   },

packages/jupyter-ai/src/index.ts

@@ -1,5 +1,5 @@
 import { IAutocompletionRegistry } from '@jupyter/chat';
-import { IGlobalAwareness } from '@jupyter/collaboration';
+import { IGlobalAwareness, UsersItem } from '@jupyter/collaboration';
 import {
   JupyterFrontEnd,
   JupyterFrontEndPlugin,
@@ -10,18 +10,23 @@ import {
   ReactWidget,
   IThemeManager,
   MainAreaWidget,
-  ICommandPalette
+  ICommandPalette,
+  IToolbarWidgetRegistry
 } from '@jupyterlab/apputils';
 import { IDocumentWidget } from '@jupyterlab/docregistry';
 import { IRenderMimeRegistry } from '@jupyterlab/rendermime';
 import { Signal } from '@lumino/signaling';
+import {
+  CollaborativeChatPanel,
+  IChatFactory
+} from 'jupyterlab-collaborative-chat';
 import type { Awareness } from 'y-protocols/awareness';
 
 import { ChatHandler } from './chat_handler';
 import { completionPlugin } from './completions';
 import { ActiveCellManager } from './contexts/active-cell-context';
-import { SelectionWatcher } from './selection-watcher';
 import { menuPlugin } from './plugins/menu-plugin';
+import { SelectionWatcher } from './selection-watcher';
 import { autocompletion } from './slash-autocompletion';
 import { statusItemPlugin } from './status';
 import {
@@ -30,6 +35,7 @@ import {
   IJaiMessageFooter,
   IJaiTelemetryHandler
 } from './tokens';
+import { userIconRenderer } from './user-icon';
 import { buildErrorWidget } from './widgets/chat-error';
 import { buildChatSidebar } from './widgets/chat-sidebar';
 import { buildAiSettings } from './widgets/settings-widget';
@@ -186,7 +192,7 @@ const plugin: JupyterFrontEndPlugin<IJaiCore> = {
 /**
  * Add slash commands to collaborative chat.
  */
-const collaborative_autocompletion: JupyterFrontEndPlugin<void> = {
+const collaborativeAutocompletion: JupyterFrontEndPlugin<void> = {
   id: '@jupyter-ai/core:autocompletion',
   autoStart: true,
   requires: [IAutocompletionRegistry],
@@ -198,12 +204,37 @@ const collaborative_autocompletion: JupyterFrontEndPlugin<void> = {
   }
 };
 
+/**
+ * Customize users item toolbar widget.
+ */
+const usersItem: JupyterFrontEndPlugin<void> = {
+  id: '@jupyter-ai/core:users_item',
+  autoStart: true,
+  requires: [IChatFactory, IToolbarWidgetRegistry],
+  activate: async (
+    app: JupyterFrontEnd,
+    chatFactory: IChatFactory,
+    toolbarRegistry: IToolbarWidgetRegistry
+  ) => {
+    toolbarRegistry.addFactory<CollaborativeChatPanel>(
+      'Chat',
+      'usersItem',
+      panel =>
+        UsersItem.createWidget({
+          model: panel.model,
+          iconRenderer: userIconRenderer
+        })
+    );
+  }
+};
+
 export default [
   plugin,
   statusItemPlugin,
   completionPlugin,
   menuPlugin,
-  collaborative_autocompletion
+  collaborativeAutocompletion,
+  usersItem
 ];
 
 export * from './contexts';

packages/jupyter-ai/src/user-icon.tsx

@@ -0,0 +1,75 @@
+import { DefaultIconRenderer, UsersItem } from '@jupyter/collaboration';
+import {
+  CollaborativeChatModel,
+  IChatChanges,
+  YChat
+} from 'jupyterlab-collaborative-chat';
+import React, { useEffect, useRef, useState } from 'react';
+
+/**
+ * The user icon renderer.
+ */
+export const userIconRenderer = (
+  props: UsersItem.IIconRendererProps
+): JSX.Element => {
+  const { user } = props;
+
+  if (user.userData.name === 'Jupyternaut') {
+    return <AgentIconRenderer {...props} />;
+  }
+
+  return <DefaultIconRenderer user={user} />;
+};
+
+/**
+ * The user icon renderer for an AI agent.
+ * It modify the metadata of the ydocument to enable/disable the AI agent.
+ */
+const AgentIconRenderer = (
+  props: UsersItem.IIconRendererProps
+): JSX.Element => {
+  const { user, model } = props;
+  const [iconClass, setIconClass] = useState<string>('');
+  const sharedModel = useRef<YChat>(
+    (model as CollaborativeChatModel).sharedModel
+  );
+
+  useEffect(() => {
+    // Update the icon class.
+    const updateStatus = () => {
+      const agents =
+        (sharedModel.current.getMetadata('agents') as string[]) || [];
+      setIconClass(
+        agents.includes(user.userData.username) ? '' : 'disconnected'
+      );
+    };
+
+    const onChange = (_: YChat, changes: IChatChanges) => {
+      if (changes.metadataChanges) {
+        updateStatus();
+      }
+    };
+
+    sharedModel.current.changed.connect(onChange);
+    updateStatus();
+    return () => {
+      sharedModel.current.changed.disconnect(updateStatus);
+    };
+  }, [model]);
+
+  const onclick = () => {
+    const agents =
+      (sharedModel.current.getMetadata('agents') as string[]) || [];
+    const index = agents.indexOf(user.userData.username);
+    if (index > -1) {
+      agents.splice(index, 1);
+    } else {
+      agents.push(user.userData.username);
+    }
+    sharedModel.current.setMetadata('agents', agents);
+  };
+
+  return (
+    <DefaultIconRenderer user={user} onClick={onclick} className={iconClass} />
+  );
+};

packages/jupyter-ai/style/base.css

@@ -7,3 +7,4 @@
 @import './expandable-text-field.css';
 @import './chat-settings.css';
 @import './rendermime-markdown.css';
+@import './user-icon.css';

packages/jupyter-ai/style/user-icon.css

@@ -0,0 +1,13 @@
+.jp-toolbar-users-item .jp-MenuBar-imageIcon.disconnected {
+  opacity: 0.5;
+}
+
+.jp-toolbar-users-item .jp-MenuBar-imageIcon.disconnected:before {
+  position: absolute;
+  left: 11px;
+  content: '';
+  height: 28px;
+  width: 2px;
+  background-color: #333;
+  transform: rotate(45deg);
+}