diff --git a/.gitignore b/.gitignore index f302760ca..e2e77ef3b 100644 --- a/.gitignore +++ b/.gitignore @@ -127,4 +127,6 @@ playground/ # reserve path for a dev script dev.sh -.vscode \ No newline at end of file +.vscode + +.jupyter_ystore.db diff --git a/packages/jupyter-ai/jupyter_ai/extension.py b/packages/jupyter-ai/jupyter_ai/extension.py index d934ce13e..e59d43c61 100644 --- a/packages/jupyter-ai/jupyter_ai/extension.py +++ b/packages/jupyter-ai/jupyter_ai/extension.py @@ -1,7 +1,7 @@ import queue from jupyter_server.extension.application import ExtensionApp from langchain import ConversationChain -from .handlers import ChatHandler, ChatHistoryHandler, PromptAPIHandler, TaskAPIHandler, ChatAPIHandler +from .handlers import ChatHandler, ChatHistoryHandler, PromptAPIHandler, TaskAPIHandler from importlib_metadata import entry_points import inspect from .engine import BaseModelEngine @@ -20,7 +20,6 @@ class AiExtension(ExtensionApp): name = "jupyter_ai" handlers = [ ("api/ai/prompt", PromptAPIHandler), - (r"api/ai/chat/?", ChatAPIHandler), (r"api/ai/tasks/?", TaskAPIHandler), (r"api/ai/tasks/([\w\-:]*)", TaskAPIHandler), (r"api/ai/chats/?", ChatHandler), @@ -113,5 +112,7 @@ def initialize_settings(self): # Store chat clients in a dictionary self.settings["chat_clients"] = {} + self.settings["chat_handlers"] = {} - \ No newline at end of file + # store chat messages in memory for now + self.settings["chat_history"] = [] diff --git a/packages/jupyter-ai/jupyter_ai/handlers.py b/packages/jupyter-ai/jupyter_ai/handlers.py index 6ef431ace..d8d7849cd 100644 --- a/packages/jupyter-ai/jupyter_ai/handlers.py +++ b/packages/jupyter-ai/jupyter_ai/handlers.py @@ -1,8 +1,9 @@ from dataclasses import asdict import json -from typing import Optional - +from typing import Dict, List import tornado +import uuid + from tornado.web import HTTPError from pydantic import ValidationError @@ -12,8 +13,8 @@ from jupyter_server.utils import ensure_async from .task_manager import TaskManager -from .models import ChatHistory, PromptRequest, ChatRequest -from langchain.schema import _message_to_dict, HumanMessage, AIMessage +from .models import ChatHistory, PromptRequest, ChatRequest, ChatMessage, AgentChatMessage, HumanChatMessage, ConnectionMessage, ChatClient +from langchain.schema import HumanMessage class APIHandler(BaseAPIHandler): @property @@ -60,26 +61,6 @@ async def post(self): "insertion_mode": task.insertion_mode })) -class ChatAPIHandler(APIHandler): - @tornado.web.authenticated - async def post(self): - try: - request = ChatRequest(**self.get_json_body()) - except ValidationError as e: - self.log.exception(e) - raise HTTPError(500, str(e)) from e - - if not self.openai_chat: - raise HTTPError(500, "No chat models available.") - - result = await ensure_async(self.openai_chat.agenerate([request.prompt])) - output = result.generations[0][0].text - self.openai_chat.append_exchange(request.prompt, output) - - self.finish(json.dumps({ - "output": output, - })) - class TaskAPIHandler(APIHandler): @tornado.web.authenticated async def get(self, id=None): @@ -106,25 +87,24 @@ def chat_provider(self): if self._chat_provider is None: self._chat_provider = self.settings["chat_provider"] return self._chat_provider - - @property - def messages(self): - self._messages = self.chat_provider.memory.chat_memory.messages or [] - return self._messages + @property + def chat_history(self): + return self.settings["chat_history"] + + @chat_history.setter + def _chat_history_setter(self, new_history): + self.settings["chat_history"] = new_history + @tornado.web.authenticated async def get(self): - messages = [] - for message in self.messages: - messages.append(message) - history = ChatHistory(messages=messages) - - self.finish(history.json(models_as_dict=False)) + history = ChatHistory(messages=self.chat_history) + self.finish(history.json()) @tornado.web.authenticated async def delete(self): self.chat_provider.memory.chat_memory.clear() - self.messages = [] + self.chat_history = [] self.set_status(204) self.finish() @@ -154,18 +134,26 @@ def chat_message_queue(self): return self._chat_message_queue @property - def messages(self): - self._messages = self.chat_provider.memory.chat_memory.messages or [] - return self._messages + def chat_handlers(self) -> Dict[str, 'ChatHandler']: + """Dictionary mapping client IDs to their WebSocket handler + instances.""" + return self.settings["chat_handlers"] - def add_chat_client(self, username): - self.settings["chat_clients"][username] = self - self.log.debug("Clients are : %s", self.settings["chat_clients"].keys()) + @property + def chat_clients(self) -> Dict[str, ChatClient]: + """Dictionary mapping client IDs to their ChatClient objects that store + metadata.""" + return self.settings["chat_clients"] - def remove_chat_client(self, username): - self.settings["chat_clients"][username] = None - self.log.debug("Chat clients: %s", self.settings['chat_clients'].keys()) + @property + def chat_client(self) -> ChatClient: + """Returns ChatClient object associated with the current connection.""" + return self.chat_clients[self.client_id] + @property + def chat_history(self) -> List[ChatMessage]: + return self.settings["chat_history"] + def initialize(self): self.log.debug("Initializing websocket connection %s", self.request.path) @@ -188,24 +176,45 @@ async def get(self, *args, **kwargs): res = super().get(*args, **kwargs) await res + def generate_client_id(self): + """Generates a client ID to identify the current WS connection.""" + # if collaborative mode is enabled, each client already has a UUID + # collaborative = self.config.get("LabApp", {}).get("collaborative", False) + # if collaborative: + # return self.current_user.username + + # if collaborative mode is not enabled, each client is assigned a UUID + return uuid.uuid4().hex + def open(self): - self.log.debug("Client with user %s connected...", self.current_user.username) - self.add_chat_client(self.current_user.username) + """Handles opening of a WebSocket connection. Client ID can be retrieved + from `self.client_id`.""" + + client_id = self.generate_client_id() + chat_client_kwargs = {k: v for k, v in asdict(self.current_user).items() if k != "username"} - def broadcast_message(self, message: any, exclude_current_user: Optional[bool] = False): - """Broadcasts message to all connected clients, - optionally excluding the current user + self.chat_handlers[client_id] = self + self.chat_clients[client_id] = ChatClient(**chat_client_kwargs, id=client_id) + self.client_id = client_id + self.write_message(ConnectionMessage(client_id=client_id).dict()) + + self.log.info(f"Client connected. ID: {client_id}") + self.log.debug("Clients are : %s", self.chat_handlers.keys()) + + def broadcast_message(self, message: ChatMessage): + """Broadcasts message to all connected clients, optionally excluding the + current user. Appends message to `self.chat_history`. """ self.log.debug("Broadcasting message: %s to all clients...", message) - client_names = self.settings["chat_clients"].keys() - if exclude_current_user: - client_names = client_names - [self.current_user.username] + client_ids = self.chat_handlers.keys() - for username in client_names: - client = self.settings["chat_clients"][username] + for client_id in client_ids: + client = self.chat_handlers[client_id] if client: - client.write_message(message) + client.write_message(message.dict()) + + self.chat_history.append(message) async def on_message(self, message): self.log.debug("Message recieved: %s", message) @@ -217,24 +226,39 @@ async def on_message(self, message): self.log.error(e) return + # message sent to the agent instance message = HumanMessage( content=chat_request.prompt, additional_kwargs=dict(user=asdict(self.current_user)) ) - data = json.dumps(_message_to_dict(message)) + # message broadcast to chat clients + chat_message_id = str(uuid.uuid4()) + chat_message = HumanChatMessage( + id=chat_message_id, + body=chat_request.prompt, + client=self.chat_client, + ) + # broadcast the message to other clients - self.broadcast_message(message=data, exclude_current_user=True) + self.broadcast_message(message=chat_message) # process the message response = await ensure_async(self.chat_provider.apredict(input=message.content)) - - response = AIMessage( - content=response + agent_message = AgentChatMessage( + id=str(uuid.uuid4()), + body=response, + reply_to=chat_message_id ) + # broadcast to all clients - self.broadcast_message(message=json.dumps(_message_to_dict(response))) + self.broadcast_message(message=agent_message) def on_close(self): - self.log.debug("Disconnecting client with user %s", self.current_user.username) - self.remove_chat_client(self.current_user.username) + self.log.debug("Disconnecting client with user %s", self.client_id) + + self.chat_handlers.pop(self.client_id, None) + self.chat_clients.pop(self.client_id, None) + + self.log.info(f"Client disconnected. ID: {self.client_id}") + self.log.debug("Chat clients: %s", self.chat_handlers.keys()) diff --git a/packages/jupyter-ai/jupyter_ai/models.py b/packages/jupyter-ai/jupyter_ai/models.py index e4ca60cfe..71e2e8740 100644 --- a/packages/jupyter-ai/jupyter_ai/models.py +++ b/packages/jupyter-ai/jupyter_ai/models.py @@ -1,16 +1,52 @@ -from pydantic import BaseModel, validator -from typing import Dict, List, Literal - -from langchain.schema import BaseMessage, _message_to_dict +from pydantic import BaseModel +from typing import Dict, List, Union, Literal, Optional class PromptRequest(BaseModel): task_id: str engine_id: str prompt_variables: Dict[str, str] +# the type of message used to chat with the agent class ChatRequest(BaseModel): prompt: str +class ChatClient(BaseModel): + id: str + initials: str + name: str + display_name: str + color: Optional[str] + avatar_url: Optional[str] + +class AgentChatMessage(BaseModel): + type: Literal["agent"] = "agent" + id: str + body: str + # message ID of the HumanChatMessage it is replying to + reply_to: str + +class HumanChatMessage(BaseModel): + type: Literal["human"] = "human" + id: str + body: str + client: ChatClient + +class ConnectionMessage(BaseModel): + type: Literal["connection"] = "connection" + client_id: str + +# the type of messages being broadcast to clients +ChatMessage = Union[ + AgentChatMessage, + HumanChatMessage, +] + +Message = Union[ + AgentChatMessage, + HumanChatMessage, + ConnectionMessage +] + class ListEnginesEntry(BaseModel): id: str name: str @@ -30,9 +66,4 @@ class DescribeTaskResponse(BaseModel): class ChatHistory(BaseModel): """History of chat messages""" - messages: List[BaseMessage] - - class Config: - json_encoders = { - BaseMessage: lambda v: _message_to_dict(v) - } \ No newline at end of file + messages: List[ChatMessage] diff --git a/packages/jupyter-ai/src/chat_handler.ts b/packages/jupyter-ai/src/chat_handler.ts index e94c77eb8..c8f029f55 100644 --- a/packages/jupyter-ai/src/chat_handler.ts +++ b/packages/jupyter-ai/src/chat_handler.ts @@ -1,27 +1,104 @@ import { IDisposable } from '@lumino/disposable'; import { ServerConnection } from '@jupyterlab/services'; import { URLExt } from '@jupyterlab/coreutils'; -import { Poll } from '@lumino/polling'; import { AiService, requestAPI } from './handler'; const CHAT_SERVICE_URL = 'api/ai/chats'; export class ChatHandler implements IDisposable { + /** + * The server settings used to make API requests. + */ + readonly serverSettings: ServerConnection.ISettings; + + /** + * ID of the connection. Requires `await initialize()`. + */ + id: string = ''; + /** * Create a new chat handler. */ constructor(options: AiService.IOptions = {}) { this.serverSettings = options.serverSettings ?? ServerConnection.makeSettings(); + } + + /** + * Initializes the WebSocket connection to the Chat backend. Promise is + * resolved when server acknowledges connection and sends the client ID. This + * must be awaited before calling any other method. + */ + public initialize(): Promise { + return new Promise((resolve, reject) => { + if (this.isDisposed) { + return; + } + const { token, WebSocket, wsUrl } = this.serverSettings; + const url = + URLExt.join(wsUrl, CHAT_SERVICE_URL) + + (token ? `?token=${encodeURIComponent(token)}` : ''); + + const socket = (this._socket = new WebSocket(url)); + socket.onmessage = msg => + msg.data && this._onMessage(JSON.parse(msg.data)); - this._poll = new Poll({ factory: () => this._subscribe() }); - this._poll.start(); + const listenForConnection = (message: AiService.Message) => { + if (message.type !== 'connection') { + return; + } + this.id = message.client_id; + resolve(); + this.removeListener(listenForConnection); + }; + + this.addListener(listenForConnection); + }); } /** - * The server settings used to make API requests. + * Sends a message across the WebSocket. Promise resolves to the message ID + * when the server sends the same message back, acknowledging receipt. */ - readonly serverSettings: ServerConnection.ISettings; + public sendMessage(message: AiService.ChatRequest): Promise { + return new Promise(resolve => { + this._socket?.send(JSON.stringify(message)); + this._sendResolverQueue.push(resolve); + }); + } + + /** + * Returns a Promise that resolves to the agent's reply, given the message ID + * of the human message. Should only be called once per message. + */ + public replyFor(messageId: string): Promise { + return new Promise(resolve => { + this._replyForResolverDict[messageId] = resolve; + }); + } + + public addListener(handler: (message: AiService.Message) => void): void { + this._listeners.push(handler); + } + + public removeListener(handler: (message: AiService.Message) => void): void { + const index = this._listeners.indexOf(handler); + if (index > -1) { + this._listeners.splice(index, 1); + } + } + + public async getHistory(): Promise { + let data: AiService.ChatHistory = { messages: [] }; + try { + data = await requestAPI('chats/history', { + method: 'GET' + }); + } catch (e) { + return Promise.reject(e); + } + return data; + } /** * Whether the chat handler is disposed. @@ -39,9 +116,6 @@ export class ChatHandler implements IDisposable { } this._isDisposed = true; - // Clean up poll. - this._poll.dispose(); - this._listeners = []; // Clean up socket. @@ -56,58 +130,40 @@ export class ChatHandler implements IDisposable { } } - public addListener(handler: (message: AiService.ChatMessage) => void): void { - this._listeners.push(handler); - } - - public removeListener( - handler: (message: AiService.ChatMessage) => void - ): void { - const index = this._listeners.indexOf(handler); - if (index > -1) { - this._listeners.splice(index, 1); + private _onMessage(message: AiService.Message): void { + // resolve promise from `sendMessage()` + if (message.type === 'human' && message.client.id === this.id) { + this._sendResolverQueue.shift()?.(message.id); } - } - public sendMessage(message: AiService.ChatRequest): void { - this._socket?.send(JSON.stringify(message)); - } - - public async getHistory(): Promise { - let data: AiService.ChatHistory = { messages: [] }; - try { - data = await requestAPI('chats/history', { - method: 'GET' - }); - } catch (e) { - return Promise.reject(e); + // resolve promise from `replyFor()` if it exists + if ( + message.type === 'agent' && + message.reply_to in this._replyForResolverDict + ) { + this._replyForResolverDict[message.reply_to](message); + delete this._replyForResolverDict[message.reply_to]; } - return data; - } - private _onMessage(message: AiService.ChatMessage): void { + // call listeners in serial this._listeners.forEach(listener => listener(message)); } - private _subscribe(): Promise { - return new Promise((_, reject) => { - if (this.isDisposed) { - return; - } - const { token, WebSocket, wsUrl } = this.serverSettings; - const url = - URLExt.join(wsUrl, CHAT_SERVICE_URL) + - (token ? `?token=${encodeURIComponent(token)}` : ''); - const socket = (this._socket = new WebSocket(url)); + /** + * Queue of Promise resolvers pushed onto by `send()` + */ + private _sendResolverQueue: ((value: string) => void)[] = []; - socket.onclose = () => reject(new Error('ChatHandler socket closed')); - socket.onmessage = msg => - msg.data && this._onMessage(JSON.parse(msg.data)); - }); - } + /** + * Dictionary mapping message IDs to Promise resolvers, inserted into by + * `replyFor()`. + */ + private _replyForResolverDict: Record< + string, + (value: AiService.AgentChatMessage) => void + > = {}; private _isDisposed = false; - private _poll: Poll; private _socket: WebSocket | null = null; private _listeners: ((msg: any) => void)[] = []; } diff --git a/packages/jupyter-ai/src/components/chat-input.tsx b/packages/jupyter-ai/src/components/chat-input.tsx index 31127dd84..1c6a63730 100644 --- a/packages/jupyter-ai/src/components/chat-input.tsx +++ b/packages/jupyter-ai/src/components/chat-input.tsx @@ -13,7 +13,6 @@ import { import SendIcon from '@mui/icons-material/Send'; type ChatInputProps = { - loading: boolean; value: string; onChange: (newValue: string) => unknown; onSend: () => unknown; @@ -39,7 +38,7 @@ export function ChatInput(props: ChatInputProps): JSX.Element { size="large" color="primary" onClick={props.onSend} - disabled={props.loading || !props.value.trim().length} + disabled={!props.value.trim().length} > diff --git a/packages/jupyter-ai/src/components/chat.tsx b/packages/jupyter-ai/src/components/chat.tsx index 431dc84d8..fbcc708f2 100644 --- a/packages/jupyter-ai/src/components/chat.tsx +++ b/packages/jupyter-ai/src/components/chat.tsx @@ -1,8 +1,4 @@ -import React, { - useState - // useMemo, - // useEffect -} from 'react'; +import React, { useState, useEffect } from 'react'; import { Box } from '@mui/system'; @@ -15,105 +11,95 @@ import { useSelectionContext } from '../contexts/selection-context'; import { SelectionWatcher } from '../selection-watcher'; -// import { ChatHandler } from '../chat_handler'; +import { ChatHandler } from '../chat_handler'; type ChatMessageGroup = { sender: 'self' | 'ai' | string; messages: string[]; }; -function ChatBody(): JSX.Element { +type ChatBodyProps = { + chatHandler: ChatHandler; +}; + +function ChatBody({ chatHandler }: ChatBodyProps): JSX.Element { const [messageGroups, setMessageGroups] = useState([]); - const [loading, setLoading] = useState(false); const [includeSelection, setIncludeSelection] = useState(true); const [replaceSelection, setReplaceSelection] = useState(false); const [input, setInput] = useState(''); const [selection, replaceSelectionFn] = useSelectionContext(); - // TODO: connect to websockets. - // const chatHandler = useMemo(() => new ChatHandler(), []); - // - // /** - // * Effect: fetch history on initial render - // */ - // useEffect(() => { - // async function fetchHistory() { - // const history = await chatHandler.getHistory(); - // const messages = history.messages; - // if (!messages.length) { - // return; - // } - - // const newMessageGroups = messages.map( - // (message: AiService.ChatMessage): ChatMessageGroup => ({ - // sender: message.type === 'ai' ? 'ai' : 'self', - // messages: [message.data.content] - // }) - // ); - // setMessageGroups(newMessageGroups); - // } - - // fetchHistory(); - // }, [chatHandler]); - - // /** - // * Effect: listen to chat messages - // */ - // useEffect(() => { - // function handleChatEvents(message: AiService.ChatMessage) { - // setMessageGroups(messageGroups => [ - // ...messageGroups, - // { - // sender: message.type === 'ai' ? 'ai' : 'self', - // messages: [message.data.content] - // } - // ]); - // } - - // chatHandler.addListener(handleChatEvents); - - // return function cleanup() { - // chatHandler.removeListener(handleChatEvents); - // }; - // }, [chatHandler]); + /** + * Effect: fetch history on initial render + */ + useEffect(() => { + async function fetchHistory() { + const history = await chatHandler.getHistory(); + const messages = history.messages; + if (!messages.length) { + return; + } + + const newMessageGroups = messages.map( + (message: AiService.ChatMessage): ChatMessageGroup => ({ + sender: message.type === 'agent' ? 'ai' : 'self', + messages: [message.body] + }) + ); + setMessageGroups(newMessageGroups); + } + + fetchHistory(); + }, [chatHandler]); + + /** + * Effect: listen to chat messages + */ + useEffect(() => { + function handleChatEvents(message: AiService.Message) { + if (message.type === 'connection') { + return; + } + + setMessageGroups(messageGroups => [ + ...messageGroups, + { + sender: message.type === 'agent' ? 'ai' : 'self', + messages: [message.body] + } + ]); + } + chatHandler.addListener(handleChatEvents); + return function cleanup() { + chatHandler.removeListener(handleChatEvents); + }; + }, [chatHandler]); + + // no need to append to messageGroups imperatively here. all of that is + // handled by the listeners registered in the effect hooks above. const onSend = async () => { - setLoading(true); setInput(''); - console.log({ - includeSelection, - replaceSelection - }); - const newMessages = [input]; - if (includeSelection && selection) { - newMessages.push('```\n' + selection + '\n```'); - } - setMessageGroups(messageGroups => [ - ...messageGroups, - { sender: 'self', messages: newMessages } - ]); - - let response: AiService.ChatResponse; - - const prompt = input + (selection ? '\n--\n' + selection : ''); - try { - response = await AiService.sendChat({ prompt }); - } finally { - setLoading(false); - } + const prompt = + input + + (includeSelection && selection?.text ? '\n--\n' + selection.text : ''); + + // send message to backend + const messageId = await chatHandler.sendMessage({ prompt }); + + // await reply from agent + // no need to append to messageGroups state variable, since that's already + // handled in the effect hooks. + const reply = await chatHandler.replyFor(messageId); if (replaceSelection && selection) { const { cellId, ...selectionProps } = selection; replaceSelectionFn({ ...selectionProps, ...(cellId && { cellId }), - text: response.output + text: reply.body }); } - setMessageGroups(messageGroups => [ - ...messageGroups, - { sender: 'ai', messages: [response.output] } - ]); }; return ( @@ -148,7 +134,6 @@ function ChatBody(): JSX.Element { ))} - + ); diff --git a/packages/jupyter-ai/src/handler.ts b/packages/jupyter-ai/src/handler.ts index 5b9ed2088..7887d6275 100644 --- a/packages/jupyter-ai/src/handler.ts +++ b/packages/jupyter-ai/src/handler.ts @@ -64,20 +64,37 @@ export namespace AiService { prompt: string; }; - export type ChatResponse = { - output: string; + export type ChatClient = { + id: string; + initials: string; + name: string; + display_name: string; + color?: string; + avatar_url?: string; + }; + + export type AgentChatMessage = { + type: 'agent'; + id: string; + body: string; + reply_to: string; }; - export type ChatMessageData = { - content: string; - additional_kwargs: { [key: string]: any }; + export type HumanChatMessage = { + type: 'human'; + id: string; + body: string; + client: ChatClient; }; - export type ChatMessage = { - type: string; - data: ChatMessageData; + export type ConnectionMessage = { + type: 'connection'; + client_id: string; }; + export type ChatMessage = AgentChatMessage | HumanChatMessage; + export type Message = AgentChatMessage | HumanChatMessage | ConnectionMessage; + export type ChatHistory = { messages: ChatMessage[]; }; @@ -103,20 +120,6 @@ export namespace AiService { return data as IPromptResponse; } - export async function sendChat(request: ChatRequest): Promise { - let data; - - try { - data = await requestAPI('chat', { - method: 'POST', - body: JSON.stringify(request) - }); - } catch (e) { - return Promise.reject(e); - } - return data as IPromptResponse; - } - export type ListTasksEntry = { id: string; name: string; diff --git a/packages/jupyter-ai/src/index.ts b/packages/jupyter-ai/src/index.ts index e78c89e96..2e6e764c9 100644 --- a/packages/jupyter-ai/src/index.ts +++ b/packages/jupyter-ai/src/index.ts @@ -17,6 +17,7 @@ import { psychologyIcon } from './icons'; import { getTextSelection } from './utils'; import { buildChatSidebar } from './widgets/chat-sidebar'; import { SelectionWatcher } from './selection-watcher'; +import { ChatHandler } from './chat_handler'; export enum NotebookTasks { GenerateCode = 'generate-code-in-cells-below', @@ -47,7 +48,7 @@ const plugin: JupyterFrontEndPlugin = { id: 'jupyter_ai:plugin', autoStart: true, requires: [INotebookTracker, IEditorTracker], - activate: ( + activate: async ( app: JupyterFrontEnd, notebookTracker: INotebookTracker, editorTracker: IEditorTracker @@ -81,10 +82,16 @@ const plugin: JupyterFrontEndPlugin = { */ const selectionWatcher = new SelectionWatcher(shell); + /** + * Initialize chat handler, open WS connection + */ + const chatHandler = new ChatHandler(); + await chatHandler.initialize(); + /** * Add Chat widget to right sidebar */ - shell.add(buildChatSidebar(selectionWatcher), 'right'); + shell.add(buildChatSidebar(selectionWatcher, chatHandler), 'right'); /** * Register inserters diff --git a/packages/jupyter-ai/src/widgets/chat-sidebar.tsx b/packages/jupyter-ai/src/widgets/chat-sidebar.tsx index 3551cd776..101ffd906 100644 --- a/packages/jupyter-ai/src/widgets/chat-sidebar.tsx +++ b/packages/jupyter-ai/src/widgets/chat-sidebar.tsx @@ -4,10 +4,14 @@ import { ReactWidget } from '@jupyterlab/apputils'; import { Chat } from '../components/chat'; import { psychologyIcon } from '../icons'; import { SelectionWatcher } from '../selection-watcher'; +import { ChatHandler } from '../chat_handler'; -export function buildChatSidebar(selectionWatcher: SelectionWatcher) { +export function buildChatSidebar( + selectionWatcher: SelectionWatcher, + chatHandler: ChatHandler +) { const ChatWidget = ReactWidget.create( - + ); ChatWidget.id = 'jupyter-ai::chat'; ChatWidget.title.icon = psychologyIcon;