mirror of
https://github.com/outbackdingo/optimclaw.git
synced 2026-08-26 23:50:17 +00:00
ce5961b1ec
* fix(libsql): support flexible embedding dimensions (#494) The libSQL schema hardcoded F32_BLOB(1536) for the embedding column, preventing use of models with other dimensions (e.g. 768-dim nomic-embed-text). This adds incremental migration support to the libSQL backend and a V9 migration that rebuilds the memory_chunks table with a plain BLOB column accepting any dimension. - Add incremental migration infrastructure (INCREMENTAL_MIGRATIONS array + run_incremental() runner tracked via _migrations table) - V9 migration rebuilds memory_chunks with BLOB column, drops the vector index (which requires fixed-dimension F32_BLOB) - Update base schema for fresh installs (BLOB, no vector index) - Vector search gracefully falls back to FTS-only when the index is absent (matches PostgreSQL behavior after its V9 migration) - Remove now-incorrect "dimension is not 1536" warnings Existing embeddings are preserved during migration. Users only need to re-embed if they change their embedding model/dimension. [skip-regression-check] Co-Authored-By: Claude Opus 4.6 <[email protected]> * fix: wrap incremental migrations in transaction for atomicity Address PR review feedback: if the process crashes after executing migration SQL but before recording it in _migrations, the migration would be applied but not marked complete. Wrapping both operations in a transaction ensures they succeed or fail together. [skip-regression-check] Co-Authored-By: Claude Opus 4.6 <[email protected]> * chore: merge main and fix formatting drift Co-Authored-By: Claude Opus 4.6 <[email protected]> --------- Co-authored-by: Claude Opus 4.6 <[email protected]>
620 lines
20 KiB
Rust
620 lines
20 KiB
Rust
//! 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,
|
|
reciprocal_rank_fusion,
|
|
};
|
|
|
|
use chrono::Utc;
|
|
|
|
#[async_trait]
|
|
impl WorkspaceStore for LibSqlBackend {
|
|
async fn get_document_by_path(
|
|
&self,
|
|
user_id: &str,
|
|
agent_id: Option<Uuid>,
|
|
path: &str,
|
|
) -> Result<MemoryDocument, 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 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<MemoryDocument, WorkspaceError> {
|
|
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<Uuid>,
|
|
path: &str,
|
|
) -> Result<MemoryDocument, WorkspaceError> {
|
|
// 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<Uuid>,
|
|
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<Uuid>,
|
|
directory: &str,
|
|
) -> Result<Vec<WorkspaceEntry>, 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<String, WorkspaceEntry> = 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<WorkspaceEntry> = 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<Uuid>,
|
|
) -> Result<Vec<String>, 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<Uuid>,
|
|
) -> Result<Vec<MemoryDocument>, 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<Uuid, WorkspaceError> {
|
|
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<u8> = 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<u8> = 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<Uuid>,
|
|
limit: usize,
|
|
) -> Result<Vec<MemoryChunk>, 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<Uuid>,
|
|
query: &str,
|
|
embedding: Option<&[f32]>,
|
|
config: &SearchConfig,
|
|
) -> Result<Vec<SearchResult>, 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::<Vec<_>>()
|
|
.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(reciprocal_rank_fusion(fts_results, vector_results, config))
|
|
}
|
|
}
|