use crate::serve::templates::ChatMessage;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::path::PathBuf;
use std::sync::{Arc, RwLock};
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ConversationMeta {
pub id: String,
pub title: String,
pub model: String,
pub created: u64,
pub updated: u64,
pub message_count: usize,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Conversation {
pub meta: ConversationMeta,
pub messages: Vec<ChatMessage>,
}
pub struct ConversationStore {
conversations: RwLock<HashMap<String, Conversation>>,
data_dir: Option<PathBuf>,
counter: std::sync::atomic::AtomicU64,
}
impl ConversationStore {
#[must_use]
pub fn in_memory() -> Arc<Self> {
Arc::new(Self {
conversations: RwLock::new(HashMap::new()),
data_dir: None,
counter: std::sync::atomic::AtomicU64::new(0),
})
}
#[must_use]
pub fn with_data_dir(dir: PathBuf) -> Arc<Self> {
let _ = std::fs::create_dir_all(&dir);
let mut conversations = HashMap::new();
let mut max_seq = 0u64;
if let Ok(entries) = std::fs::read_dir(&dir) {
for entry in entries.flatten() {
let path = entry.path();
if path.extension().and_then(|e| e.to_str()) == Some("jsonl") {
let conv_id =
path.file_stem().and_then(|s| s.to_str()).unwrap_or("unknown").to_string();
let mut messages = Vec::new();
if let Ok(content) = std::fs::read_to_string(&path) {
for line in content.lines() {
if let Ok(msg) = serde_json::from_str::<ChatMessage>(line) {
messages.push(msg);
}
}
}
if let Some(seq_str) = conv_id.rsplit('-').next() {
if let Ok(seq) = seq_str.parse::<u64>() {
max_seq = max_seq.max(seq + 1);
}
}
let title = messages
.first()
.filter(|m| matches!(m.role, crate::serve::templates::Role::User))
.map(|m| auto_title(&m.content))
.unwrap_or_else(|| "Loaded conversation".to_string());
let conv = Conversation {
meta: ConversationMeta {
id: conv_id.clone(),
title,
model: "unknown".to_string(),
created: epoch_secs(),
updated: epoch_secs(),
message_count: messages.len(),
},
messages,
};
conversations.insert(conv_id, conv);
}
}
}
let loaded = conversations.len();
if loaded > 0 {
eprintln!("[banco] Loaded {loaded} conversations from {}", dir.display());
}
Arc::new(Self {
conversations: RwLock::new(conversations),
data_dir: Some(dir),
counter: std::sync::atomic::AtomicU64::new(max_seq),
})
}
pub fn create(&self, model: &str) -> String {
let seq = self.counter.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
let id = format!("conv-{}-{seq}", epoch_secs());
let meta = ConversationMeta {
id: id.clone(),
title: "New conversation".to_string(),
model: model.to_string(),
created: epoch_secs(),
updated: epoch_secs(),
message_count: 0,
};
let conv = Conversation { meta, messages: Vec::new() };
if let Ok(mut store) = self.conversations.write() {
store.insert(id.clone(), conv);
}
id
}
pub fn append(&self, id: &str, message: ChatMessage) -> Result<(), ConversationError> {
let mut store = self.conversations.write().map_err(|_| ConversationError::LockPoisoned)?;
let conv = store.get_mut(id).ok_or(ConversationError::NotFound(id.to_string()))?;
if conv.messages.is_empty()
&& conv.meta.title == "New conversation"
&& matches!(message.role, crate::serve::templates::Role::User)
{
conv.meta.title = auto_title(&message.content);
}
conv.messages.push(message);
conv.meta.message_count = conv.messages.len();
conv.meta.updated = epoch_secs();
if let Some(ref dir) = self.data_dir {
let path = dir.join(format!("{id}.jsonl"));
let json = serde_json::to_string(&conv.messages.last().expect("just pushed"))
.unwrap_or_default();
let _ = std::fs::OpenOptions::new().create(true).append(true).open(path).and_then(
|mut f| {
use std::io::Write;
writeln!(f, "{json}")
},
);
}
Ok(())
}
#[must_use]
pub fn list(&self) -> Vec<ConversationMeta> {
let store = self.conversations.read().unwrap_or_else(|e| e.into_inner());
let mut metas: Vec<ConversationMeta> = store.values().map(|c| c.meta.clone()).collect();
metas.sort_by(|a, b| b.updated.cmp(&a.updated));
metas
}
#[must_use]
pub fn get(&self, id: &str) -> Option<Conversation> {
let store = self.conversations.read().unwrap_or_else(|e| e.into_inner());
store.get(id).cloned()
}
pub fn rename(&self, id: &str, title: &str) -> Result<(), ConversationError> {
let mut store = self.conversations.write().map_err(|_| ConversationError::LockPoisoned)?;
let conv = store.get_mut(id).ok_or(ConversationError::NotFound(id.to_string()))?;
conv.meta.title = title.to_string();
conv.meta.updated = epoch_secs();
Ok(())
}
pub fn delete(&self, id: &str) -> Result<(), ConversationError> {
let mut store = self.conversations.write().map_err(|_| ConversationError::LockPoisoned)?;
store.remove(id).ok_or(ConversationError::NotFound(id.to_string()))?;
if let Some(ref dir) = self.data_dir {
let _ = std::fs::remove_file(dir.join(format!("{id}.jsonl")));
}
Ok(())
}
#[must_use]
pub fn len(&self) -> usize {
self.conversations.read().map(|s| s.len()).unwrap_or(0)
}
#[must_use]
pub fn is_empty(&self) -> bool {
self.len() == 0
}
#[must_use]
pub fn search(&self, query: &str) -> Vec<ConversationMeta> {
let store = self.conversations.read().unwrap_or_else(|e| e.into_inner());
let query_lower = query.to_lowercase();
let mut results: Vec<ConversationMeta> = store
.values()
.filter(|c| {
c.meta.title.to_lowercase().contains(&query_lower)
|| c.messages.iter().any(|m| m.content.to_lowercase().contains(&query_lower))
})
.map(|c| c.meta.clone())
.collect();
results.sort_by(|a, b| b.updated.cmp(&a.updated));
results
}
#[must_use]
pub fn export_all(&self) -> Vec<Conversation> {
let store = self.conversations.read().unwrap_or_else(|e| e.into_inner());
let mut convs: Vec<Conversation> = store.values().cloned().collect();
convs.sort_by(|a, b| b.meta.updated.cmp(&a.meta.updated));
convs
}
pub fn import_all(&self, conversations: Vec<Conversation>) -> usize {
let mut store = self.conversations.write().unwrap_or_else(|e| e.into_inner());
let count = conversations.len();
for conv in conversations {
store.insert(conv.meta.id.clone(), conv);
}
count
}
}
fn auto_title(content: &str) -> String {
let words: Vec<&str> = content.split_whitespace().take(5).collect();
if words.is_empty() {
"New conversation".to_string()
} else {
let mut title = words.join(" ");
if content.split_whitespace().count() > 5 {
title.push_str("...");
}
title
}
}
fn epoch_secs() -> u64 {
std::time::SystemTime::now().duration_since(std::time::UNIX_EPOCH).unwrap_or_default().as_secs()
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ConversationError {
NotFound(String),
LockPoisoned,
}
impl std::fmt::Display for ConversationError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::NotFound(id) => write!(f, "Conversation not found: {id}"),
Self::LockPoisoned => write!(f, "Internal lock error"),
}
}
}
impl std::error::Error for ConversationError {}