Implement message deletion

heymans/heymans.py

@@ -54,7 +54,8 @@ def __init__(self, user_id: str, persistent: bool = False,
         self.answer_tools = [getattr(tools, t)(self) for t in answer_tools]
         self.tools = self.answer_tools
 
-    def send_user_message(self, message: str) -> GeneratorType:
+    def send_user_message(self, message: str,
+                          message_id: str=None) -> GeneratorType:
         """The main function that takes a user message and returns one or 
         replies. This is a generator function where each yield gives a tuple.
         
@@ -68,7 +69,7 @@ def send_user_message(self, message: str) -> GeneratorType:
             yield config.max_tokens_per_hour_exceeded_message, \
                     self.messages.metadata()
             return
-        self.messages.append('user', message)
+        self.messages.append('user', message, message_id)
         if self._search_first:
             for reply, metadata in self._search(message):
                 yield reply, metadata

heymans/messages.py

@@ -33,28 +33,54 @@ def init_conversation(self):
         metadata = self.metadata()
         metadata['search_model'] = 'Welcome message'
         self._conversation_title = config.default_conversation_title
-        self._message_history = [('assistant', self.welcome_message(), 
-                                  metadata)]
+        self._message_history = [['assistant', self.welcome_message(), 
+                                  metadata]]
         self._condensed_message_history = [
-            (role, content) for role, content, metadata
+            [role, content] for role, content, metadata
             in self._message_history[:]]        
 
-    def metadata(self):
-        return {'timestamp': utils.current_datetime(),
+    def metadata(self, message_id=None):
+        return {'message_id': str(uuid.uuid4()) 
+                    if message_id is None else message_id,
+                'timestamp': utils.current_datetime(),
                 'sources': self._heymans.documentation.to_json(),
                 'search_model': config.search_model,
                 'condense_model': config.condense_model,
                 'answer_model': config.answer_model}
 
-    def append(self, role, message):
-        metadata = self.metadata()
-        self._message_history.append((role, message, metadata))
-        self._condensed_message_history.append((role, message))
+    def append(self, role, message, message_id=None):
+        metadata = self.metadata(message_id=message_id)
+        self._message_history.append([role, message, metadata])
+        self._condensed_message_history.append([role, message])
         self._condense_message_history()
         if self._persistent:
             self.save()
         return metadata
 
+    def delete(self, message_id):
+        message_to_remove = None
+        condensed_message_to_remove = None
+        for role, message, metadata in self._message_history:
+            if metadata['message_id'] == message_id:
+                message_to_remove = [role, message, metadata]
+                condensed_message_to_remove = [role, message]
+                break
+        else:
+            logger.info(f'message not found for deletion: {message_id}')
+            return
+        if message_to_remove:
+            logger.info(f'deleting message: {message_id}')
+            self._message_history.remove(message_to_remove)
+        if condensed_message_to_remove:
+            try:
+                self._condensed_message_history.remove(
+                    condensed_message_to_remove)
+            except ValueError:
+                logger.error(
+                    f'Could not find condensed message to remove for {message_id}')
+        if self._persistent:
+            self.save()
+
     def prompt(self, system_prompt=None):
         """The prompt consists of the system prompt followed by a sequence of
         AI and user messages. Transient messages are special messages that are
@@ -171,11 +197,18 @@ def load(self):
         if not conversation['message_history']:
             self.init_conversation()
             return
+        modified = False
+        for _, _, metadata in conversation['message_history']:
+            if 'message_id' not in metadata:
+                metadata['message_id'] = str(uuid.uuid4())
+                modified = True
         self._conversation_title = conversation['title']
         self._message_history = conversation['message_history']
         self._condensed_text = conversation['condensed_text']
         self._condensed_message_history = \
             conversation['condensed_message_history']
+        if self._persistent and modified:
+            self.save()
 
     def save(self):
         self._update_title()

heymans/routes/api.py

@@ -19,6 +19,7 @@ def api_chat_start():
     data = request.json
     session['user_message'] = data.get('message', '')
     session['search_first'] = data.get('search_first', True)
+    session['message_id'] = data.get('message_id', None)
     heymans = get_heymans()
     redis_client.delete(f'stream_cancel_{heymans.user_id}')
     return '{}'
@@ -37,10 +38,11 @@ def api_chat_cancel_stream():
 @login_required
 def api_chat_stream():
     message = session['user_message']
+    message_id = session['message_id']
     heymans = get_heymans()
     logger.info(f'starting stream for {heymans.user_id}')
     def generate():
-        for reply, metadata in heymans.send_user_message(message):
+        for reply, metadata in heymans.send_user_message(message, message_id):
             if isinstance(reply, dict):
                 reply = json.dumps(reply)
             else:
@@ -156,3 +158,11 @@ def get_attachment(attachment_id):
         return response
     return jsonify(success=False,
                    message="Attachment not found or access denied"), 404
+
+
+@api_blueprint.route('/message/delete/<message_id>', methods=['DELETE'])
+@login_required
+def delete_message(message_id):
+    heymans = get_heymans()
+    heymans.messages.delete(message_id)
+    return jsonify(success=True)

heymans/routes/app.py

@@ -34,6 +34,8 @@ def chat_page():
     previous_timestamp = None
     previous_answer_model = None
     for role, message, metadata in heymans.messages:
+        message_id = metadata.get('message_id', 0)
+        delete_button = f'<button class="message-delete" onclick="deleteMessage(\'{message_id}\')"><i class="fas fa-trash"></i></button>'
         if role == 'assistant':
             html_body = utils.md(
                 f'{config.ai_name}: {config.process_ai_message(message)}')
@@ -62,7 +64,7 @@ def chat_page():
         else:
             answer_model_div = ''
 
-        html_content += f'<div class="{html_class}">{html_body}{timestamp_div}{answer_model_div}{sources_div}</div>'
+        html_content += f'<div class="message {html_class}" data-message-id="{message_id}">{delete_button}{html_body}{timestamp_div}{answer_model_div}{sources_div}</div>'
     return utils.render('chat.html', message_history=html_content)
 
 

heymans/templates/main.js

@@ -56,6 +56,12 @@ function initMain(event) {
     });
 }
 
