//! Workspace-related WorkspaceStore implementation for LibSqlBackend. use std::collections::HashMap; use async_trait::async_trait; use libsql::params; use uuid::Uuid; use super::{ LibSqlBackend, fmt_ts, get_i64, get_opt_text, get_opt_ts, get_text, get_ts, row_to_memory_document, }; use crate::db::WorkspaceStore; use crate::error::WorkspaceError; use crate::workspace::{ MemoryChunk, MemoryDocument, RankedResult, SearchConfig, SearchResult, WorkspaceEntry, fuse_results, }; use chrono::Utc; #[async_trait] impl WorkspaceStore for LibSqlBackend { async fn get_document_by_path( &self, user_id: &str, agent_id: Option, path: &str, ) -> Result { let conn = self .connect() .await .map_err(|e| WorkspaceError::SearchFailed { reason: e.to_string(), })?; let agent_id_str = agent_id.map(|id| id.to_string()); let mut rows = conn .query( r#" SELECT id, user_id, agent_id, path, content, created_at, updated_at, metadata FROM memory_documents WHERE user_id = ?1 AND agent_id IS ?2 AND path = ?3 "#, params![user_id, agent_id_str.as_deref(), path], ) .await .map_err(|e| WorkspaceError::SearchFailed { reason: format!("Query failed: {}", e), })?; match rows .next() .await .map_err(|e| WorkspaceError::SearchFailed { reason: format!("Query failed: {}", e), })? { Some(row) => Ok(row_to_memory_document(&row)), None => Err(WorkspaceError::DocumentNotFound { doc_type: path.to_string(), user_id: user_id.to_string(), }), } } async fn get_document_by_id(&self, id: Uuid) -> Result { let conn = self .connect() .await .map_err(|e| WorkspaceError::SearchFailed { reason: e.to_string(), })?; let mut rows = conn .query( r#" SELECT id, user_id, agent_id, path, content, created_at, updated_at, metadata FROM memory_documents WHERE id = ?1 "#, params![id.to_string()], ) .await .map_err(|e| WorkspaceError::SearchFailed { reason: format!("Query failed: {}", e), })?; match rows .next() .await .map_err(|e| WorkspaceError::SearchFailed { reason: format!("Query failed: {}", e), })? { Some(row) => Ok(row_to_memory_document(&row)), None => Err(WorkspaceError::DocumentNotFound { doc_type: "unknown".to_string(), user_id: "unknown".to_string(), }), } } async fn get_or_create_document_by_path( &self, user_id: &str, agent_id: Option, path: &str, ) -> Result { // Try get match self.get_document_by_path(user_id, agent_id, path).await { Ok(doc) => return Ok(doc), Err(WorkspaceError::DocumentNotFound { .. }) => {} Err(e) => return Err(e), } // Create let conn = self .connect() .await .map_err(|e| WorkspaceError::SearchFailed { reason: e.to_string(), })?; let id = Uuid::new_v4(); let agent_id_str = agent_id.map(|id| id.to_string()); conn.execute( r#" INSERT INTO memory_documents (id, user_id, agent_id, path, content, metadata) VALUES (?1, ?2, ?3, ?4, '', '{}') ON CONFLICT (user_id, agent_id, path) DO NOTHING "#, params![id.to_string(), user_id, agent_id_str.as_deref(), path], ) .await .map_err(|e| WorkspaceError::SearchFailed { reason: format!("Insert failed: {}", e), })?; self.get_document_by_path(user_id, agent_id, path).await } async fn update_document(&self, id: Uuid, content: &str) -> Result<(), WorkspaceError> { let conn = self .connect() .await .map_err(|e| WorkspaceError::SearchFailed { reason: e.to_string(), })?; let now = fmt_ts(&Utc::now()); conn.execute( "UPDATE memory_documents SET content = ?2, updated_at = ?3 WHERE id = ?1", params![id.to_string(), content, now], ) .await .map_err(|e| WorkspaceError::SearchFailed { reason: format!("Update failed: {}", e), })?; Ok(()) } async fn delete_document_by_path( &self, user_id: &str, agent_id: Option, path: &str, ) -> Result<(), WorkspaceError> { let doc = self.get_document_by_path(user_id, agent_id, path).await?; self.delete_chunks(doc.id).await?; let conn = self .connect() .await .map_err(|e| WorkspaceError::SearchFailed { reason: e.to_string(), })?; let agent_id_str = agent_id.map(|id| id.to_string()); conn.execute( "DELETE FROM memory_documents WHERE user_id = ?1 AND agent_id IS ?2 AND path = ?3", params![user_id, agent_id_str.as_deref(), path], ) .await .map_err(|e| WorkspaceError::SearchFailed { reason: format!("Delete failed: {}", e), })?; Ok(()) } async fn list_directory( &self, user_id: &str, agent_id: Option, directory: &str, ) -> Result, WorkspaceError> { let conn = self .connect() .await .map_err(|e| WorkspaceError::SearchFailed { reason: e.to_string(), })?; let dir = if !directory.is_empty() && !directory.ends_with('/') { format!("{}/", directory) } else { directory.to_string() }; let agent_id_str = agent_id.map(|id| id.to_string()); let pattern = if dir.is_empty() { "%".to_string() } else { format!("{}%", dir) }; let mut rows = conn .query( r#" SELECT path, updated_at, substr(content, 1, 200) as content_preview FROM memory_documents WHERE user_id = ?1 AND agent_id IS ?2 AND (?3 = '%' OR path LIKE ?3) ORDER BY path "#, params![user_id, agent_id_str.as_deref(), pattern], ) .await .map_err(|e| WorkspaceError::SearchFailed { reason: format!("List directory failed: {}", e), })?; let mut entries_map: HashMap = HashMap::new(); while let Some(row) = rows .next() .await .map_err(|e| WorkspaceError::SearchFailed { reason: format!("Query failed: {}", e), })? { let full_path = get_text(&row, 0); let updated_at = get_opt_ts(&row, 1); let content_preview = get_opt_text(&row, 2); let relative = if dir.is_empty() { &full_path } else if let Some(stripped) = full_path.strip_prefix(&dir) { stripped } else { continue; }; let child_name = if let Some(slash_pos) = relative.find('/') { &relative[..slash_pos] } else { relative }; if child_name.is_empty() { continue; } let is_dir = relative.contains('/'); let entry_path = if dir.is_empty() { child_name.to_string() } else { format!("{}{}", dir, child_name) }; entries_map .entry(child_name.to_string()) .and_modify(|e| { if is_dir { e.is_directory = true; e.content_preview = None; } if let (Some(existing), Some(new)) = (&e.updated_at, &updated_at) && new > existing { e.updated_at = Some(*new); } }) .or_insert(WorkspaceEntry { path: entry_path, is_directory: is_dir, updated_at, content_preview: if is_dir { None } else { content_preview }, }); } let mut entries: Vec = entries_map.into_values().collect(); entries.sort_by(|a, b| a.path.cmp(&b.path)); Ok(entries) } async fn list_all_paths( &self, user_id: &str, agent_id: Option, ) -> Result, WorkspaceError> { let conn = self .connect() .await .map_err(|e| WorkspaceError::SearchFailed { reason: e.to_string(), })?; let agent_id_str = agent_id.map(|id| id.to_string()); let mut rows = conn .query( "SELECT path FROM memory_documents WHERE user_id = ?1 AND agent_id IS ?2 ORDER BY path", params![user_id, agent_id_str.as_deref()], ) .await .map_err(|e| WorkspaceError::SearchFailed { reason: format!("List paths failed: {}", e), })?; let mut paths = Vec::new(); while let Some(row) = rows .next() .await .map_err(|e| WorkspaceError::SearchFailed { reason: format!("Query failed: {}", e), })? { paths.push(get_text(&row, 0)); } Ok(paths) } async fn list_documents( &self, user_id: &str, agent_id: Option, ) -> Result, WorkspaceError> { let conn = self .connect() .await .map_err(|e| WorkspaceError::SearchFailed { reason: e.to_string(), })?; let agent_id_str = agent_id.map(|id| id.to_string()); let mut rows = conn .query( r#" SELECT id, user_id, agent_id, path, content, created_at, updated_at, metadata FROM memory_documents WHERE user_id = ?1 AND agent_id IS ?2 ORDER BY updated_at DESC "#, params![user_id, agent_id_str.as_deref()], ) .await .map_err(|e| WorkspaceError::SearchFailed { reason: format!("Query failed: {}", e), })?; let mut docs = Vec::new(); while let Some(row) = rows .next() .await .map_err(|e| WorkspaceError::SearchFailed { reason: format!("Query failed: {}", e), })? { docs.push(row_to_memory_document(&row)); } Ok(docs) } async fn delete_chunks(&self, document_id: Uuid) -> Result<(), WorkspaceError> { let conn = self .connect() .await .map_err(|e| WorkspaceError::ChunkingFailed { reason: e.to_string(), })?; conn.execute( "DELETE FROM memory_chunks WHERE document_id = ?1", params![document_id.to_string()], ) .await .map_err(|e| WorkspaceError::ChunkingFailed { reason: format!("Delete failed: {}", e), })?; Ok(()) } async fn insert_chunk( &self, document_id: Uuid, chunk_index: i32, content: &str, embedding: Option<&[f32]>, ) -> Result { let conn = self .connect() .await .map_err(|e| WorkspaceError::ChunkingFailed { reason: e.to_string(), })?; let id = Uuid::new_v4(); let embedding_blob = embedding.map(|e| { let bytes: Vec = e.iter().flat_map(|f| f.to_le_bytes()).collect(); bytes }); conn.execute( r#" INSERT INTO memory_chunks (id, document_id, chunk_index, content, embedding) VALUES (?1, ?2, ?3, ?4, ?5) "#, params![ id.to_string(), document_id.to_string(), chunk_index as i64, content, embedding_blob.map(libsql::Value::Blob), ], ) .await .map_err(|e| WorkspaceError::ChunkingFailed { reason: format!("Insert failed: {}", e), })?; Ok(id) } async fn update_chunk_embedding( &self, chunk_id: Uuid, embedding: &[f32], ) -> Result<(), WorkspaceError> { let conn = self .connect() .await .map_err(|e| WorkspaceError::EmbeddingFailed { reason: e.to_string(), })?; let bytes: Vec = embedding.iter().flat_map(|f| f.to_le_bytes()).collect(); conn.execute( "UPDATE memory_chunks SET embedding = ?2 WHERE id = ?1", params![chunk_id.to_string(), libsql::Value::Blob(bytes)], ) .await .map_err(|e| WorkspaceError::EmbeddingFailed { reason: format!("Update failed: {}", e), })?; Ok(()) } async fn get_chunks_without_embeddings( &self, user_id: &str, agent_id: Option, limit: usize, ) -> Result, WorkspaceError> { let conn = self .connect() .await .map_err(|e| WorkspaceError::SearchFailed { reason: e.to_string(), })?; let agent_id_str = agent_id.map(|id| id.to_string()); let mut rows = conn .query( r#" SELECT c.id, c.document_id, c.chunk_index, c.content, c.created_at FROM memory_chunks c JOIN memory_documents d ON d.id = c.document_id WHERE d.user_id = ?1 AND d.agent_id IS ?2 AND c.embedding IS NULL LIMIT ?3 "#, params![user_id, agent_id_str.as_deref(), limit as i64], ) .await .map_err(|e| WorkspaceError::SearchFailed { reason: format!("Query failed: {}", e), })?; let mut chunks = Vec::new(); while let Some(row) = rows .next() .await .map_err(|e| WorkspaceError::SearchFailed { reason: format!("Query failed: {}", e), })? { chunks.push(MemoryChunk { id: get_text(&row, 0).parse().unwrap_or_default(), document_id: get_text(&row, 1).parse().unwrap_or_default(), chunk_index: get_i64(&row, 2) as i32, content: get_text(&row, 3), embedding: None, created_at: get_ts(&row, 4), }); } Ok(chunks) } async fn hybrid_search( &self, user_id: &str, agent_id: Option, query: &str, embedding: Option<&[f32]>, config: &SearchConfig, ) -> Result, WorkspaceError> { let conn = self .connect() .await .map_err(|e| WorkspaceError::SearchFailed { reason: e.to_string(), })?; let agent_id_str = agent_id.map(|id| id.to_string()); let pre_limit = config.pre_fusion_limit as i64; let fts_results = if config.use_fts { let mut rows = conn .query( r#" SELECT c.id, c.document_id, d.path, c.content FROM memory_chunks_fts fts JOIN memory_chunks c ON c._rowid = fts.rowid JOIN memory_documents d ON d.id = c.document_id WHERE d.user_id = ?1 AND d.agent_id IS ?2 AND memory_chunks_fts MATCH ?3 ORDER BY rank LIMIT ?4 "#, params![user_id, agent_id_str.as_deref(), query, pre_limit], ) .await .map_err(|e| WorkspaceError::SearchFailed { reason: format!("FTS query failed: {}", e), })?; let mut results = Vec::new(); while let Some(row) = rows .next() .await .map_err(|e| WorkspaceError::SearchFailed { reason: format!("FTS row fetch failed: {}", e), })? { results.push(RankedResult { chunk_id: get_text(&row, 0).parse().unwrap_or_default(), document_id: get_text(&row, 1).parse().unwrap_or_default(), document_path: get_text(&row, 2), content: get_text(&row, 3), rank: results.len() as u32 + 1, }); } results } else { Vec::new() }; let vector_results = if let (true, Some(emb)) = (config.use_vector, embedding) { let vector_json = format!( "[{}]", emb.iter() .map(|f| f.to_string()) .collect::>() .join(",") ); // vector_top_k requires a libsql_vector_idx index. After the V9 // migration the index is dropped (to support flexible embedding // dimensions), so this query may fail. Fall back to FTS-only. match conn .query( r#" SELECT c.id, c.document_id, d.path, c.content FROM vector_top_k('idx_memory_chunks_embedding', vector(?1), ?2) AS top_k JOIN memory_chunks c ON c._rowid = top_k.id JOIN memory_documents d ON d.id = c.document_id WHERE d.user_id = ?3 AND d.agent_id IS ?4 "#, params![vector_json, pre_limit, user_id, agent_id_str.as_deref()], ) .await { Ok(mut rows) => { let mut results = Vec::new(); while let Some(row) = rows.next() .await .map_err(|e| WorkspaceError::SearchFailed { reason: format!("Vector row fetch failed: {}", e), })? { results.push(RankedResult { chunk_id: get_text(&row, 0).parse().unwrap_or_default(), document_id: get_text(&row, 1).parse().unwrap_or_default(), document_path: get_text(&row, 2), content: get_text(&row, 3), rank: results.len() as u32 + 1, }); } results } Err(e) => { tracing::debug!( "Vector index query failed (expected after V9 migration), \ falling back to FTS-only: {e}" ); Vec::new() } } } else { Vec::new() }; if embedding.is_some() && !config.use_vector { tracing::warn!( "Embedding provided but vector search is disabled in config; using FTS-only results" ); } Ok(fuse_results(fts_results, vector_results, config)) } }