Add ability to delete messages + start new chat session

packages/jupyter-ai/jupyter_ai/chat_handlers/clear.py

@@ -23,9 +23,4 @@ async def process_message(self, _):
                 continue
 
             handler.broadcast_message(ClearMessage())
-            self._chat_history.clear()
-            self.llm_chat_memory.clear()
             break
-
-        # re-send help message
-        self.send_help_message()

packages/jupyter-ai/jupyter_ai/handlers.py

@@ -24,6 +24,8 @@
     ChatMessage,
     ChatRequest,
     ChatUser,
+    ClearMessage,
+    ClearRequest,
     ClosePendingMessage,
     ConnectionMessage,
     HumanChatMessage,
@@ -40,6 +42,8 @@
     from jupyter_ai_magics.embedding_providers import BaseEmbeddingsProvider
     from jupyter_ai_magics.providers import BaseProvider
 
+    from .history import BoundChatHistory
+
 
 class ChatHistoryHandler(BaseAPIHandler):
     """Handler to return message history"""
@@ -98,6 +102,10 @@ def chat_history(self) -> List[ChatMessage]:
     def chat_history(self, new_history):
         self.settings["chat_history"] = new_history
 
+    @property
+    def llm_chat_memory(self) -> "BoundChatHistory":
+        return self.settings["llm_chat_memory"]
+
     @property
     def loop(self) -> AbstractEventLoop:
         return self.settings["jai_event_loop"]
@@ -246,17 +254,33 @@ def broadcast_message(self, message: Message):
             self.pending_messages = list(
                 filter(lambda m: m.id != message.id, self.pending_messages)
             )
+        elif isinstance(message, ClearMessage):
+            if message.target:
+                self._clear_chat_history_at(message.target)
+            else:
+                self.chat_history.clear()
+                self.pending_messages.clear()
+                self.llm_chat_memory.clear()
+                self.settings["jai_chat_handlers"]["default"].send_help_message()
 
     async def on_message(self, message):
         self.log.debug("Message received: %s", message)
 
         try:
             message = json.loads(message)
-            chat_request = ChatRequest(**message)
+            if message.get("type") == "clear":
+                request = ClearRequest(**message)
+            else:
+                request = ChatRequest(**message)
         except ValidationError as e:
             self.log.error(e)
             return
 
+        if isinstance(request, ClearRequest):
+            self.broadcast_message(ClearMessage(target=request.target))
+            return
+
+        chat_request = request
         message_body = chat_request.prompt
         if chat_request.selection:
             message_body += f"\n\n```\n{chat_request.selection.source}\n```\n"
@@ -302,6 +326,22 @@ 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):
+        """Clears the chat history at a specific message ID."""
+        target_msg = None
+        for msg in self.chat_history:
+            if msg.id == msg_id:
+                target_msg = msg
+
+        if target_msg is not None:
+            self.chat_history[:] = [
+                msg for msg in self.chat_history if msg.time < target_msg.time
+            ]
+            self.pending_messages[:] = [
+                msg for msg in self.pending_messages if msg.time < target_msg.time
+            ]
+            self.llm_chat_memory.clear(target_msg.time)
+
     def on_close(self):
         self.log.debug("Disconnecting client with user %s", self.client_id)
 

packages/jupyter-ai/jupyter_ai/history.py

@@ -7,6 +7,8 @@
 
 from .models import HumanChatMessage
 