+function generateUUID() {
+    return ([1e7]+-1e3+-4e3+-8e3+-1e11).replace(/[018]/g, c =>
+        (c ^ crypto.getRandomValues(new Uint8Array(1))[0] & 15 >> c / 4).toString(16)
+    );
+}
+
 async function fetchWithRetry(url, options, retries = 3) {
     for (let i = 0; i < retries; i++) {
         try {
@@ -69,13 +75,21 @@ async function fetchWithRetry(url, options, retries = 3) {
 
 async function sendMessage(message) {
     console.log('user message: ' + message)
+    const user_message_id = generateUUID()
     messageCounter.innerText = ''
     // Show the user's message
     if (message) {
         const userMessageBox = document.createElement('div');
         userMessageBox.innerText = 'You: ' + message;
-        userMessageBox.className = 'message-user';
+        userMessageBox.className = 'message-user message';
+        userMessageBox.setAttribute('data-message-id', user_message_id);
         responseDiv.appendChild(userMessageBox);
+        // Create a div for the delete button
+        const deleteButton = document.createElement('button');
+        deleteButton.className = 'message-delete';
+        deleteButton.innerHTML = '<i class="fas fa-trash"></i>';
+        deleteButton.onclick = () => deleteMessage(user_message_id);
+        userMessageBox.prepend(deleteButton);
     }
     // Show the loading indicator and animate it
     const loadingMessageBox = document.createElement('div');
@@ -105,7 +119,7 @@ async function sendMessage(message) {
     await fetchWithRetry('{{ server_url }}/api/chat/start', {
         method: 'POST',
         headers: {'Content-Type': 'application/json'},
-        body: requestBody(message, searchFirst)
+        body: requestBody(message, searchFirst, user_message_id)
     }).catch(e => {
         console.error('Failed to start chat session:', e);
     });
@@ -132,14 +146,20 @@ async function sendMessage(message) {
             }
             return
         }
+        const metadata = data.metadata
         // If the message is an actual message, we add it to the chat window
         // Create and append the message elements as in your existing code
         const aiMessage = document.createElement('div');
-        aiMessage.className = 'message-ai';
+        aiMessage.className = 'message-ai message';
+        aiMessage.setAttribute('data-message-id', metadata['message_id']);
         aiMessage.innerHTML = data.response;
         responseDiv.appendChild(aiMessage);
-        // Parse the JSON string to an object
-        const metadata = data.metadata
+        // Create a div for the delete button
+        const deleteButton = document.createElement('button');
+        deleteButton.className = 'message-delete';
+        deleteButton.innerHTML = '<i class="fas fa-trash"></i>';
+        deleteButton.onclick = () => deleteMessage(metadata['message_id']);
+        aiMessage.prepend(deleteButton);
         // Create a div for timestamp
         const timestampDiv = document.createElement('div');
         timestampDiv.className = 'message-timestamp';
@@ -206,8 +226,30 @@ async function sendMessage(message) {
     cancelButton.onclick = cancelStreaming;
 }
 
-function requestBody(message, searchFirst) {
-    return JSON.stringify({message: message, search_first: searchFirst})
+function deleteMessage(messageId) {
+    fetch(`/api/message/delete/${messageId}`, {
+        method: 'DELETE',
+    })
+    .then(response => response.json())
+    .then(data => {
+        if (data.success) {
+            const messageElement = document.querySelector(`.message[data-message-id="${messageId}"]`);
+            if (messageElement) {
+                messageElement.remove();
+            }
+        } else {
+            console.error('Failed to delete message:', data.error);
+        }
+    }) .catch(error => console.error('Error deleting message:', error)); }
+
+
+
+function requestBody(message, searchFirst, user_message_id) {
+    return JSON.stringify({
+        message: message,
+        search_first: searchFirst,
+        message_id:user_message_id
+    })
 }
 
 document.addEventListener('DOMContentLoaded', globalElements)

heymans/templates/stylesheet.css.jinja

@@ -279,3 +279,18 @@ button#start {
 .hidden-data {
     display: none;
 }
+
+.message-delete {
+    float: right;
+    padding: 0px;
+    border: none;
+    background: none;
+    display: none;
+    font-size: 0.6em;
+    color: #90A4AE;    
+}
+
+.message-user:hover .message-delete,
+.message-ai:hover .message-delete {
+    display: block;
+}