Stop streaming (#61)

* Add a button to stop streaming

* Update dependency to jupyter-chat 0.9.0

Author: brichet

package.json

@@ -56,7 +56,7 @@
         "watch:labextension": "jupyter labextension watch ."
     },
     "dependencies": {
-        "@jupyter/chat": "^0.8.1",
+        "@jupyter/chat": "^0.9.0",
         "@jupyterlab/application": "^4.4.0-alpha.0",
         "@jupyterlab/apputils": "^4.5.0-alpha.0",
         "@jupyterlab/completer": "^4.4.0-alpha.0",

src/chat-handler.ts

@@ -139,9 +139,11 @@ export class ChatHandler extends ChatModel {
 
     let content = '';
 
+    this._controller = new AbortController();
     try {
       for await (const chunk of await this._providerRegistry.currentChatModel.stream(
-        messages
+        messages,
+        { signal: this._controller.signal }
       )) {
         content += chunk.content ?? chunk;
         botMsg.body = content;
@@ -162,6 +164,7 @@ export class ChatHandler extends ChatModel {
       return false;
     } finally {
       this.updateWriters([]);
+      this._controller = null;
     }
   }
 
@@ -177,12 +180,17 @@ export class ChatHandler extends ChatModel {
     super.messageAdded(message);
   }
 
+  stopStreaming(): void {
+    this._controller?.abort();
+  }
+
   private _providerRegistry: IAIProviderRegistry;
   private _personaName = 'AI';
   private _prompt: string;
   private _errorMessage: string = '';
   private _history: IChatHistory = { messages: [] };
   private _defaultErrorMessage = 'AI provider not configured';
+  private _controller: AbortController | null = null;
 }
 
 export namespace ChatHandler {

src/components/stop-button.tsx

@@ -0,0 +1,56 @@
+/*
+ * Copyright (c) Jupyter Development Team.
+ * Distributed under the terms of the Modified BSD License.
+ */
+
+import StopIcon from '@mui/icons-material/Stop';
+import React from 'react';
+
+import { InputToolbarRegistry, TooltippedButton } from '@jupyter/chat';
+
+/**
+ * Properties of the stop button.
+ */
+export interface IStopButtonProps
+  extends InputToolbarRegistry.IToolbarItemProps {
+  /**
+   * The function to stop streaming.
+   */
+  stopStreaming: () => void;
+}
+
+/**
+ * The stop button.
+ */
+export function StopButton(props: IStopButtonProps): JSX.Element {
+  const tooltip = 'Stop streaming';
+  return (
+    <TooltippedButton
+      onClick={props.stopStreaming}
+      tooltip={tooltip}
+      buttonProps={{
+        size: 'small',
+        variant: 'contained',
+        title: tooltip
+      }}
+    >
+      <StopIcon />
+    </TooltippedButton>
+  );
+}
+
+/**
+ * factory returning the toolbar item.
+ */
+export function stopItem(
+  stopStreaming: () => void
+): InputToolbarRegistry.IToolbarItem {
+  return {
+    element: (props: InputToolbarRegistry.IToolbarItemProps) => {
+      const stopProps: IStopButtonProps = { ...props, stopStreaming };
+      return StopButton(stopProps);
+    },
+    position: 50,
+    hidden: true /* hidden by default */
+  };
+}

src/index.ts

@@ -4,7 +4,8 @@ import {
   buildErrorWidget,
   ChatCommandRegistry,
   IActiveCellManager,
-  IChatCommandRegistry
+  IChatCommandRegistry,
+  InputToolbarRegistry
 } from '@jupyter/chat';
 import {
   JupyterFrontEnd,
@@ -28,6 +29,7 @@ import { defaultProviderPlugins } from './default-providers';
 import { AIProviderRegistry } from './provider';
 import { aiSettingsRenderer, SettingConnector } from './settings';
 import { IAIProviderRegistry } from './tokens';
+import { stopItem } from './components/stop-button';
 
 const chatCommandRegistryPlugin: JupyterFrontEndPlugin<IChatCommandRegistry> = {
   id: '@jupyterlite/ai:autocompletion-registry',
@@ -102,12 +104,30 @@ const chatPlugin: JupyterFrontEndPlugin<void> = {
       });
 
     let chatWidget: ReactWidget | null = null;
+
+    const inputToolbarRegistry = InputToolbarRegistry.defaultToolbarRegistry();
+    const stopButton = stopItem(() => chatHandler.stopStreaming());
+    inputToolbarRegistry.addItem('stop', stopButton);
+
+    chatHandler.writersChanged.connect((_, users) => {
+      if (
+        users.filter(user => user.username === chatHandler.personaName).length
+      ) {
+        inputToolbarRegistry.hide('send');
+        inputToolbarRegistry.show('stop');
+      } else {
+        inputToolbarRegistry.hide('stop');
+        inputToolbarRegistry.show('send');
+      }
+    });
+
     try {
       chatWidget = buildChatSidebar({
         model: chatHandler,
         themeManager,
         rmRegistry,
-        chatCommandRegistry
+        chatCommandRegistry,
+        inputToolbarRegistry
       });
       chatWidget.title.caption = 'Jupyterlite AI Chat';
     } catch (e) {

yarn.lock

@@ -906,9 +906,9 @@ __metadata:
   languageName: node
   linkType: hard
 
-"@jupyter/chat@npm:^0.8.1":
-  version: 0.8.1
-  resolution: "@jupyter/chat@npm:0.8.1"
+"@jupyter/chat@npm:^0.9.0":
+  version: 0.9.0
+  resolution: "@jupyter/chat@npm:0.9.0"
   dependencies:
     "@emotion/react": ^11.10.5
     "@emotion/styled": ^11.10.5
@@ -930,7 +930,7 @@ __metadata:
     clsx: ^2.1.0
     react: ^18.2.0
     react-dom: ^18.2.0
-  checksum: d4a3635cd419642c62ca462ef9ec5327178bed7f898684e31a193246b081707685ba825ef0a83e3b6618316a7e7ffd68141194759cb1f2b107215cbe33ef5d59
+  checksum: dee9a9e02ea8a6d25e1caebf4e9bece42c82fa40baa84dc655540bb64a676b4b1890373390c291f35eb1fe8f821d15935ec21fdd0e12063a497e575f4ddbcd78
   languageName: node
   linkType: hard
 
@@ -2108,7 +2108,7 @@ __metadata:
   version: 0.0.0-use.local
   resolution: "@jupyterlite/ai@workspace:."
   dependencies:
-    "@jupyter/chat": ^0.8.1
+    "@jupyter/chat": ^0.9.0
     "@jupyterlab/application": ^4.4.0-alpha.0
     "@jupyterlab/apputils": ^4.5.0-alpha.0
     "@jupyterlab/builder": ^4.0.0