From 4fab757f2daca174180927640c0e39c99f1f2937 Mon Sep 17 00:00:00 2001
From: dal <dallin@buster.so>
Date: Wed, 5 Mar 2025 10:31:35 -0700
Subject: [PATCH] message ids

---
 .../handlers/src/chats/get_chat_handler.rs    | 29 +++----
 .../handlers/src/chats/post_chat_handler.rs   | 87 ++++++++++---------
 api/libs/handlers/src/chats/types.rs          | 84 +++++++++++++++++-
 3 files changed, 141 insertions(+), 59 deletions(-)

diff --git a/api/libs/handlers/src/chats/get_chat_handler.rs b/api/libs/handlers/src/chats/get_chat_handler.rs
index 737e12ad7..935466cb0 100644
--- a/api/libs/handlers/src/chats/get_chat_handler.rs
+++ b/api/libs/handlers/src/chats/get_chat_handler.rs
@@ -119,7 +119,7 @@ pub async fn get_chat_handler(chat_id: &Uuid, user_id: &Uuid) -> Result<ChatWith
     )?;
 
     // Transform messages into ThreadMessage format
-    let thread_messages = messages
+    let thread_messages: Vec<ChatMessage> = messages
         .into_iter()
         .map(|msg| {
             let sender_avatar = msg
@@ -148,15 +148,13 @@ pub async fn get_chat_handler(chat_id: &Uuid, user_id: &Uuid) -> Result<ChatWith
                 sender_avatar,
             };
 
-            let chat_message = ChatMessage::new_with_messages(
+            ChatMessage::new_with_messages(
                 msg.id,
                 request_message,
                 response_messages,
                 reasoning,
                 None,
-            );
-
-            chat_message
+            )
         })
         .collect();
 
@@ -167,17 +165,14 @@ pub async fn get_chat_handler(chat_id: &Uuid, user_id: &Uuid) -> Result<ChatWith
         .and_then(|v| v.as_str())
         .map(String::from);
 
-    // Construct and return the ThreadWithMessages
-    Ok(ChatWithMessages {
-        id: thread.id,
-        title: thread.title,
-        is_favorited: false, // Not implemented in current schema
-        messages: thread_messages,
-        created_at: thread.created_at.to_string(),
-        updated_at: thread.updated_at.to_string(),
-        created_by: thread.user_email,
-        created_by_id: thread.user_id.to_string(),
-        created_by_name: thread.user_name.unwrap_or_else(|| "Unknown".to_string()),
+    // Construct and return the ThreadWithMessages using new_with_messages
+    Ok(ChatWithMessages::new_with_messages(
+        thread.id,
+        thread.title,
+        thread_messages,
+        false, // is_favorited not implemented in current schema
+        thread.user_id.to_string(),
+        thread.user_name.unwrap_or_else(|| "Unknown".to_string()),
         created_by_avatar,
-    })
+    ))
 }
