mirror of
https://github.com/outbackdingo/optimclaw.git
synced 2026-08-25 14:53:34 +00:00
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:
@@ -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.
|
||||
|
||||
@@ -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
File diff suppressed because it is too large
Load Diff
@@ -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
@@ -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 | ✅ | ❌ | |
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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))
|
||||
}
|
||||
}
|
||||
@@ -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
@@ -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.
|
||||
|
||||
@@ -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
@@ -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.
|
||||
|
||||
@@ -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 { .. }));
|
||||
}
|
||||
}
|
||||
@@ -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;
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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);
|
||||
}
|
||||
Reference in New Issue
Block a user