feat: add LanceDB support for workspace semantic search

- Introduced optional LanceDB vector store for semantic search, configurable via environment variables.
- Updated `.env.example` and `Cargo.toml` to include LanceDB settings.
- Enhanced `DatabaseConfig` to support vector backend selection and LanceDB path configuration.
- Implemented `VectorBackend` enum to manage vector store options.
- Added functionality to connect to LanceDB in the database connection logic.
- Updated relevant documentation to reflect new features and configuration options.

This change allows users to leverage LanceDB as an alternative to pgvector/libsql for improved search capabilities.
This commit is contained in:
ILGIN KANAT
2026-02-18 11:06:09 +04:00
parent 96d5fc0d39
commit 327e009622
15 changed files with 4661 additions and 18 deletions
+6
View File
@@ -2,6 +2,12 @@
DATABASE_URL=postgres://localhost/ironclaw
DATABASE_POOL_SIZE=10
# Vector store for workspace memory (optional)
# When set to "lancedb", uses LanceDB for semantic search instead of pgvector/libsql
# VECTOR_BACKEND=pgvector # default: use database's built-in (pgvector or libsql_vector_idx)
# VECTOR_BACKEND=lancedb
# LANCEDB_PATH=~/.ironclaw/lancedb # path for LanceDB when VECTOR_BACKEND=lancedb
# LLM Provider (NEAR AI)
# NEAR AI provides a unified interface to all models with user authentication
# Session token is stored in ~/.ironclaw/session.json and managed automatically.
+6
View File
@@ -401,6 +401,12 @@ LIBSQL_PATH=~/.ironclaw/ironclaw.db # Default path
# libSQL (Turso cloud sync)
LIBSQL_URL=libsql://your-db.turso.io
LIBSQL_AUTH_TOKEN=your-token # Required when LIBSQL_URL is set
# Vector store for workspace semantic search (optional)
# When lancedb: uses LanceDB instead of pgvector/libsql for vector search
VECTOR_BACKEND=builtin # default: use database built-in
# VECTOR_BACKEND=lancedb # requires: cargo build --features lancedb
# LANCEDB_PATH=~/.ironclaw/lancedb # default when VECTOR_BACKEND=lancedb
```
### Current Limitations (libSQL backend)
Generated
+2967 -10
View File
File diff suppressed because it is too large Load Diff
+4
View File
@@ -108,6 +108,9 @@ open = "5"
# The postgres feature provides ToSql/FromSql for postgres-types (shared by tokio-postgres)
pgvector = { version = "0.4", features = ["postgres"], optional = true }
# LanceDB vector store (optional alternative to pgvector/libsql for workspace search)
lancedb = { version = "0.26", optional = true }
# WASM sandbox for untrusted tool execution
wasmtime = { version = "28", features = ["component-model"] }
wasmtime-wasi = "28" # WASI support for component model
@@ -153,6 +156,7 @@ tempfile = "3"
[features]
default = ["postgres", "libsql"]
lancedb = ["dep:lancedb"]
postgres = [
"dep:deadpool-postgres",
"dep:tokio-postgres",
+1 -1
View File
@@ -251,7 +251,7 @@ This document tracks feature parity between IronClaw (Rust implementation) and O
| Gemini embeddings | ✅ | ❌ | |
| Local embeddings | ✅ | ❌ | |
| SQLite-vec backend | ✅ | ❌ | IronClaw uses PostgreSQL |
| LanceDB backend | ✅ | | |
| LanceDB backend | ✅ | | VectorStore + LanceDbVectorStore, DbWithLanceVectorStore wrapper, VECTOR_BACKEND=lancedb |
| QMD backend | ✅ | ❌ | |
| Atomic reindexing | ✅ | ✅ | |
| Embeddings batching | ✅ | ❌ | |
+81
View File
@@ -310,6 +310,33 @@ impl std::str::FromStr for DatabaseBackend {
}
}
/// Which vector store to use for workspace semantic search.
///
/// When None, uses the database backend's built-in (pgvector or libsql_vector_idx).
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum VectorBackend {
/// Use database built-in (pgvector or libsql_vector_idx).
#[default]
Builtin,
/// Use LanceDB for vector search (requires feature "lancedb").
LanceDb,
}
impl std::str::FromStr for VectorBackend {
type Err = String;
fn from_str(s: &str) -> Result<Self, Self::Err> {
match s.to_lowercase().as_str() {
"builtin" | "pgvector" | "libsql" | "" => Ok(Self::Builtin),
"lancedb" | "lance" => Ok(Self::LanceDb),
_ => Err(format!(
"invalid vector backend '{}', expected 'builtin' or 'lancedb'",
s
)),
}
}
}
/// Database configuration.
#[derive(Debug, Clone)]
pub struct DatabaseConfig {
@@ -327,6 +354,12 @@ pub struct DatabaseConfig {
pub libsql_url: Option<String>,
/// Turso auth token (required when libsql_url is set).
pub libsql_auth_token: Option<SecretString>,
// -- Vector store (workspace semantic search) --
/// Override vector backend (default: use database built-in).
pub vector_backend: VectorBackend,
/// Path for LanceDB when vector_backend=LanceDb (default: ~/.ironclaw/lancedb).
pub lancedb_path: Option<PathBuf>,
}
impl DatabaseConfig {
@@ -376,6 +409,19 @@ impl DatabaseConfig {
});
}
let vector_backend: VectorBackend =
optional_env("VECTOR_BACKEND")?
.and_then(|s| s.parse().ok())
.unwrap_or_default();
let lancedb_path = optional_env("LANCEDB_PATH")?.map(PathBuf::from).or_else(|| {
if vector_backend == VectorBackend::LanceDb {
Some(default_lancedb_path())
} else {
None
}
});
Ok(Self {
backend,
url: SecretString::from(url),
@@ -383,6 +429,8 @@ impl DatabaseConfig {
libsql_path,
libsql_url,
libsql_auth_token,
vector_backend,
lancedb_path,
})
}
@@ -400,6 +448,39 @@ pub fn default_libsql_path() -> PathBuf {
.join("ironclaw.db")
}
/// Default LanceDB path (~/.ironclaw/lancedb).
pub fn default_lancedb_path() -> PathBuf {
dirs::home_dir()
.unwrap_or_else(|| PathBuf::from("."))
.join(".ironclaw")
.join("lancedb")
}
#[cfg(test)]
mod tests {
use super::VectorBackend;
#[test]
fn test_vector_backend_parse() {
assert_eq!("builtin".parse::<VectorBackend>().unwrap(), VectorBackend::Builtin);
assert_eq!("pgvector".parse::<VectorBackend>().unwrap(), VectorBackend::Builtin);
assert_eq!("libsql".parse::<VectorBackend>().unwrap(), VectorBackend::Builtin);
assert_eq!("".parse::<VectorBackend>().unwrap(), VectorBackend::Builtin);
assert_eq!("lancedb".parse::<VectorBackend>().unwrap(), VectorBackend::LanceDb);
assert_eq!("lance".parse::<VectorBackend>().unwrap(), VectorBackend::LanceDb);
assert_eq!("Lancedb".parse::<VectorBackend>().unwrap(), VectorBackend::LanceDb);
assert!("invalid".parse::<VectorBackend>().is_err());
}
#[test]
fn test_default_lancedb_path() {
let path = super::default_lancedb_path();
assert!(path.to_string_lossy().ends_with(".ironclaw/lancedb"));
}
}
/// Which LLM backend to use.
///
/// Defaults to `NearAi` to keep IronClaw close to the NEAR ecosystem.
+664
View File
@@ -0,0 +1,664 @@
//! Database wrapper that uses LanceDB for vector search.
//!
//! When `VECTOR_BACKEND=lancedb`, this wraps the main database (Postgres or libSQL)
//! and delegates vector search to LanceDB while using the inner DB for FTS and all
//! other operations.
use std::sync::Arc;
use async_trait::async_trait;
use uuid::Uuid;
use crate::agent::BrokenTool;
use crate::agent::routine::{Routine, RoutineRun, RunStatus};
use crate::context::{ActionRecord, JobContext, JobState};
use crate::db::Database;
use crate::error::{DatabaseError, WorkspaceError};
use crate::history::{
ConversationMessage, ConversationSummary, JobEventRecord, LlmCallRecord, SandboxJobRecord,
SandboxJobSummary, SettingRow,
};
use crate::workspace::{
MemoryChunk, MemoryDocument, SearchConfig, SearchResult, WorkspaceEntry,
};
use crate::workspace::search::{RankedResult, reciprocal_rank_fusion};
use crate::workspace::lancedb_store::VectorStore;
/// Wraps a Database with a LanceDB VectorStore for hybrid search.
///
/// Documents and chunks stay in the inner DB. Vector search uses LanceDB.
pub struct DbWithLanceVectorStore {
inner: Arc<dyn Database>,
vector_store: Arc<dyn VectorStore>,
}
impl DbWithLanceVectorStore {
pub fn new(inner: Arc<dyn Database>, vector_store: Arc<dyn VectorStore>) -> Self {
Self { inner, vector_store }
}
}
#[async_trait]
impl Database for DbWithLanceVectorStore {
async fn run_migrations(&self) -> Result<(), DatabaseError> {
self.inner.run_migrations().await
}
async fn create_conversation(
&self,
channel: &str,
user_id: &str,
thread_id: Option<&str>,
) -> Result<Uuid, DatabaseError> {
self.inner
.create_conversation(channel, user_id, thread_id)
.await
}
async fn touch_conversation(&self, id: Uuid) -> Result<(), DatabaseError> {
self.inner.touch_conversation(id).await
}
async fn add_conversation_message(
&self,
conversation_id: Uuid,
role: &str,
content: &str,
) -> Result<Uuid, DatabaseError> {
self.inner
.add_conversation_message(conversation_id, role, content)
.await
}
async fn ensure_conversation(
&self,
id: Uuid,
channel: &str,
user_id: &str,
thread_id: Option<&str>,
) -> Result<(), DatabaseError> {
self.inner
.ensure_conversation(id, channel, user_id, thread_id)
.await
}
async fn list_conversations_with_preview(
&self,
user_id: &str,
channel: &str,
limit: i64,
) -> Result<Vec<ConversationSummary>, DatabaseError> {
self.inner
.list_conversations_with_preview(user_id, channel, limit)
.await
}
async fn get_or_create_assistant_conversation(
&self,
user_id: &str,
channel: &str,
) -> Result<Uuid, DatabaseError> {
self.inner
.get_or_create_assistant_conversation(user_id, channel)
.await
}
async fn create_conversation_with_metadata(
&self,
user_id: &str,
channel: &str,
metadata: &serde_json::Value,
) -> Result<Uuid, DatabaseError> {
self.inner
.create_conversation_with_metadata(user_id, channel, metadata)
.await
}
async fn list_conversation_messages_paginated(
&self,
conversation_id: Uuid,
before: Option<chrono::DateTime<chrono::Utc>>,
limit: i64,
) -> Result<(Vec<ConversationMessage>, bool), DatabaseError> {
self.inner
.list_conversation_messages_paginated(conversation_id, before, limit)
.await
}
async fn update_conversation_metadata_field(
&self,
id: Uuid,
key: &str,
value: &serde_json::Value,
) -> Result<(), DatabaseError> {
self.inner
.update_conversation_metadata_field(id, key, value)
.await
}
async fn get_conversation_metadata(
&self,
id: Uuid,
) -> Result<Option<serde_json::Value>, DatabaseError> {
self.inner.get_conversation_metadata(id).await
}
async fn list_conversation_messages(
&self,
conversation_id: Uuid,
) -> Result<Vec<ConversationMessage>, DatabaseError> {
self.inner.list_conversation_messages(conversation_id).await
}
async fn conversation_belongs_to_user(
&self,
conversation_id: Uuid,
user_id: &str,
) -> Result<bool, DatabaseError> {
self.inner
.conversation_belongs_to_user(conversation_id, user_id)
.await
}
async fn save_job(&self, ctx: &JobContext) -> Result<(), DatabaseError> {
self.inner.save_job(ctx).await
}
async fn get_job(&self, id: Uuid) -> Result<Option<JobContext>, DatabaseError> {
self.inner.get_job(id).await
}
async fn update_job_status(
&self,
id: Uuid,
status: crate::context::JobState,
failure_reason: Option<&str>,
) -> Result<(), DatabaseError> {
self.inner.update_job_status(id, status, failure_reason).await
}
async fn mark_job_stuck(&self, id: Uuid) -> Result<(), DatabaseError> {
self.inner.mark_job_stuck(id).await
}
async fn get_stuck_jobs(&self) -> Result<Vec<Uuid>, DatabaseError> {
self.inner.get_stuck_jobs().await
}
async fn save_action(&self, job_id: Uuid, action: &ActionRecord) -> Result<(), DatabaseError> {
self.inner.save_action(job_id, action).await
}
async fn get_job_actions(&self, job_id: Uuid) -> Result<Vec<ActionRecord>, DatabaseError> {
self.inner.get_job_actions(job_id).await
}
async fn record_llm_call(&self, record: &LlmCallRecord<'_>) -> Result<Uuid, DatabaseError> {
self.inner.record_llm_call(record).await
}
async fn save_estimation_snapshot(
&self,
job_id: Uuid,
category: &str,
tool_names: &[String],
estimated_cost: rust_decimal::Decimal,
estimated_time_secs: i32,
estimated_value: rust_decimal::Decimal,
) -> Result<Uuid, DatabaseError> {
self.inner
.save_estimation_snapshot(
job_id,
category,
tool_names,
estimated_cost,
estimated_time_secs,
estimated_value,
)
.await
}
async fn update_estimation_actuals(
&self,
id: Uuid,
actual_cost: rust_decimal::Decimal,
actual_time_secs: i32,
actual_value: Option<rust_decimal::Decimal>,
) -> Result<(), DatabaseError> {
self.inner
.update_estimation_actuals(id, actual_cost, actual_time_secs, actual_value)
.await
}
async fn save_sandbox_job(&self, job: &SandboxJobRecord) -> Result<(), DatabaseError> {
self.inner.save_sandbox_job(job).await
}
async fn get_sandbox_job(&self, id: Uuid) -> Result<Option<SandboxJobRecord>, DatabaseError> {
self.inner.get_sandbox_job(id).await
}
async fn list_sandbox_jobs(&self) -> Result<Vec<SandboxJobRecord>, DatabaseError> {
self.inner.list_sandbox_jobs().await
}
async fn update_sandbox_job_status(
&self,
id: Uuid,
status: &str,
success: Option<bool>,
message: Option<&str>,
started_at: Option<chrono::DateTime<chrono::Utc>>,
completed_at: Option<chrono::DateTime<chrono::Utc>>,
) -> Result<(), DatabaseError> {
self.inner
.update_sandbox_job_status(id, status, success, message, started_at, completed_at)
.await
}
async fn cleanup_stale_sandbox_jobs(&self) -> Result<u64, DatabaseError> {
self.inner.cleanup_stale_sandbox_jobs().await
}
async fn sandbox_job_summary(&self) -> Result<SandboxJobSummary, DatabaseError> {
self.inner.sandbox_job_summary().await
}
async fn list_sandbox_jobs_for_user(
&self,
user_id: &str,
) -> Result<Vec<SandboxJobRecord>, DatabaseError> {
self.inner.list_sandbox_jobs_for_user(user_id).await
}
async fn sandbox_job_summary_for_user(
&self,
user_id: &str,
) -> Result<SandboxJobSummary, DatabaseError> {
self.inner.sandbox_job_summary_for_user(user_id).await
}
async fn sandbox_job_belongs_to_user(
&self,
job_id: Uuid,
user_id: &str,
) -> Result<bool, DatabaseError> {
self.inner
.sandbox_job_belongs_to_user(job_id, user_id)
.await
}
async fn update_sandbox_job_mode(&self, id: Uuid, mode: &str) -> Result<(), DatabaseError> {
self.inner.update_sandbox_job_mode(id, mode).await
}
async fn get_sandbox_job_mode(&self, id: Uuid) -> Result<Option<String>, DatabaseError> {
self.inner.get_sandbox_job_mode(id).await
}
async fn save_job_event(
&self,
job_id: Uuid,
event_type: &str,
data: &serde_json::Value,
) -> Result<(), DatabaseError> {
self.inner.save_job_event(job_id, event_type, data).await
}
async fn list_job_events(
&self,
job_id: Uuid,
limit: Option<i64>,
) -> Result<Vec<crate::history::JobEventRecord>, DatabaseError> {
self.inner.list_job_events(job_id, limit).await
}
async fn create_routine(&self, routine: &crate::agent::routine::Routine) -> Result<(), DatabaseError> {
self.inner.create_routine(routine).await
}
async fn get_routine(&self, id: Uuid) -> Result<Option<crate::agent::routine::Routine>, DatabaseError> {
self.inner.get_routine(id).await
}
async fn get_routine_by_name(
&self,
user_id: &str,
name: &str,
) -> Result<Option<crate::agent::routine::Routine>, DatabaseError> {
self.inner.get_routine_by_name(user_id, name).await
}
async fn list_routines(
&self,
user_id: &str,
) -> Result<Vec<crate::agent::routine::Routine>, DatabaseError> {
self.inner.list_routines(user_id).await
}
async fn list_event_routines(&self) -> Result<Vec<crate::agent::routine::Routine>, DatabaseError> {
self.inner.list_event_routines().await
}
async fn list_due_cron_routines(&self) -> Result<Vec<crate::agent::routine::Routine>, DatabaseError> {
self.inner.list_due_cron_routines().await
}
async fn update_routine(&self, routine: &crate::agent::routine::Routine) -> Result<(), DatabaseError> {
self.inner.update_routine(routine).await
}
async fn update_routine_runtime(
&self,
id: Uuid,
last_run_at: chrono::DateTime<chrono::Utc>,
next_fire_at: Option<chrono::DateTime<chrono::Utc>>,
run_count: u64,
consecutive_failures: u32,
state: &serde_json::Value,
) -> Result<(), DatabaseError> {
self.inner
.update_routine_runtime(
id,
last_run_at,
next_fire_at,
run_count,
consecutive_failures,
state,
)
.await
}
async fn delete_routine(&self, id: Uuid) -> Result<bool, DatabaseError> {
self.inner.delete_routine(id).await
}
async fn create_routine_run(&self, run: &RoutineRun) -> Result<Uuid, DatabaseError> {
self.inner.create_routine_run(run).await
}
async fn complete_routine_run(
&self,
id: Uuid,
status: RunStatus,
result_summary: Option<&str>,
tokens_used: Option<i32>,
) -> Result<(), DatabaseError> {
self.inner
.complete_routine_run(id, status, result_summary, tokens_used)
.await
}
async fn list_routine_runs(
&self,
routine_id: Uuid,
limit: i64,
) -> Result<Vec<RoutineRun>, DatabaseError> {
self.inner.list_routine_runs(routine_id, limit).await
}
async fn count_running_routine_runs(&self, routine_id: Uuid) -> Result<i64, DatabaseError> {
self.inner.count_running_routine_runs(routine_id).await
}
async fn record_tool_failure(
&self,
tool_name: &str,
error_message: &str,
) -> Result<(), DatabaseError> {
self.inner.record_tool_failure(tool_name, error_message).await
}
async fn get_broken_tools(&self, threshold: i32) -> Result<Vec<BrokenTool>, DatabaseError> {
self.inner.get_broken_tools(threshold).await
}
async fn mark_tool_repaired(&self, tool_name: &str) -> Result<(), DatabaseError> {
self.inner.mark_tool_repaired(tool_name).await
}
async fn increment_repair_attempts(&self, tool_name: &str) -> Result<(), DatabaseError> {
self.inner.increment_repair_attempts(tool_name).await
}
async fn get_setting(
&self,
user_id: &str,
key: &str,
) -> Result<Option<serde_json::Value>, DatabaseError> {
self.inner.get_setting(user_id, key).await
}
async fn get_setting_full(
&self,
user_id: &str,
key: &str,
) -> Result<Option<SettingRow>, DatabaseError> {
self.inner.get_setting_full(user_id, key).await
}
async fn set_setting(
&self,
user_id: &str,
key: &str,
value: &serde_json::Value,
) -> Result<(), DatabaseError> {
self.inner.set_setting(user_id, key, value).await
}
async fn delete_setting(&self, user_id: &str, key: &str) -> Result<bool, DatabaseError> {
self.inner.delete_setting(user_id, key).await
}
async fn list_settings(&self, user_id: &str) -> Result<Vec<SettingRow>, DatabaseError> {
self.inner.list_settings(user_id).await
}
async fn get_all_settings(
&self,
user_id: &str,
) -> Result<std::collections::HashMap<String, serde_json::Value>, DatabaseError> {
self.inner.get_all_settings(user_id).await
}
async fn set_all_settings(
&self,
user_id: &str,
settings: &std::collections::HashMap<String, serde_json::Value>,
) -> Result<(), DatabaseError> {
self.inner.set_all_settings(user_id, settings).await
}
async fn has_settings(&self, user_id: &str) -> Result<bool, DatabaseError> {
self.inner.has_settings(user_id).await
}
async fn get_document_by_path(
&self,
user_id: &str,
agent_id: Option<Uuid>,
path: &str,
) -> Result<MemoryDocument, WorkspaceError> {
self.inner
.get_document_by_path(user_id, agent_id, path)
.await
}
async fn get_document_by_id(&self, id: Uuid) -> Result<MemoryDocument, WorkspaceError> {
self.inner.get_document_by_id(id).await
}
async fn get_or_create_document_by_path(
&self,
user_id: &str,
agent_id: Option<Uuid>,
path: &str,
) -> Result<MemoryDocument, WorkspaceError> {
self.inner
.get_or_create_document_by_path(user_id, agent_id, path)
.await
}
async fn update_document(&self, id: Uuid, content: &str) -> Result<(), WorkspaceError> {
self.inner.update_document(id, content).await
}
async fn delete_document_by_path(
&self,
user_id: &str,
agent_id: Option<Uuid>,
path: &str,
) -> Result<(), WorkspaceError> {
self.inner
.delete_document_by_path(user_id, agent_id, path)
.await
}
async fn list_directory(
&self,
user_id: &str,
agent_id: Option<Uuid>,
directory: &str,
) -> Result<Vec<WorkspaceEntry>, WorkspaceError> {
self.inner
.list_directory(user_id, agent_id, directory)
.await
}
async fn list_all_paths(
&self,
user_id: &str,
agent_id: Option<Uuid>,
) -> Result<Vec<String>, WorkspaceError> {
self.inner.list_all_paths(user_id, agent_id).await
}
async fn list_documents(
&self,
user_id: &str,
agent_id: Option<Uuid>,
) -> Result<Vec<MemoryDocument>, WorkspaceError> {
self.inner.list_documents(user_id, agent_id).await
}
async fn delete_chunks(&self, document_id: Uuid) -> Result<(), WorkspaceError> {
self.inner.delete_chunks(document_id).await?;
self.vector_store.delete_chunks(document_id).await?;
Ok(())
}
async fn insert_chunk(
&self,
document_id: Uuid,
chunk_index: i32,
content: &str,
embedding: Option<&[f32]>,
) -> Result<Uuid, WorkspaceError> {
let chunk_id = self
.inner
.insert_chunk(document_id, chunk_index, content, embedding)
.await?;
if let Some(emb) = embedding {
let doc = self.inner.get_document_by_id(document_id).await?;
self.vector_store
.insert_chunk(
chunk_id,
document_id,
&doc.user_id,
doc.agent_id,
content,
emb,
)
.await?;
}
Ok(chunk_id)
}
async fn update_chunk_embedding(
&self,
chunk_id: Uuid,
embedding: &[f32],
) -> Result<(), WorkspaceError> {
self.inner.update_chunk_embedding(chunk_id, embedding).await?;
if let Some(chunk) = self.inner.get_chunk_by_id(chunk_id).await? {
let doc = self.inner.get_document_by_id(chunk.document_id).await?;
self.vector_store
.update_chunk_embedding(
chunk_id,
chunk.document_id,
&doc.user_id,
doc.agent_id,
&chunk.content,
embedding,
)
.await?;
}
Ok(())
}
async fn get_chunks_without_embeddings(
&self,
user_id: &str,
agent_id: Option<Uuid>,
limit: usize,
) -> Result<Vec<MemoryChunk>, WorkspaceError> {
self.inner
.get_chunks_without_embeddings(user_id, agent_id, limit)
.await
}
async fn get_chunk_by_id(
&self,
chunk_id: Uuid,
) -> Result<Option<MemoryChunk>, WorkspaceError> {
self.inner.get_chunk_by_id(chunk_id).await
}
async fn hybrid_search(
&self,
user_id: &str,
agent_id: Option<Uuid>,
query: &str,
embedding: Option<&[f32]>,
config: &SearchConfig,
) -> Result<Vec<SearchResult>, WorkspaceError> {
let fts_config = SearchConfig {
use_fts: true,
use_vector: false,
..*config
};
let fts_search_results = self
.inner
.hybrid_search(user_id, agent_id, query, None, &fts_config)
.await?;
let fts_results: Vec<RankedResult> = fts_search_results
.iter()
.enumerate()
.map(|(i, r)| RankedResult {
chunk_id: r.chunk_id,
document_id: r.document_id,
content: r.content.clone(),
rank: (i + 1) as u32,
})
.collect();
let vector_results = if config.use_vector {
if let Some(emb) = embedding {
self.vector_store
.vector_search(user_id, agent_id, emb, config.pre_fusion_limit)
.await?
} else {
Vec::new()
}
} else {
Vec::new()
};
Ok(reciprocal_rank_fusion(fts_results, vector_results, config))
}
}
+40
View File
@@ -2486,6 +2486,46 @@ impl Database for LibSqlBackend {
Ok(chunks)
}
async fn get_chunk_by_id(
&self,
chunk_id: Uuid,
) -> Result<Option<MemoryChunk>, WorkspaceError> {
let conn = self
.connect()
.await
.map_err(|e| WorkspaceError::SearchFailed {
reason: e.to_string(),
})?;
let mut rows = conn
.query(
r#"
SELECT id, document_id, chunk_index, content, created_at
FROM memory_chunks WHERE id = ?1
"#,
params![chunk_id.to_string()],
)
.await
.map_err(|e| WorkspaceError::SearchFailed {
reason: format!("Query failed: {}", e),
})?;
let row = rows
.next()
.await
.map_err(|e| WorkspaceError::SearchFailed {
reason: format!("Row fetch failed: {}", e),
})?;
Ok(row.map(|r| MemoryChunk {
id: get_text(&r, 0).parse().unwrap_or_default(),
document_id: get_text(&r, 1).parse().unwrap_or_default(),
chunk_index: get_i64(&r, 2) as i32,
content: get_text(&r, 3),
embedding: None,
created_at: get_ts(&r, 4),
}))
}
// ==================== Workspace: Search ====================
async fn hybrid_search(
+32 -6
View File
@@ -18,6 +18,9 @@ pub mod libsql_backend;
#[cfg(feature = "libsql")]
pub mod libsql_migrations;
#[cfg(all(feature = "lancedb", any(feature = "postgres", feature = "libsql")))]
pub mod lancedb_wrapper;
use std::collections::HashMap;
use std::sync::Arc;
@@ -48,7 +51,7 @@ use crate::workspace::{SearchConfig, SearchResult};
pub async fn connect_from_config(
config: &crate::config::DatabaseConfig,
) -> Result<Arc<dyn Database>, DatabaseError> {
match config.backend {
let inner: Arc<dyn Database> = match config.backend {
#[cfg(feature = "libsql")]
crate::config::DatabaseBackend::LibSql => {
use secrecy::ExposeSecret as _;
@@ -75,7 +78,7 @@ pub async fn connect_from_config(
.map_err(|e| DatabaseError::Pool(e.to_string()))?
};
backend.run_migrations().await?;
Ok(Arc::new(backend))
Arc::new(backend)
}
#[cfg(feature = "postgres")]
_ => {
@@ -83,13 +86,33 @@ pub async fn connect_from_config(
.await
.map_err(|e| DatabaseError::Pool(e.to_string()))?;
pg.run_migrations().await?;
Ok(Arc::new(pg))
Arc::new(pg)
}
#[cfg(not(feature = "postgres"))]
_ => Err(DatabaseError::Pool(
"No database backend available. Enable 'postgres' or 'libsql' feature.".to_string(),
)),
_ => {
return Err(DatabaseError::Pool(
"No database backend available. Enable 'postgres' or 'libsql' feature."
.to_string(),
));
}
};
#[cfg(feature = "lancedb")]
if config.vector_backend == crate::config::VectorBackend::LanceDb {
let path = config
.lancedb_path
.clone()
.unwrap_or_else(crate::config::default_lancedb_path);
let store = crate::workspace::LanceDbVectorStore::new(path)
.await
.map_err(|e| DatabaseError::Pool(format!("LanceDB: {}", e)))?;
return Ok(Arc::new(lancedb_wrapper::DbWithLanceVectorStore::new(
inner,
Arc::new(store),
)) as Arc<dyn Database>);
}
Ok(inner)
}
/// Backend-agnostic database trait.
@@ -528,6 +551,9 @@ pub trait Database: Send + Sync {
limit: usize,
) -> Result<Vec<MemoryChunk>, WorkspaceError>;
/// Get a chunk by ID (for LanceDB update flow).
async fn get_chunk_by_id(&self, chunk_id: Uuid) -> Result<Option<MemoryChunk>, WorkspaceError>;
// ==================== Workspace: Search ====================
/// Perform hybrid search combining FTS and vector similarity.
+7
View File
@@ -614,6 +614,13 @@ impl Database for PgBackend {
.await
}
async fn get_chunk_by_id(
&self,
chunk_id: Uuid,
) -> Result<Option<MemoryChunk>, WorkspaceError> {
self.repo.get_chunk_by_id(chunk_id).await
}
// ==================== Workspace: Search ====================
async fn hybrid_search(
+37 -1
View File
@@ -386,7 +386,7 @@ async fn main() -> anyhow::Result<()> {
#[cfg(feature = "libsql")]
let mut libsql_db: Option<std::sync::Arc<libsql::Database>> = None;
let db: Option<Arc<dyn ironclaw::db::Database>> = if cli.no_db {
let mut inner_db: Option<Arc<dyn ironclaw::db::Database>> = if cli.no_db {
tracing::warn!("Running without database connection");
None
} else {
@@ -443,6 +443,42 @@ async fn main() -> anyhow::Result<()> {
}
};
// Optionally wrap with LanceDB vector store for workspace semantic search
let db: Option<Arc<dyn ironclaw::db::Database>> = if let Some(inner) = inner_db.take() {
#[cfg(feature = "lancedb")]
{
if config.database.vector_backend == ironclaw::config::VectorBackend::LanceDb {
let path = config
.database
.lancedb_path
.clone()
.unwrap_or_else(ironclaw::config::default_lancedb_path);
let store = ironclaw::workspace::LanceDbVectorStore::new(path)
.await
.map_err(|e| anyhow::anyhow!("LanceDB: {}", e))?;
tracing::info!("LanceDB vector store connected for workspace search");
Some(Arc::new(ironclaw::db::lancedb_wrapper::DbWithLanceVectorStore::new(
inner,
Arc::new(store),
)) as Arc<dyn ironclaw::db::Database>)
} else {
Some(inner)
}
}
#[cfg(not(feature = "lancedb"))]
{
if config.database.vector_backend == ironclaw::config::VectorBackend::LanceDb {
anyhow::bail!(
"VECTOR_BACKEND=lancedb requires the 'lancedb' feature. Build with: cargo build --features lancedb"
);
} else {
Some(inner)
}
}
} else {
None
};
// Post-init operations using the database
if let Some(ref db) = db {
// One-time migration: move disk config files into the DB settings table.
+623
View File
@@ -0,0 +1,623 @@
//! LanceDB-backed vector store for workspace memory chunks.
//!
//! Provides an alternative to pgvector/libsql for semantic search when the
//! `lancedb` feature is enabled. Documents and metadata stay in the main
//! database; this store holds chunk embeddings for vector similarity search.
//!
//! Configuration:
//! LANCEDB_PATH=~/.ironclaw/lancedb # Default
//! VECTOR_BACKEND=lancedb # Use LanceDB for vector search
use std::sync::Arc;
use async_trait::async_trait;
use uuid::Uuid;
use crate::error::WorkspaceError;
use crate::workspace::search::RankedResult;
/// Embedding dimension (text-embedding-3-small default).
/// Must match the embedding model used.
pub const DEFAULT_EMBEDDING_DIM: i32 = 1536;
/// Vector store abstraction for semantic search.
///
/// Implementations: pgvector/libsql (embedded in Database), LanceDB (this module).
#[async_trait]
pub trait VectorStore: Send + Sync {
/// Insert a chunk with its embedding.
async fn insert_chunk(
&self,
chunk_id: Uuid,
document_id: Uuid,
user_id: &str,
agent_id: Option<Uuid>,
content: &str,
embedding: &[f32],
) -> Result<(), WorkspaceError>;
/// Update an existing chunk's embedding.
///
/// For LanceDB, this performs delete+insert since LanceDB has limited update
/// support. Caller must provide full chunk metadata for the re-insert.
async fn update_chunk_embedding(
&self,
chunk_id: Uuid,
document_id: Uuid,
user_id: &str,
agent_id: Option<Uuid>,
content: &str,
embedding: &[f32],
) -> Result<(), WorkspaceError>;
/// Delete all chunks for a document.
async fn delete_chunks(&self, document_id: Uuid) -> Result<(), WorkspaceError>;
/// Vector similarity search, filtered by user and agent.
async fn vector_search(
&self,
user_id: &str,
agent_id: Option<Uuid>,
embedding: &[f32],
limit: usize,
) -> Result<Vec<RankedResult>, WorkspaceError>;
}
#[cfg(feature = "lancedb")]
mod impl_lancedb {
use std::sync::Arc;
use arrow_array::types::Float32Type;
use arrow_array::{
Array, FixedSizeListArray, RecordBatch, RecordBatchIterator, StringArray,
};
use arrow_schema::{DataType, Field, Schema};
use async_trait::async_trait;
use futures::StreamExt;
use lancedb::index::Index;
use uuid::Uuid;
use super::{RankedResult, VectorStore, DEFAULT_EMBEDDING_DIM};
use crate::error::WorkspaceError;
const TABLE_NAME: &str = "memory_chunks";
/// LanceDB-backed vector store.
pub struct LanceDbVectorStore {
db: Arc<lancedb::Connection>,
table_name: String,
embedding_dim: i32,
}
impl LanceDbVectorStore {
/// Create a new LanceDB store at the given path.
pub async fn new(path: impl AsRef<std::path::Path>) -> Result<Self, WorkspaceError> {
let path_str = path
.as_ref()
.to_str()
.ok_or_else(|| WorkspaceError::SearchFailed {
reason: "Invalid LanceDB path".to_string(),
})?;
let db = lancedb::connect(path_str)
.execute()
.await
.map_err(|e| WorkspaceError::SearchFailed {
reason: format!("Failed to connect to LanceDB: {}", e),
})?;
let store = Self {
db: Arc::new(db),
table_name: TABLE_NAME.to_string(),
embedding_dim: DEFAULT_EMBEDDING_DIM,
};
store.ensure_table().await?;
Ok(store)
}
async fn ensure_table(&self) -> Result<(), WorkspaceError> {
let tables = self.db.table_names().execute().await.map_err(|e| {
WorkspaceError::SearchFailed {
reason: format!("Failed to list tables: {}", e),
}
})?;
if tables.iter().any(|t| t == &self.table_name) {
return Ok(());
}
let schema = Arc::new(self.schema());
self.db
.create_empty_table(&self.table_name, schema.clone())
.execute()
.await
.map_err(|e| WorkspaceError::SearchFailed {
reason: format!("Failed to create table: {}", e),
})?;
let table = self.db.open_table(&self.table_name).execute().await.map_err(|e| {
WorkspaceError::SearchFailed {
reason: format!("Failed to open table: {}", e),
}
})?;
table
.create_index(&["vector"], Index::Auto)
.execute()
.await
.map_err(|e| WorkspaceError::SearchFailed {
reason: format!("Failed to create vector index: {}", e),
})?;
Ok(())
}
fn schema(&self) -> Schema {
Schema::new(vec![
Field::new("chunk_id", DataType::Utf8, false),
Field::new("document_id", DataType::Utf8, false),
Field::new("user_id", DataType::Utf8, false),
Field::new("agent_id", DataType::Utf8, true),
Field::new("content", DataType::Utf8, false),
Field::new(
"vector",
DataType::FixedSizeList(
Arc::new(Field::new("item", DataType::Float32, true)),
self.embedding_dim,
),
false,
),
])
}
}
#[async_trait]
impl VectorStore for LanceDbVectorStore {
async fn insert_chunk(
&self,
chunk_id: Uuid,
document_id: Uuid,
user_id: &str,
agent_id: Option<Uuid>,
content: &str,
embedding: &[f32],
) -> Result<(), WorkspaceError> {
if embedding.len() != self.embedding_dim as usize {
return Err(WorkspaceError::EmbeddingFailed {
reason: format!(
"Embedding dimension {} does not match expected {}",
embedding.len(),
self.embedding_dim
),
});
}
let table = self.db.open_table(&self.table_name).execute().await.map_err(|e| {
WorkspaceError::SearchFailed {
reason: format!("Failed to open table: {}", e),
}
})?;
let chunk_ids = StringArray::from(vec![chunk_id.to_string()]);
let document_ids = StringArray::from(vec![document_id.to_string()]);
let user_ids = StringArray::from(vec![user_id]);
let agent_ids = StringArray::from(vec![agent_id.map(|a| a.to_string())]);
let contents = StringArray::from(vec![content]);
let vec_values: Vec<Option<f32>> =
embedding.iter().map(|&x| Some(x)).collect();
let vectors = FixedSizeListArray::from_iter_primitive::<Float32Type, _, _>(
vec![Some(vec_values)].into_iter(),
self.embedding_dim,
);
let batch = RecordBatch::try_new(
Arc::new(self.schema()),
vec![
Arc::new(chunk_ids),
Arc::new(document_ids),
Arc::new(user_ids),
Arc::new(agent_ids),
Arc::new(contents),
Arc::new(vectors),
],
)
.map_err(|e| WorkspaceError::ChunkingFailed {
reason: format!("Failed to create record batch: {}", e),
})?;
let batches = RecordBatchIterator::new(
vec![Ok(batch)].into_iter(),
Arc::new(self.schema()),
);
table
.add(Box::new(batches))
.execute()
.await
.map_err(|e| WorkspaceError::ChunkingFailed {
reason: format!("Failed to insert chunk: {}", e),
})?;
Ok(())
}
async fn update_chunk_embedding(
&self,
chunk_id: Uuid,
document_id: Uuid,
user_id: &str,
agent_id: Option<Uuid>,
content: &str,
embedding: &[f32],
) -> Result<(), WorkspaceError> {
let table = self.db.open_table(&self.table_name).execute().await.map_err(|e| {
WorkspaceError::SearchFailed {
reason: format!("Failed to open table: {}", e),
}
})?;
table
.delete(&format!("chunk_id = '{}'", chunk_id))
.await
.map_err(|e| WorkspaceError::EmbeddingFailed {
reason: format!("Failed to delete chunk for update: {}", e),
})?;
self.insert_chunk(chunk_id, document_id, user_id, agent_id, content, embedding)
.await
}
async fn delete_chunks(&self, document_id: Uuid) -> Result<(), WorkspaceError> {
let table = self.db.open_table(&self.table_name).execute().await.map_err(|e| {
WorkspaceError::SearchFailed {
reason: format!("Failed to open table: {}", e),
}
})?;
table
.delete(&format!("document_id = '{}'", document_id))
.await
.map_err(|e| WorkspaceError::ChunkingFailed {
reason: format!("Failed to delete chunks: {}", e),
})?;
Ok(())
}
async fn vector_search(
&self,
user_id: &str,
agent_id: Option<Uuid>,
embedding: &[f32],
limit: usize,
) -> Result<Vec<RankedResult>, WorkspaceError> {
let table = self.db.open_table(&self.table_name).execute().await.map_err(|e| {
WorkspaceError::SearchFailed {
reason: format!("Failed to open table: {}", e),
}
})?;
let filter = if let Some(aid) = agent_id {
format!("user_id = '{}' AND agent_id = '{}'", user_id, aid)
} else {
format!("user_id = '{}' AND agent_id IS NULL", user_id)
};
let mut stream = table
.query()
.nearest_to(embedding)
.map_err(|e| WorkspaceError::SearchFailed {
reason: format!("Invalid query vector: {}", e),
})?
.only_if(&filter)
.limit(limit as u32)
.execute()
.await
.map_err(|e| WorkspaceError::SearchFailed {
reason: format!("Vector search failed: {}", e),
})?;
let mut results = Vec::new();
let mut rank: u32 = 1;
while let Some(batch) = stream.next().await {
let batch = batch.map_err(|e| WorkspaceError::SearchFailed {
reason: format!("Stream error: {}", e),
})?;
let chunk_id_col = batch
.column_by_name("chunk_id")
.ok_or_else(|| WorkspaceError::SearchFailed {
reason: "chunk_id column missing".to_string(),
})?;
let document_id_col = batch
.column_by_name("document_id")
.ok_or_else(|| WorkspaceError::SearchFailed {
reason: "document_id column missing".to_string(),
})?;
let content_col = batch
.column_by_name("content")
.ok_or_else(|| WorkspaceError::SearchFailed {
reason: "content column missing".to_string(),
})?;
let chunk_ids = chunk_id_col.as_any().downcast_ref::<StringArray>().ok_or_else(|| {
WorkspaceError::SearchFailed {
reason: "chunk_id wrong type".to_string(),
}
})?;
let document_ids = document_id_col.as_any().downcast_ref::<StringArray>().ok_or_else(|| {
WorkspaceError::SearchFailed {
reason: "document_id wrong type".to_string(),
}
})?;
let contents = content_col.as_any().downcast_ref::<StringArray>().ok_or_else(|| {
WorkspaceError::SearchFailed {
reason: "content wrong type".to_string(),
}
})?;
for i in 0..batch.num_rows() {
let chunk_id = chunk_ids
.value(i)
.parse()
.unwrap_or_else(|_| Uuid::nil());
let document_id = document_ids
.value(i)
.parse()
.unwrap_or_else(|_| Uuid::nil());
let content = contents.value(i).to_string();
results.push(RankedResult {
chunk_id,
document_id,
content,
rank,
});
rank += 1;
}
}
Ok(results)
}
}
}
#[cfg(feature = "lancedb")]
pub use impl_lancedb::LanceDbVectorStore;
#[cfg(all(test, feature = "lancedb"))]
mod tests {
use std::sync::Arc;
use tempfile::TempDir;
use uuid::Uuid;
use super::{LanceDbVectorStore, VectorStore, DEFAULT_EMBEDDING_DIM};
fn make_embedding(seed: f32) -> Vec<f32> {
(0..DEFAULT_EMBEDDING_DIM as usize)
.map(|i| (seed * (i as f32 + 1.0)).sin())
.collect()
}
#[tokio::test]
async fn test_insert_and_vector_search() {
let dir = TempDir::new().unwrap();
let store = LanceDbVectorStore::new(dir.path()).await.unwrap();
let chunk_id = Uuid::new_v4();
let document_id = Uuid::new_v4();
let user_id = "user1";
let content = "Rust is a systems programming language";
let embedding = make_embedding(1.0);
store
.insert_chunk(
chunk_id,
document_id,
user_id,
None,
content,
&embedding,
)
.await
.unwrap();
let results = store
.vector_search(user_id, None, &embedding, 5)
.await
.unwrap();
assert_eq!(results.len(), 1);
assert_eq!(results[0].chunk_id, chunk_id);
assert_eq!(results[0].document_id, document_id);
assert_eq!(results[0].content, content);
assert_eq!(results[0].rank, 1);
}
#[tokio::test]
async fn test_insert_multiple_and_search_returns_ordered() {
let dir = TempDir::new().unwrap();
let store = LanceDbVectorStore::new(dir.path()).await.unwrap();
let doc_id = Uuid::new_v4();
let user_id = "user1";
// Insert 3 chunks with different embeddings
for (i, seed) in [1.0, 2.0, 3.0].iter().enumerate() {
store
.insert_chunk(
Uuid::new_v4(),
doc_id,
user_id,
None,
&format!("content {}", i),
&make_embedding(*seed),
)
.await
.unwrap();
}
// Search returns all 3, ordered by similarity
let query_emb = make_embedding(2.0);
let results = store
.vector_search(user_id, None, &query_emb, 5)
.await
.unwrap();
assert_eq!(results.len(), 3);
let contents: Vec<_> = results.iter().map(|r| r.content.as_str()).collect();
assert!(contents.contains(&"content 0"));
assert!(contents.contains(&"content 1"));
assert!(contents.contains(&"content 2"));
}
#[tokio::test]
async fn test_delete_chunks() {
let dir = TempDir::new().unwrap();
let store = LanceDbVectorStore::new(dir.path()).await.unwrap();
let doc_id = Uuid::new_v4();
let user_id = "user1";
store
.insert_chunk(
Uuid::new_v4(),
doc_id,
user_id,
None,
"content",
&make_embedding(1.0),
)
.await
.unwrap();
let results = store
.vector_search(user_id, None, &make_embedding(1.0), 5)
.await
.unwrap();
assert_eq!(results.len(), 1);
store.delete_chunks(doc_id).await.unwrap();
let results_after = store
.vector_search(user_id, None, &make_embedding(1.0), 5)
.await
.unwrap();
assert!(results_after.is_empty());
}
#[tokio::test]
async fn test_update_chunk_embedding() {
let dir = TempDir::new().unwrap();
let store = LanceDbVectorStore::new(dir.path()).await.unwrap();
let chunk_id = Uuid::new_v4();
let doc_id = Uuid::new_v4();
let user_id = "user1";
let content = "original content";
store
.insert_chunk(
chunk_id,
doc_id,
user_id,
None,
content,
&make_embedding(1.0),
)
.await
.unwrap();
// Update with new embedding
let new_embedding = make_embedding(5.0);
store
.update_chunk_embedding(chunk_id, doc_id, user_id, None, content, &new_embedding)
.await
.unwrap();
// Search with new embedding should find it
let results = store
.vector_search(user_id, None, &new_embedding, 5)
.await
.unwrap();
assert_eq!(results.len(), 1);
assert_eq!(results[0].chunk_id, chunk_id);
}
#[tokio::test]
async fn test_vector_search_filters_by_user_and_agent() {
let dir = TempDir::new().unwrap();
let store = LanceDbVectorStore::new(dir.path()).await.unwrap();
let doc_id = Uuid::new_v4();
let embedding = make_embedding(1.0);
store
.insert_chunk(
Uuid::new_v4(),
doc_id,
"user1",
None,
"user1 content",
&embedding,
)
.await
.unwrap();
store
.insert_chunk(
Uuid::new_v4(),
doc_id,
"user2",
None,
"user2 content",
&embedding,
)
.await
.unwrap();
let results_user1 = store
.vector_search("user1", None, &embedding, 5)
.await
.unwrap();
assert_eq!(results_user1.len(), 1);
assert_eq!(results_user1[0].content, "user1 content");
let results_user2 = store
.vector_search("user2", None, &embedding, 5)
.await
.unwrap();
assert_eq!(results_user2.len(), 1);
assert_eq!(results_user2[0].content, "user2 content");
let results_wrong_user = store
.vector_search("user3", None, &embedding, 5)
.await
.unwrap();
assert!(results_wrong_user.is_empty());
}
#[tokio::test]
async fn test_insert_rejects_wrong_embedding_dim() {
let dir = TempDir::new().unwrap();
let store = LanceDbVectorStore::new(dir.path()).await.unwrap();
let wrong_dim: Vec<f32> = vec![1.0; 100];
let err = store
.insert_chunk(
Uuid::new_v4(),
Uuid::new_v4(),
"user1",
None,
"content",
&wrong_dim,
)
.await
.unwrap_err();
assert!(matches!(err, crate::error::WorkspaceError::EmbeddingFailed { .. }));
}
}
+4
View File
@@ -44,6 +44,8 @@ mod chunker;
mod document;
mod embeddings;
pub mod hygiene;
#[cfg(feature = "lancedb")]
pub mod lancedb_store;
#[cfg(feature = "postgres")]
mod repository;
mod search;
@@ -53,6 +55,8 @@ pub use document::{MemoryChunk, MemoryDocument, WorkspaceEntry, paths};
pub use embeddings::{EmbeddingProvider, MockEmbeddings, NearAiEmbeddings, OpenAiEmbeddings};
#[cfg(feature = "postgres")]
pub use repository::Repository;
#[cfg(feature = "lancedb")]
pub use lancedb_store::{LanceDbVectorStore, VectorStore, DEFAULT_EMBEDDING_DIM};
pub use search::{RankedResult, SearchConfig, SearchResult, reciprocal_rank_fusion};
use std::sync::Arc;
+30
View File
@@ -347,6 +347,36 @@ impl Repository {
Ok(())
}
/// Get a chunk by ID.
pub async fn get_chunk_by_id(
&self,
chunk_id: Uuid,
) -> Result<Option<MemoryChunk>, WorkspaceError> {
let conn = self.conn().await?;
let row = conn
.query_opt(
r#"
SELECT id, document_id, chunk_index, content, created_at
FROM memory_chunks WHERE id = $1
"#,
&[&chunk_id],
)
.await
.map_err(|e| WorkspaceError::SearchFailed {
reason: format!("Query failed: {}", e),
})?;
Ok(row.map(|r| MemoryChunk {
id: r.get("id"),
document_id: r.get("document_id"),
chunk_index: r.get("chunk_index"),
content: r.get("content"),
embedding: None,
created_at: r.get("created_at"),
}))
}
/// Get chunks without embeddings for backfilling.
pub async fn get_chunks_without_embeddings(
&self,
+159
View File
@@ -0,0 +1,159 @@
//! Integration tests for LanceDB vector store with Database wrapper.
//!
//! Requires: cargo test --features "libsql,lancedb"
//!
//! Verifies DbWithLanceVectorStore: document + chunk insert, hybrid search
//! (FTS from libSQL, vector from LanceDB), delete_chunks sync.
#![cfg(all(feature = "libsql", feature = "lancedb"))]
use std::sync::Arc;
use ironclaw::db::lancedb_wrapper::DbWithLanceVectorStore;
use ironclaw::db::Database;
use ironclaw::db::libsql_backend::LibSqlBackend;
use ironclaw::workspace::{LanceDbVectorStore, SearchConfig};
use tempfile::TempDir;
use uuid::Uuid;
const EMBEDDING_DIM: usize = 1536;
fn make_embedding(seed: f32) -> Vec<f32> {
(0..EMBEDDING_DIM)
.map(|i| (seed * (i as f32 + 1.0)).sin())
.collect()
}
async fn setup_wrapped_db() -> (Arc<dyn Database>, TempDir) {
let libsql = LibSqlBackend::new_memory().await.unwrap();
libsql.run_migrations().await.unwrap();
let lancedb_dir = TempDir::new().unwrap();
let store = LanceDbVectorStore::new(lancedb_dir.path())
.await
.unwrap();
let db = Arc::new(DbWithLanceVectorStore::new(
Arc::new(libsql) as Arc<dyn Database>,
Arc::new(store),
)) as Arc<dyn Database>;
(db, lancedb_dir)
}
#[tokio::test]
async fn test_wrapper_hybrid_search_combines_fts_and_vector() {
let (db, _) = setup_wrapped_db().await;
let user_id = "test_user";
let agent_id: Option<Uuid> = None;
// Create document
let doc = db
.get_or_create_document_by_path(user_id, agent_id, "context/rust.md")
.await
.unwrap();
// Write content for FTS
db.update_document(doc.id, "Rust is a systems programming language focused on safety and performance.")
.await
.unwrap();
// Chunk and insert with embedding (triggers sync to LanceDB)
let content = "Rust is a systems programming language focused on safety.";
let embedding = make_embedding(1.0);
let chunk_id = db
.insert_chunk(doc.id, 0, content, Some(&embedding))
.await
.unwrap();
// Hybrid search: FTS for "Rust" + vector for semantic
let config = SearchConfig::default().with_limit(5);
let results = db
.hybrid_search(
user_id,
agent_id,
"Rust",
Some(&embedding),
&config,
)
.await
.unwrap();
assert!(!results.is_empty(), "hybrid search should return results");
assert_eq!(results[0].chunk_id, chunk_id);
assert!(results[0].content.contains("Rust"));
}
#[tokio::test]
async fn test_wrapper_delete_chunks_removes_from_both() {
let (db, _) = setup_wrapped_db().await;
let user_id = "test_user";
let agent_id: Option<Uuid> = None;
let doc = db
.get_or_create_document_by_path(user_id, agent_id, "notes/deleted.md")
.await
.unwrap();
db.update_document(doc.id, "Content to be deleted.").await.unwrap();
let embedding = make_embedding(2.0);
db.insert_chunk(doc.id, 0, "Content to be deleted.", Some(&embedding))
.await
.unwrap();
let before = db
.hybrid_search(user_id, agent_id, "deleted", Some(&embedding), &SearchConfig::default())
.await
.unwrap();
assert_eq!(before.len(), 1);
db.delete_chunks(doc.id).await.unwrap();
let after = db
.hybrid_search(user_id, agent_id, "deleted", Some(&embedding), &SearchConfig::default())
.await
.unwrap();
assert!(after.is_empty());
}
#[tokio::test]
async fn test_wrapper_insert_chunk_syncs_to_lancedb() {
let (db, _) = setup_wrapped_db().await;
let user_id = "sync_user";
let agent_id: Option<Uuid> = None;
let doc = db
.get_or_create_document_by_path(user_id, agent_id, "sync/test.md")
.await
.unwrap();
let content = "Semantic content for vector search";
let embedding = make_embedding(3.0);
let chunk_id = db
.insert_chunk(doc.id, 0, content, Some(&embedding))
.await
.unwrap();
// Vector-only search (no FTS query match) - should still find via LanceDB
let config = SearchConfig::default().vector_only().with_limit(5);
let results = db
.hybrid_search(
user_id,
agent_id,
"nonexistent_fts_term",
Some(&embedding),
&config,
)
.await
.unwrap();
assert_eq!(results.len(), 1);
assert_eq!(results[0].chunk_id, chunk_id);
assert_eq!(results[0].content, content);
}