wip - initial/ experimental insert of messages to db

serverlessnext · Jul 12, 2024 · b1bbcb8 · b1bbcb8
1 parent f1d204d
commit b1bbcb8
Show file tree

Hide file tree

Showing 3 changed files with 34 additions and 5 deletions.
diff --git a/lumni/src/apps/builtin/llm/prompt/src/chat/db/schema.rs b/lumni/src/apps/builtin/llm/prompt/src/chat/db/schema.rs
@@ -258,6 +258,35 @@ impl InMemoryDatabase {
             .unwrap_or_default()
     }
 
+    pub fn finalize_last_exchange(
+        &mut self,
+        db_store: &Arc<Mutex<ConversationDatabaseStore>>,
+        conversation_id: ConversationId,
+    ) {
+        let exchange = self.get_last_exchange(conversation_id);
+        if let Some(exchange) = exchange {
+            let messages = self.get_exchange_messages(exchange.id);
+            let attachments = messages
+                .iter()
+                .flat_map(|message| {
+                    self.get_message_attachments(message.id)
+                })
+                .collect::<Vec<_>>();
+
+            // Convert Vec<&Message> to Vec<Message> and Vec<&Attachment> to Vec<Attachment>
+            let owned_messages: Vec<Message> = messages.into_iter().cloned().collect();
+            let owned_attachments: Vec<Attachment> = attachments.into_iter().cloned().collect();
+
+            eprintln!("Owned messages: {:?}", owned_messages);
+
+            let mut db_lock_store = db_store.lock().unwrap();
+            db_lock_store.store_finalized_exchange(&exchange, &owned_messages, &owned_attachments);
+            let result = db_lock_store.commit_queued_operations();
+            eprintln!("Commit result: {:?}", result);
+        }
+
+    }
+
     pub fn get_last_message_of_exchange(
         &self,
         exchange_id: ExchangeId,

diff --git a/lumni/src/apps/builtin/llm/prompt/src/chat/db/store.rs b/lumni/src/apps/builtin/llm/prompt/src/chat/db/store.rs
@@ -173,9 +173,6 @@ impl ConversationDatabaseStore {
             );
             self.db.queue_operation(attachment_sql);
         }
-
-        // Commit the transaction
-        //self.db.queue_operation("COMMIT;".to_string());
     }
 
     pub fn commit_queued_operations(&mut self) -> Result<(), SqliteError> {

diff --git a/lumni/src/apps/builtin/llm/prompt/src/chat/instruction.rs b/lumni/src/apps/builtin/llm/prompt/src/chat/instruction.rs
@@ -127,7 +127,7 @@ impl PromptInstruction {
         db_conn: &ConversationDatabase,
     ) {
         ExchangeHandler::put_last_response(
-            &mut db_conn.cache.lock().unwrap(),
+            db_conn,
             self.current_conversation_id,
             answer,
             tokens_predicted,
@@ -483,11 +483,13 @@ impl ExchangeHandler {
     }
 
     pub fn put_last_response(
-        db_lock: &mut MutexGuard<InMemoryDatabase>,
+        db_conn: &ConversationDatabase,
         current_conversation_id: ConversationId,
         answer: &str,
         tokens_predicted: Option<usize>,
     ) {
+        let db_lock = &mut db_conn.cache.lock().unwrap();
+
         let (message_id, is_assistant) = if let Some(last_exchange) =
             db_lock.get_last_exchange(current_conversation_id)
         {
@@ -510,6 +512,7 @@ impl ExchangeHandler {
         if let (Some(id), true) = (message_id, is_assistant) {
             let token_length = tokens_predicted.map(|t| t as i32);
             db_lock.update_message_by_id(id, answer, token_length);
+            db_lock.finalize_last_exchange(&db_conn.store, current_conversation_id);
         }
     }
 }