diff --git a/api/libs/handlers/src/chats/post_chat_handler.rs b/api/libs/handlers/src/chats/post_chat_handler.rs
index 0d1fe7ea6..129c6e514 100644
--- a/api/libs/handlers/src/chats/post_chat_handler.rs
+++ b/api/libs/handlers/src/chats/post_chat_handler.rs
@@ -301,20 +301,23 @@ pub async fn post_chat_handler(
         prepare_final_message_state(&all_transformed_containers)?;
 
     // Update chat_with_messages with final state
-    if let Some(chat_message) = chat_with_messages.messages.first_mut() {
-        *chat_message = ChatMessage::new_with_messages(
-            message_id,
-            chat_message.request_message.clone(),
-            response_messages.clone(),
-            reasoning_messages.clone(),
-            None,
-        );
-    }
-
-    let final_reasoning_message = format!("Reasoned for {} seconds", reasoning_duration);
+    let message = ChatMessage::new_with_messages(
+        message_id,
+        ChatUserMessage {
+            request: request.prompt.clone(),
+            sender_id: user.id.clone(),
+            sender_name: user.name.clone().unwrap_or_default(),
+            sender_avatar: None,
+        },
+        response_messages.clone(),
+        reasoning_messages.clone(),
+        Some(format!("Reasoned for {} seconds", reasoning_duration).to_string()),
+    );
+    
+    chat_with_messages.update_message(message);
 
     // Create and store message in the database with final state
-    let message = Message {
+    let db_message = Message {
         id: message_id,
         request_message: request.prompt,
         chat_id,
@@ -324,20 +327,20 @@ pub async fn post_chat_handler(
         deleted_at: None,
         response_messages: serde_json::to_value(&response_messages)?,
         reasoning: serde_json::to_value(&reasoning_messages)?,
-        final_reasoning_message,
+        final_reasoning_message: format!("Reasoned for {} seconds", reasoning_duration),
         title: title.title.clone().unwrap_or_default(),
         raw_llm_messages: serde_json::to_value(&raw_llm_messages)?,
     };
 
     // Insert message into database
     insert_into(messages::table)
-        .values(&message)
+        .values(&db_message)
         .execute(&mut conn)
         .await?;
 
     // First process completed files (database updates only)
     let _ =
-        process_completed_files(&mut conn, &message, &all_messages, &user_org_id, &user.id).await?;
+        process_completed_files(&mut conn, &db_message, &all_messages, &user_org_id, &user.id).await?;
 
     // Then send text response messages
     if let Some(tx) = &tx {
@@ -1619,8 +1622,8 @@ async fn initialize_chat(
         // Get existing chat - no need to create new chat in DB
         let mut existing_chat = get_chat_handler(&existing_chat_id, &user.id).await?;
 
-        // Add new message to existing chat
-        existing_chat.messages.push(ChatMessage::new_with_messages(
+        // Create new message
+        let message = ChatMessage::new_with_messages(
             message_id,
             ChatUserMessage {
                 request: request.prompt.clone(),
@@ -1631,7 +1634,10 @@ async fn initialize_chat(
             Vec::new(),
             Vec::new(),
             None,
-        ));
+        );
+
+        // Add message to existing chat
+        existing_chat.add_message(message);
 
         Ok((existing_chat_id, message_id, existing_chat))
     } else {
@@ -1648,29 +1654,28 @@ async fn initialize_chat(
             updated_by: user.id.clone(),
         };
 
-        let chat_with_messages = ChatWithMessages {
-            id: chat_id,
-            title: request.prompt.clone(),
-            is_favorited: false,
-            messages: vec![ChatMessage::new_with_messages(
-                message_id,
-                ChatUserMessage {
-                    request: request.prompt.clone(),
-                    sender_id: user.id.clone(),
-                    sender_name: user.name.clone().unwrap_or_default(),
-                    sender_avatar: None,
-                },
-                Vec::new(),
-                Vec::new(),
-                None,
-            )],
-            created_at: Utc::now().to_string(),
-            updated_at: Utc::now().to_string(),
-            created_by: user.id.to_string(),
-            created_by_id: user.id.to_string(),
-            created_by_name: user.name.clone().unwrap_or_default(),
-            created_by_avatar: None,
-        };
+        // Create initial message
+        let message = ChatMessage::new_with_messages(
+            message_id,
+            ChatUserMessage {
+                request: request.prompt.clone(),
+                sender_id: user.id.clone(),
+                sender_name: user.name.clone().unwrap_or_default(),
+                sender_avatar: None,
+            },
+            Vec::new(),
+            Vec::new(),
+            None,
+        );
+
+        let mut chat_with_messages = ChatWithMessages::new(
+            request.prompt.clone(),
+            user.id.to_string(),
+            user.name.clone().unwrap_or_default(),
+            None,
+        );
+        chat_with_messages.id = chat_id;
+        chat_with_messages.add_message(message);
 
         // Only create new chat in DB if this is a new chat
         let mut conn = get_pg_pool().get().await?;
diff --git a/api/libs/handlers/src/chats/types.rs b/api/libs/handlers/src/chats/types.rs
index dfc97d570..bf3a2b039 100644
--- a/api/libs/handlers/src/chats/types.rs
+++ b/api/libs/handlers/src/chats/types.rs
@@ -1,3 +1,5 @@
+use std::collections::HashMap;
+
 use serde::{Deserialize, Serialize};
 use uuid::Uuid;
 
@@ -8,7 +10,8 @@ pub struct ChatWithMessages {
     pub id: Uuid,
     pub title: String,
     pub is_favorited: bool,
-    pub messages: Vec<ChatMessage>,
+    pub message_ids: Vec<String>,
+    pub messages: HashMap<String, ChatMessage>,
     pub created_at: String,
     pub updated_at: String,
     pub created_by: String,
@@ -17,3 +20,82 @@ pub struct ChatWithMessages {
     pub created_by_avatar: Option<String>,
 }
 
+impl ChatWithMessages {
+    pub fn new(
+        title: String,
+        created_by_id: String,
+        created_by_name: String,
+        created_by_avatar: Option<String>,
+    ) -> Self {
+        let now = chrono::Utc::now().to_rfc3339();
+        Self {
+            id: uuid::Uuid::new_v4(),
+            title,
+            is_favorited: false,
+            message_ids: Vec::new(),
+            messages: HashMap::new(),
+            created_at: now.clone(),
+            updated_at: now,
+            created_by: created_by_id.clone(),
+            created_by_id,
+            created_by_name,
+            created_by_avatar,
+        }
+    }
+
+    pub fn new_with_messages(
+        id: uuid::Uuid,
+        title: String,
+        messages: Vec<ChatMessage>,
+        is_favorited: bool,
+        created_by_id: String,
+        created_by_name: String,
+        created_by_avatar: Option<String>,
+    ) -> Self {
+        let now = chrono::Utc::now().to_rfc3339();
+        
+        // Convert messages into a HashMap and collect their IDs
+        let message_ids: Vec<String> = messages.iter()
+            .map(|msg| msg.id.to_string())
+            .collect();
+            
+        let messages_map: HashMap<String, ChatMessage> = messages.into_iter()
+            .map(|msg| (msg.id.to_string(), msg))
+            .collect();
+
+        Self {
+            id,
+            title,
+            is_favorited,
+            message_ids,
+            messages: messages_map,
+            created_at: now.clone(),
+            updated_at: now,
+            created_by: created_by_id.clone(),
+            created_by_id,
+            created_by_name,
+            created_by_avatar,
+        }
+    }
+
+    pub fn add_message(&mut self, message: ChatMessage) {
+        let message_id = message.id.to_string();
+        if !self.message_ids.contains(&message_id) {
+            self.message_ids.push(message_id.clone());
+        }
+        self.messages.insert(message_id, message);
+        self.updated_at = chrono::Utc::now().to_rfc3339();
+    }
+
+    pub fn update_message(&mut self, message: ChatMessage) {
+        let message_id = message.id.to_string();
+        if !self.message_ids.contains(&message_id) {
+            self.message_ids.push(message_id.clone());
+        }
+        self.messages.insert(message_id, message);
+        self.updated_at = chrono::Utc::now().to_rfc3339();
+    }
+}
+
+
+