From e7b562d762d0a586b206072879e14c95b40e22b9 Mon Sep 17 00:00:00 2001 From: Nicolas Brichet Date: Fri, 13 Sep 2024 17:09:06 +0200 Subject: [PATCH 1/2] Allow for stream messages --- .../jupyter_ai/chat_handlers/default.py | 43 ++++++++++--------- packages/jupyter-ai/jupyter_ai/extension.py | 19 ++++++-- 2 files changed, 38 insertions(+), 24 deletions(-) diff --git a/packages/jupyter-ai/jupyter_ai/chat_handlers/default.py b/packages/jupyter-ai/jupyter_ai/chat_handlers/default.py index 480fd4565..7e90fa2f0 100644 --- a/packages/jupyter-ai/jupyter_ai/chat_handlers/default.py +++ b/packages/jupyter-ai/jupyter_ai/chat_handlers/default.py @@ -11,6 +11,7 @@ from langchain_core.messages import AIMessageChunk from langchain_core.runnables import ConfigurableFieldSpec from langchain_core.runnables.history import RunnableWithMessageHistory +from langchain_core.outputs.generation import GenerationChunk try: from jupyterlab_collaborative_chat.ychat import YChat @@ -62,27 +63,29 @@ def create_llm_chain( ) self.llm_chain = runnable - def _start_stream(self, human_msg: HumanChatMessage) -> str: + def _start_stream(self, human_msg: HumanChatMessage, chat: YChat | None) -> str: """ Sends an `agent-stream` message to indicate the start of a response stream. Returns the ID of the message, denoted as the `stream_id`. """ - stream_id = uuid4().hex - stream_msg = AgentStreamMessage( - id=stream_id, - time=time.time(), - body="", - reply_to=human_msg.id, - persona=self.persona, - complete=False, - ) - - for handler in self._root_chat_handlers.values(): - if not handler: - continue - - handler.broadcast_message(stream_msg) - break + if chat is not None: + stream_id = self.write_message(chat, "") + else: + stream_id = uuid4().hex + stream_msg = AgentStreamMessage( + id=stream_id, + time=time.time(), + body="", + reply_to=human_msg.id, + persona=self.persona, + complete=False, + ) + for handler in self._root_chat_handlers.values(): + if not handler: + continue + + handler.broadcast_message(stream_msg) + break return stream_id @@ -92,7 +95,7 @@ def _send_stream_chunk(self, stream_id: str, content: str, chat: YChat | None, c appended to an existing `agent-stream` message with ID `stream_id`. """ if chat is not None: - self.write_message(chat, content) + self.write_message(chat, content, stream_id) else: stream_chunk_msg = AgentStreamChunkMessage( id=stream_id, content=content, stream_complete=complete @@ -121,11 +124,11 @@ async def process_message(self, message: HumanChatMessage, chat: YChat | None): # when receiving the first chunk, close the pending message and # start the stream. self.close_pending(pending_message) - stream_id = self._start_stream(human_msg=message) + stream_id = self._start_stream(message, chat) received_first_chunk = True if isinstance(chunk, AIMessageChunk): - self._send_stream_chunk(stream_id, chunk.content) + self._send_stream_chunk(stream_id, chunk.content, chat) elif isinstance(chunk, str): self._send_stream_chunk(stream_id, chunk, chat) else: diff --git a/packages/jupyter-ai/jupyter_ai/extension.py b/packages/jupyter-ai/jupyter_ai/extension.py index 128263726..201291019 100644 --- a/packages/jupyter-ai/jupyter_ai/extension.py +++ b/packages/jupyter-ai/jupyter_ai/extension.py @@ -229,6 +229,10 @@ def initialize(self): listener=self.connect_chat ) + # Keep the message indexes to avoid extra computation looking for a message when + # updating it. + self.messages_indexes = {} + async def connect_chat(self, logger: EventLogger, schema_id: str, data: dict) -> None: if data["room"].startswith("text:chat:") \ and data["action"] == "initialize"\ @@ -299,7 +303,7 @@ async def _route(self, message: HumanChatMessage, chat: YChat): command_readable = "Default" if command == "default" else command self.log.info(f"{command_readable} chat handler resolved in {latency_ms} ms.") - def write_message(self, chat: YChat, body: str) -> None: + def write_message(self, chat: YChat, body: str, id: str | None = None) -> str: BOT["avatar_url"] = url_path_join( self.settings.get("base_url", "/"), "api/ai/static/jupyternaut.svg" @@ -310,14 +314,21 @@ def write_message(self, chat: YChat, body: str) -> None: else: BOT["username"] = bot["username"] - chat.add_message({ + index = self.messages_indexes[id] if id else None + id = id if id else str(uuid.uuid4()) + new_index = chat.set_message({ "type": "msg", "body": body, - "id": str(uuid.uuid4()), + "id": id if id else str(uuid.uuid4()), "time": time.time(), "sender": BOT["username"], "raw_time": False - }) + }, index, True) + + if new_index != index: + self.messages_indexes[id] = new_index + + return id def initialize_settings(self): start = time.time() From c6429d445f33c1f4c7db841faf9712d4e6506eb0 Mon Sep 17 00:00:00 2001 From: Nicolas Brichet Date: Fri, 13 Sep 2024 17:50:28 +0200 Subject: [PATCH 2/2] update jupyter collaborative chat dependency --- packages/jupyter-ai/package.json | 2 +- packages/jupyter-ai/pyproject.toml | 3 +-- yarn.lock | 10 +++++----- 3 files changed, 7 insertions(+), 8 deletions(-) diff --git a/packages/jupyter-ai/package.json b/packages/jupyter-ai/package.json index a1f8029da..cc404aaa3 100644 --- a/packages/jupyter-ai/package.json +++ b/packages/jupyter-ai/package.json @@ -61,7 +61,7 @@ "dependencies": { "@emotion/react": "^11.10.5", "@emotion/styled": "^11.10.5", - "@jupyter/chat": "^0.3.1", + "@jupyter/chat": "^0.4.0", "@jupyter/collaboration": "^1", "@jupyterlab/application": "^4.2.0", "@jupyterlab/apputils": "^4.2.0", diff --git a/packages/jupyter-ai/pyproject.toml b/packages/jupyter-ai/pyproject.toml index 9b1d3d996..2353343e3 100644 --- a/packages/jupyter-ai/pyproject.toml +++ b/packages/jupyter-ai/pyproject.toml @@ -33,6 +33,7 @@ dependencies = [ "typing_extensions>=4.5.0", "traitlets>=5.0", "deepmerge>=1.0", + "jupyterlab-collaborative-chat", ] dynamic = ["version", "description", "authors", "urls", "keywords"] @@ -56,8 +57,6 @@ dev = ["jupyter_ai_magics[dev]"] all = ["jupyter_ai_magics[all]", "pypdf", "arxiv"] -collaborative = ["jupyterlab-collaborative-chat"] - [tool.hatch.version] source = "nodejs" diff --git a/yarn.lock b/yarn.lock index 49d9b3304..5d45cf685 100644 --- a/yarn.lock +++ b/yarn.lock @@ -2220,7 +2220,7 @@ __metadata: "@babel/preset-env": ^7.0.0 "@emotion/react": ^11.10.5 "@emotion/styled": ^11.10.5 - "@jupyter/chat": ^0.3.1 + "@jupyter/chat": ^0.4.0 "@jupyter/collaboration": ^1 "@jupyterlab/application": ^4.2.0 "@jupyterlab/apputils": ^4.2.0 @@ -2286,9 +2286,9 @@ __metadata: languageName: unknown linkType: soft -"@jupyter/chat@npm:^0.3.1": - version: 0.3.1 - resolution: "@jupyter/chat@npm:0.3.1" +"@jupyter/chat@npm:^0.4.0": + version: 0.4.0 + resolution: "@jupyter/chat@npm:0.4.0" dependencies: "@emotion/react": ^11.10.5 "@emotion/styled": ^11.10.5 @@ -2306,7 +2306,7 @@ __metadata: clsx: ^2.1.0 react: ^18.2.0 react-dom: ^18.2.0 - checksum: 92d1f6d6d3083be2a8c2309fc37eb3bb39bc184e1c41d18f79d51f87e54d7badba1495c5943916e7a21fc65414d6100860b044a14554add8192b545fca2748dd + checksum: 6e309c8e70cf480103eb26f3109da417c58d2e6844d5e56e63feabf71926f9dba6f9bc85caff765dfc574a8fd7ed803a8c03e5d812c28568dcf6ec918bbd2e66 languageName: node linkType: hard