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,
@@ -98,6 +100,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) -> List[ChatMessage]:
+        return self.settings["llm_chat_memory"]
+
     @property
     def loop(self) -> AbstractEventLoop:
         return self.settings["jai_event_loop"]
@@ -246,17 +252,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 +324,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):
     """
@@ -30,14 +32,31 @@ 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()
 
     async def aclear(self) -> None:
@@ -75,6 +94,7 @@ def messages(self) -> List[BaseMessage]:
     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:
+            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 empty strig, clears all.
+    """
+
+
 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, clears all.
+    """
 
 
 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,29 @@
+import React from 'react';
+import { IconButton, SxProps } from '@mui/material';
+import { Close } from '@mui/icons-material';
+
+import { AiService } from '../../handler';
+import { ChatHandler } from '../../chat_handler';
+
+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 (
+    <IconButton
+      onClick={() => props.chatHandler.sendMessage(request)}
+      sx={props.sx}
+    >
+      <Close />
+    </IconButton>
+  );
+}
+
+export default ChatMessageDelete;

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

@@ -1,8 +1,9 @@
 import React, { useState, useEffect } from 'react';
 import { Box } from '@mui/system';
-import { Button, IconButton, Stack } from '@mui/material';
+import { Button, IconButton, Stack, Tooltip } from '@mui/material';
 import SettingsIcon from '@mui/icons-material/Settings';
 import ArrowBackIcon from '@mui/icons-material/ArrowBack';
+import AddIcon from '@mui/icons-material/Add';
 import type { Awareness } from 'y-protocols/awareness';
 import type { IThemeManager } from '@jupyterlab/apputils';
 import { IRenderMimeRegistry } from '@jupyterlab/rendermime';
@@ -143,10 +144,11 @@ function ChatBody({
       <ScrollContainer sx={{ flexGrow: 1 }}>
         <ChatMessages
           messages={messages}
+          chatHandler={chatHandler}
           rmRegistry={renderMimeRegistry}
           messageFooter={messageFooter}
         />
-        <PendingMessages messages={pendingMessages} />
+        <PendingMessages messages={pendingMessages} chatHandler={chatHandler} />
       </ScrollContainer>
       <ChatInput
         chatHandler={chatHandler}
@@ -216,9 +218,20 @@ export function Chat(props: ChatProps): JSX.Element {
                   <Box />
                 )}
                 {view === ChatView.Chat ? (
-                  <IconButton onClick={() => setView(ChatView.Settings)}>
-                    <SettingsIcon />
-                  </IconButton>
+                  <Box sx={{ display: 'flex' }}>
+                    <Tooltip title="New Chat">
+                      <IconButton
+                        onClick={() =>
+                          props.chatHandler.sendMessage({ type: 'clear' })
+                        }
+                      >
+                        <AddIcon />
+                      </IconButton>
+                    </Tooltip>
+                    <IconButton onClick={() => setView(ChatView.Settings)}>
+                      <SettingsIcon />
+                    </IconButton>
+                  </Box>
                 ) : (
                   <Box />
                 )}

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

@@ -3,9 +3,11 @@ import React, { useState, useEffect } from 'react';
 import { Box, Typography } from '@mui/material';
 import { AiService } from '../handler';
 import { ChatMessageHeader } from './chat-messages';
+import { ChatHandler } from '../chat_handler';
 
 type PendingMessagesProps = {
   messages: AiService.PendingMessage[];
+  chatHandler: ChatHandler;
 };
 
 type PendingMessageElementProps = {
@@ -85,6 +87,7 @@ export function PendingMessages(
     >
       <ChatMessageHeader
         message={agentMessage}
+        chatHandler={props.chatHandler}
         timestamp={timestamp}
         sx={{
           marginBottom: 4

packages/jupyter-ai/src/handler.ts

@@ -83,6 +83,11 @@ export namespace AiService {
     selection?: Selection;
   };
 
+  export type ClearRequest = {
+    type: 'clear';
+    target?: string;
+  };
+
   export type Collaborator = {
     username: string;
     initials: string;
@@ -138,6 +143,7 @@ export namespace AiService {
 
   export type ClearMessage = {
     type: 'clear';
+    target?: string;
   };
 
   export type PendingMessage = {
@@ -166,6 +172,8 @@ export namespace AiService {
     stream_complete: boolean;
   };
 
+  export type Request = ChatRequest | ClearRequest;
+
   export type ChatMessage =
     | AgentChatMessage
     | HumanChatMessage