-
Notifications
You must be signed in to change notification settings - Fork 0
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
simplify ConversationDatabaseStore, add additional display queries
- Loading branch information
Showing
9 changed files
with
321 additions
and
225 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,78 @@ | ||
use lumni::api::error::ApplicationError; | ||
|
||
use super::schema::{ | ||
Conversation, ConversationId, Message, | ||
}; | ||
use super::ConversationDatabaseStore; | ||
pub use crate::external as lumni; | ||
|
||
impl ConversationDatabaseStore { | ||
pub async fn print_last_conversation( | ||
&self, | ||
) -> Result<(), ApplicationError> { | ||
if let Some((conversation, messages)) = | ||
self.fetch_conversation(None, None)? | ||
{ | ||
display_conversation_with_messages(&conversation, &messages); | ||
} else { | ||
println!("No conversations found."); | ||
} | ||
Ok(()) | ||
} | ||
|
||
pub async fn print_conversation_list( | ||
&self, | ||
limit: usize, | ||
) -> Result<(), ApplicationError> { | ||
let conversations = self.fetch_conversation_list(limit)?; | ||
for conversation in conversations { | ||
println!( | ||
"ID: {}, Name: {}, Updated: {}", | ||
conversation.id.0, conversation.name, conversation.updated_at | ||
); | ||
} | ||
Ok(()) | ||
} | ||
|
||
pub async fn print_conversation_by_id( | ||
&self, | ||
id: &str, | ||
) -> Result<(), ApplicationError> { | ||
let conversation_id = ConversationId(id.parse().map_err(|_| { | ||
ApplicationError::NotFound( | ||
format!("Conversation {id} not found in database"), | ||
)})?); | ||
|
||
if let Some((conversation, messages)) = | ||
self.fetch_conversation(Some(conversation_id), None)? | ||
{ | ||
display_conversation_with_messages(&conversation, &messages); | ||
} else { | ||
println!("Conversation not found."); | ||
} | ||
Ok(()) | ||
} | ||
} | ||
|
||
fn display_conversation_with_messages( | ||
conversation: &Conversation, | ||
messages: &[Message], | ||
) { | ||
println!( | ||
"Conversation: {} (ID: {})", | ||
conversation.name, conversation.id.0 | ||
); | ||
println!("Updated at: {}", conversation.updated_at); | ||
|
||
if !messages.is_empty() { | ||
println!("Messages:"); | ||
for message in messages { | ||
println!(" Role: {}", message.role); | ||
println!(" Content: {}", message.content); | ||
println!(" ---"); | ||
} | ||
} else { | ||
println!("No messages"); | ||
} | ||
println!("==============================="); | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -1,85 +1,9 @@ | ||
use std::path::PathBuf; | ||
use std::sync::{Arc, Mutex}; | ||
|
||
use rusqlite::Error as SqliteError; | ||
|
||
mod connector; | ||
mod display; | ||
mod schema; | ||
mod store; | ||
|
||
pub use schema::{ | ||
Attachment, Conversation, ConversationCache, ConversationId, Exchange, | ||
ExchangeId, Message, ModelId, | ||
}; | ||
pub use schema::{ConversationCache, Exchange, ExchangeId, ConversationId, Message, ModelId}; | ||
pub use store::ConversationDatabaseStore; | ||
|
||
pub use super::PromptRole; | ||
|
||
pub struct ConversationDatabase { | ||
pub store: Arc<Mutex<ConversationDatabaseStore>>, | ||
} | ||
|
||
impl ConversationDatabase { | ||
pub fn new(sqlite_file: &PathBuf) -> Result<Self, SqliteError> { | ||
Ok(Self { | ||
store: Arc::new(Mutex::new(ConversationDatabaseStore::new( | ||
sqlite_file, | ||
)?)), | ||
}) | ||
} | ||
|
||
pub fn new_conversation( | ||
&self, | ||
name: &str, | ||
parent_id: Option<ConversationId>, | ||
) -> Result<ConversationId, SqliteError> { | ||
let mut store = self.store.lock().unwrap(); | ||
let conversation = Conversation { | ||
id: ConversationId(-1), // Temporary ID | ||
name: name.to_string(), | ||
metadata: serde_json::Value::Null, | ||
parent_conversation_id: parent_id, | ||
fork_exchange_id: None, | ||
schema_version: 1, | ||
created_at: 0, | ||
updated_at: 0, | ||
is_deleted: false, | ||
}; | ||
let conversation_id = store.put_new_conversation(&conversation)?; | ||
Ok(conversation_id) | ||
} | ||
|
||
pub fn finalize_exchange( | ||
&self, | ||
exchange: &Exchange, | ||
cache: &ConversationCache, | ||
) -> Result<(), SqliteError> { | ||
let messages = cache.get_exchange_messages(exchange.id); | ||
let attachments = messages | ||
.iter() | ||
.flat_map(|message| cache.get_message_attachments(message.id)) | ||
.collect::<Vec<_>>(); | ||
let owned_messages: Vec<Message> = | ||
messages.into_iter().cloned().collect(); | ||
let owned_attachments: Vec<Attachment> = | ||
attachments.into_iter().cloned().collect(); | ||
let mut store = self.store.lock().unwrap(); | ||
store.put_finalized_exchange( | ||
exchange, | ||
&owned_messages, | ||
&owned_attachments, | ||
)?; | ||
Ok(()) | ||
} | ||
|
||
pub fn get_recent_conversations_with_last_exchange_and_messages( | ||
&self, | ||
limit: usize, | ||
) -> Result< | ||
Vec<(Conversation, Option<(Exchange, Vec<Message>)>)>, | ||
SqliteError, | ||
> { | ||
let mut store = self.store.lock().unwrap(); | ||
store.get_recent_conversations_with_last_exchange_and_messages(limit) | ||
} | ||
} |
Oops, something went wrong.