+MESSAGE_TIME_KEY = "_jupyter_ai_msg_time"
+
 
 class BoundedChatHistory(BaseChatMessageHistory, BaseModel):
     """
@@ -19,6 +21,7 @@ class BoundedChatHistory(BaseChatMessageHistory, BaseModel):
 
     k: int
     clear_time: float = 0.0
+    clear_after: float = 0.0
     _all_messages: List[BaseMessage] = PrivateAttr(default_factory=list)
 
     @property
@@ -30,15 +33,33 @@ async def aget_messages(self) -> List[BaseMessage]:
 
     def add_message(self, message: BaseMessage) -> None:
         """Add a self-created message to the store"""
+        # Adds a timestamp to the message as a fallback if message was not
+        # added not using WrappedBoundedChatHistory.
+        # In such a case, it possible that this message may be cleared even if
+        # the target clear message is after this one.
+        # This will occur if the current time is greater than the last_human_msg time of
+        # a future message that was added using WrappedBoundedChatHistory.
+        message.additional_kwargs[MESSAGE_TIME_KEY] = message.additional_kwargs.get(
+            MESSAGE_TIME_KEY, time.time()
+        )
         self._all_messages.append(message)
 
     async def aadd_messages(self, messages: Sequence[BaseMessage]) -> None:
         """Add messages to the store"""
         self.add_messages(messages)
 
-    def clear(self) -> None:
-        self._all_messages = []
+    def clear(self, after: float = 0.0) -> None:
+        """Clear all messages after the given time"""
+        if after:
+            self._all_messages = [
+                m
+                for m in self._all_messages
+                if m.additional_kwargs[MESSAGE_TIME_KEY] < after
+            ]
+        else:
+            self._all_messages = []
         self.clear_time = time.time()
+        self.clear_after = after
 
     async def aclear(self) -> None:
         self.clear()
@@ -73,8 +94,14 @@ def messages(self) -> List[BaseMessage]:
         return self.history.messages
 
     def add_message(self, message: BaseMessage) -> None:
-        """Prevent adding messages to the store if clear was triggered."""
-        if self.last_human_msg.time > self.history.clear_time:
+        # prevent adding pending messages to the store if clear was triggered.
+        # if partial clearing, allow adding pending messages that were not cleared.
+        if (
+            self.last_human_msg.time
+            > self.history.clear_time | self.last_human_msg.time
+            < self.history.clear_after
+        ):
+            message.additional_kwargs[MESSAGE_TIME_KEY] = self.last_human_msg.time
             self.history.add_message(message)
 
     async def aadd_messages(self, messages: Sequence[BaseMessage]) -> None:

packages/jupyter-ai/jupyter_ai/models.py

@@ -39,6 +39,15 @@ class ChatRequest(BaseModel):
     selection: Optional[Selection]
 
 
+class ClearRequest(BaseModel):
+    type: Literal["clear"]
+    target: Optional[str]
+    """
+    Message ID of the ChatMessage to clear at and all messages after.
+    If not provided, this requests the backend to clear all messages.
+    """
+
+
 class ChatUser(BaseModel):
     # User ID assigned by IdentityProvider.
     username: str
@@ -105,6 +114,11 @@ class HumanChatMessage(BaseModel):
 
 class ClearMessage(BaseModel):
     type: Literal["clear"] = "clear"
+    target: Optional[str] = None
+    """
+    Message ID of the ChatMessage to clear at and all messages after.
+    If not provided, this instructs the frontend to clear all messages.
+    """
 
 
 class PendingMessage(BaseModel):

packages/jupyter-ai/src/chat_handler.ts

@@ -39,7 +39,7 @@ export class ChatHandler implements IDisposable {
    * Sends a message across the WebSocket. Promise resolves to the message ID
    * when the server sends the same message back, acknowledging receipt.
    */
-  public sendMessage(message: AiService.ChatRequest): Promise<string> {
+  public sendMessage(message: AiService.Request): Promise<string> {
     return new Promise(resolve => {
       this._socket?.send(JSON.stringify(message));
       this._sendResolverQueue.push(resolve);
@@ -132,8 +132,22 @@ export class ChatHandler implements IDisposable {
       case 'connection':
         break;
       case 'clear':
-        this._messages = [];
-        this._pendingMessages = [];
+        if (newMessage.target) {
+          const targetMsg = this._messages.find(
+            m => m.id === newMessage.target
+          );
+          if (targetMsg) {
+            this._messages = this._messages.filter(
+              msg => msg.time < targetMsg.time
+            );
+            this._pendingMessages = this._pendingMessages.filter(
+              msg => msg.time < targetMsg.time
+            );
+          }
+        } else {
+          this._messages = [];
+          this._pendingMessages = [];
+        }
         break;
       case 'pending':
         this._pendingMessages = [...this._pendingMessages, newMessage];

packages/jupyter-ai/src/components/chat-messages.tsx

@@ -10,16 +10,20 @@ import { AiService } from '../handler';
 import { RendermimeMarkdown } from './rendermime-markdown';
 import { useCollaboratorsContext } from '../contexts/collaborators-context';
 import { ChatMessageMenu } from './chat-messages/chat-message-menu';
+import { ChatMessageDelete } from './chat-messages/chat-message-delete';
+import { ChatHandler } from '../chat_handler';
 import { IJaiMessageFooter } from '../tokens';
 
 type ChatMessagesProps = {
   rmRegistry: IRenderMimeRegistry;
   messages: AiService.ChatMessage[];
+  chatHandler: ChatHandler;
   messageFooter: IJaiMessageFooter | null;
 };
 
 type ChatMessageHeaderProps = {
   message: AiService.ChatMessage;
+  chatHandler: ChatHandler;
   timestamp: string;
   sx?: SxProps<Theme>;
 };
@@ -113,6 +117,7 @@ export function ChatMessageHeader(props: ChatMessageHeaderProps): JSX.Element {
   const shouldShowMenu =
     props.message.type === 'agent' ||
     (props.message.type === 'agent-stream' && props.message.complete);
+  const shouldShowDelete = props.message.type === 'human';
 
   return (
     <Box
@@ -154,6 +159,13 @@ export function ChatMessageHeader(props: ChatMessageHeaderProps): JSX.Element {
               sx={{ marginLeft: '4px', marginRight: '-8px' }}
             />
           )}
+          {shouldShowDelete && (
+            <ChatMessageDelete
+              message={props.message}
+              chatHandler={props.chatHandler}
+              sx={{ marginLeft: '4px', marginRight: '-8px' }}
+            />
+          )}
         </Box>
       </Box>
     </Box>
@@ -208,6 +220,7 @@ export function ChatMessages(props: ChatMessagesProps): JSX.Element {
             <ChatMessageHeader
               message={message}
               timestamp={timestamps[message.id]}
+              chatHandler={props.chatHandler}
               sx={{ marginBottom: 3 }}
             />
             <RendermimeMarkdown

packages/jupyter-ai/src/components/chat-messages/chat-message-delete.tsx

@@ -0,0 +1,31 @@
+import React from 'react';
+import { SxProps } from '@mui/material';
+import { Close } from '@mui/icons-material';
+
+import { AiService } from '../../handler';
+import { ChatHandler } from '../../chat_handler';
+import { TooltippedIconButton } from '../mui-extras/tooltipped-icon-button';
+
+type DeleteButtonProps = {
+  message: AiService.ChatMessage;
+  chatHandler: ChatHandler;
+  sx?: SxProps;
+};
+
+export function ChatMessageDelete(props: DeleteButtonProps): JSX.Element {
+  const request: AiService.ClearRequest = {
+    type: 'clear',
+    target: props.message.id
+  };
+  return (
+    <TooltippedIconButton
+      onClick={() => props.chatHandler.sendMessage(request)}
+      sx={props.sx}
+      tooltip="Delete this and all future messages"
+    >
+      <Close />
+    </TooltippedIconButton>
+  );
+}
+
+export default ChatMessageDelete;