From f4d996e66e049080434778d4ad79efa7bcf93e31 Mon Sep 17 00:00:00 2001 From: michael Date: Sun, 8 Sep 2024 11:46:30 +0800 Subject: [PATCH] support clearing all subsequent exchanges --- packages/jupyter-ai/jupyter_ai/handlers.py | 31 ++++++++++++++----- packages/jupyter-ai/jupyter_ai/history.py | 9 +++--- packages/jupyter-ai/jupyter_ai/models.py | 9 ++++-- packages/jupyter-ai/src/chat_handler.ts | 9 +++--- .../chat-messages/chat-message-delete.tsx | 3 +- packages/jupyter-ai/src/handler.ts | 3 +- 6 files changed, 43 insertions(+), 21 deletions(-) diff --git a/packages/jupyter-ai/jupyter_ai/handlers.py b/packages/jupyter-ai/jupyter_ai/handlers.py index cc34369c5..a614e3e84 100644 --- a/packages/jupyter-ai/jupyter_ai/handlers.py +++ b/packages/jupyter-ai/jupyter_ai/handlers.py @@ -255,8 +255,8 @@ def broadcast_message(self, message: Message): filter(lambda m: m.id != message.id, self.pending_messages) ) elif isinstance(message, ClearMessage): - if message.target: - self._clear_chat_history_at(message.target) + if message.targets: + self._clear_chat_history_at(message.targets) else: self.chat_history.clear() self.pending_messages.clear() @@ -277,7 +277,22 @@ async def on_message(self, message): return if isinstance(request, ClearRequest): - self.broadcast_message(ClearMessage(target=request.target)) + if not request.target: + targets = None + elif request.after: + target_msg = None + for msg in self.chat_history: + if msg.id == request.target: + target_msg = msg + if target_msg: + targets = [ + msg.id + for msg in self.chat_history + if msg.time >= target_msg.time and msg.type == "human" + ] + else: + targets = [request.target] + self.broadcast_message(ClearMessage(targets=targets)) return chat_request = request @@ -326,19 +341,19 @@ async def _route(self, message): command_readable = "Default" if command == "default" else command self.log.info(f"{command_readable} chat handler resolved in {latency_ms} ms.") - def _clear_chat_history_at(self, msg_id: str): + def _clear_chat_history_at(self, msg_ids: List[str]): """ - Clears a conversation exchange given a human message ID `msg_id`. + Clears conversation exchanges associated with list of human message IDs. """ self.chat_history[:] = [ msg for msg in self.chat_history - if msg.id != msg_id and getattr(msg, "reply_to", None) != msg_id + if msg.id not in msg_ids and getattr(msg, "reply_to", None) not in msg_ids ] self.pending_messages[:] = [ - msg for msg in self.pending_messages if msg.reply_to != msg_id + msg for msg in self.pending_messages if msg.reply_to not in msg_ids ] - self.llm_chat_memory.clear(msg_id) + self.llm_chat_memory.clear(msg_ids) def on_close(self): self.log.debug("Disconnecting client with user %s", self.client_id) diff --git a/packages/jupyter-ai/jupyter_ai/history.py b/packages/jupyter-ai/jupyter_ai/history.py index f96c80fe2..9e1064194 100644 --- a/packages/jupyter-ai/jupyter_ai/history.py +++ b/packages/jupyter-ai/jupyter_ai/history.py @@ -45,17 +45,17 @@ async def aadd_messages(self, messages: Sequence[BaseMessage]) -> None: """Add messages to the store""" self.add_messages(messages) - def clear(self, human_msg_id: Optional[str] = None) -> None: + def clear(self, human_msg_ids: Optional[List[str]] = None) -> None: """Clears conversation exchanges. If `human_msg_id` is provided, only clears the respective human message and its reply. Otherwise, clears all messages.""" - if human_msg_id: + if human_msg_ids: self._all_messages = [ m for m in self._all_messages - if m.additional_kwargs[HUMAN_MSG_ID_KEY] != human_msg_id + if m.additional_kwargs[HUMAN_MSG_ID_KEY] not in human_msg_ids ] - self.cleared_msgs.add(human_msg_id) + self.cleared_msgs.update(human_msg_ids) else: self._all_messages = [] self.cleared_msgs = set() @@ -95,7 +95,6 @@ def messages(self) -> List[BaseMessage]: def add_message(self, message: BaseMessage) -> None: # prevent adding pending messages to the store if clear was triggered. - # if targeted clearing, prevent adding target message if still pending. if ( self.last_human_msg.time > self.history.clear_time and self.last_human_msg.id not in self.history.cleared_msgs diff --git a/packages/jupyter-ai/jupyter_ai/models.py b/packages/jupyter-ai/jupyter_ai/models.py index 7e663d55d..f9098a12a 100644 --- a/packages/jupyter-ai/jupyter_ai/models.py +++ b/packages/jupyter-ai/jupyter_ai/models.py @@ -47,6 +47,11 @@ class ClearRequest(BaseModel): If not provided, this requests the backend to clear all messages. """ + after: Optional[bool] + """ + Whether to clear target and all subsequent exchanges. + """ + class ChatUser(BaseModel): # User ID assigned by IdentityProvider. @@ -114,9 +119,9 @@ class HumanChatMessage(BaseModel): class ClearMessage(BaseModel): type: Literal["clear"] = "clear" - target: Optional[str] = None + targets: Optional[List[str]] = None """ - Message ID of the HumanChatMessage to delete an exchange at. + Message IDs of the HumanChatMessage to delete an exchange at. If not provided, this instructs the frontend to clear all messages. """ diff --git a/packages/jupyter-ai/src/chat_handler.ts b/packages/jupyter-ai/src/chat_handler.ts index 2ff8b108c..76c93a851 100644 --- a/packages/jupyter-ai/src/chat_handler.ts +++ b/packages/jupyter-ai/src/chat_handler.ts @@ -132,14 +132,15 @@ export class ChatHandler implements IDisposable { case 'connection': break; case 'clear': - if (newMessage.target) { + if (newMessage.targets) { + const targets = newMessage.targets; this._messages = this._messages.filter( msg => - msg.id !== newMessage.target && - !('reply_to' in msg && msg.reply_to === newMessage.target) + !targets.includes(msg.id) && + !('reply_to' in msg && targets.includes(msg.reply_to)) ); this._pendingMessages = this._pendingMessages.filter( - msg => msg.reply_to !== newMessage.target + msg => !targets.includes(msg.reply_to) ); } else { this._messages = []; diff --git a/packages/jupyter-ai/src/components/chat-messages/chat-message-delete.tsx b/packages/jupyter-ai/src/components/chat-messages/chat-message-delete.tsx index d6fc691bd..b91e15b93 100644 --- a/packages/jupyter-ai/src/components/chat-messages/chat-message-delete.tsx +++ b/packages/jupyter-ai/src/components/chat-messages/chat-message-delete.tsx @@ -15,7 +15,8 @@ type DeleteButtonProps = { export function ChatMessageDelete(props: DeleteButtonProps): JSX.Element { const request: AiService.ClearRequest = { type: 'clear', - target: props.message.id + target: props.message.id, + after: false }; return (