Compare commits

...
Author SHA1 Message Date
[email protected]andClaude Opus 4.6 c5b1cdc2f9 fix: address PR review comments (round 2)
Human reviewer (zmanian):
- H1: Accept configurable embedding dimension in LanceDbVectorStore::new()
  instead of hardcoding 1536. Dimension is sourced from
  EmbeddingProvider::dimension() at init time.
- H2: Skip double-write of embeddings to DB when external vector store
  is active (pass None to insert_chunk for embedding column).
- H3: Update PR title from "refactor" to "feat" (net-new feature).
- H4: Document non-atomic update_embedding in struct doc comment.

Bot reviewer (Copilot):
- Cache LanceDB table handle via tokio::sync::OnceCell (avoid
  open_table per operation).
- Cache Arc<Schema> in struct (avoid rebuilding per insert).
- Fix error variants: ChunkingFailed → EmbeddingFailed for LanceDB
  store/delete operations.
- Propagate store_embedding errors in reindex_document instead of
  warn-only (prevents silent data loss).
- Prefetch document metadata map in backfill_embeddings to avoid N+1
  queries.
- Add lancedb feature + protoc to CI test matrix so LanceDB tests
  actually run on Linux.

Co-Authored-By: Claude Opus 4.6 <[email protected]>
2026-03-08 17:34:11 -07:00
[email protected] fe17196367 Merge remote-tracking branch 'origin/main' into feat/lancedb-backend 2026-03-08 14:10:48 -07:00
[email protected]andClaude Opus 4.6 e84448d5ea ci: install protoc for lancedb feature builds
LanceDB depends on lance-encoding which requires protoc for protobuf
compilation. Add arduino/setup-protoc to all CI jobs that build with
--all-features.

Co-Authored-By: Claude Opus 4.6 <[email protected]>
2026-03-08 01:23:11 -08:00
[email protected]andClaude Opus 4.6 3f92b9cb24 Merge origin/main into feat/lancedb-backend
Resolve merge conflicts from main's config refactoring (config.rs split
into config/ directory), app builder pattern (src/app.rs), module renames
(libsql_backend → libsql), and new RankedResult.document_path field.

Co-Authored-By: Claude Opus 4.6 <[email protected]>
2026-03-08 01:19:25 -08:00
[email protected]andClaude Opus 4.6 1e0494e72d refactor: replace LanceDB Database decorator with VectorStore composition
Instead of wrapping all ~80 Database trait methods in a 664-line decorator
(lancedb_wrapper.rs), introduce a 4-method VectorStore trait that any vector
backend can implement. Workspace composes FTS from the database with vector
search from the external store via RRF fusion.

- Add src/workspace/vector_store.rs with VectorStore trait
- Rewrite lancedb_store.rs to implement VectorStore (not wrap Database)
- Delete src/db/lancedb_wrapper.rs (664 lines removed)
- Remove get_chunk_by_id from Database trait and all backends
- Workspace gains with_vector_store() builder for optional composition
- Fix LanceDB tests: bypass_vector_index() for brute-force search
- Fix integration tests: use temp file DB (libSQL :memory: is per-connection)
- Merge duplicate mod tests in config.rs

Net: -724 lines. Adding a new vector backend requires 4 methods, not 80.

Co-Authored-By: Claude Opus 4.6 <[email protected]>
2026-03-08 00:06:04 -08:00
c9dbcb2627 Update .env.example
Co-authored-by: Copilot <[email protected]>
2026-02-22 19:52:40 +04:00
Ilgın KanatandGitHub 98f44754e4 Merge branch 'main' into feat/lancedb-backend 2026-02-18 17:03:53 +04:00
e239ea850f Update src/workspace/lancedb_store.rs
Co-authored-by: Copilot <[email protected]>
2026-02-18 16:58:38 +04:00
ILGIN KANAT 6489c433f6 feat: add SQL-style escaping for LanceDB predicate values
- Introduced `escape_predicate_value` function to safely escape strings for use in LanceDB predicate expressions, preventing SQL injection.
- Updated delete operations in the LanceDB vector store to utilize the new escaping function for `chunk_id` and `document_id`.
- Enhanced filtering logic to apply escaping for `user_id` and `agent_id` in query conditions.

This change improves security by ensuring that user inputs are properly sanitized before being used in database queries.
2026-02-18 16:57:46 +04:00
ILGIN KANAT 327e009622 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.
2026-02-18 11:06:09 +04:00
13 changed files with 4201 additions and 57 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=builtin # default: use database's built-in index (pgvector or libsql_vector_idx); "pgvector" is also accepted as an alias for "builtin"
# VECTOR_BACKEND=lancedb
# LANCEDB_PATH=~/.ironclaw/lancedb # path for LanceDB when VECTOR_BACKEND=lancedb
# LLM Provider
# LLM_BACKEND=nearai # default
# Possible values: nearai, ollama, openai_compatible, openai, anthropic, tinfoil
+10
View File
@@ -36,6 +36,11 @@ jobs:
uses: dtolnay/rust-toolchain@stable
with:
components: clippy
- name: Install protoc (for lancedb)
if: matrix.name == 'all-features'
uses: arduino/setup-protoc@v3
with:
repo-token: ${{ secrets.GITHUB_TOKEN }}
- uses: Swatinem/rust-cache@v2
with:
key: clippy-${{ matrix.name }}
@@ -62,6 +67,11 @@ jobs:
uses: dtolnay/rust-toolchain@stable
with:
components: clippy
- name: Install protoc (for lancedb)
if: matrix.name == 'all-features'
uses: arduino/setup-protoc@v3
with:
repo-token: ${{ secrets.GITHUB_TOKEN }}
- uses: Swatinem/rust-cache@v2
with:
key: clippy-windows-${{ matrix.name }}
+15 -1
View File
@@ -14,7 +14,7 @@ jobs:
matrix:
include:
- name: all-features
flags: "--features postgres,libsql,html-to-markdown"
flags: "--features postgres,libsql,lancedb,html-to-markdown"
- name: default
flags: ""
- name: libsql-only
@@ -26,6 +26,11 @@ jobs:
uses: dtolnay/rust-toolchain@stable
with:
targets: wasm32-wasip2
- name: Install protoc (for lancedb)
if: matrix.name == 'all-features'
uses: arduino/setup-protoc@v3
with:
repo-token: ${{ secrets.GITHUB_TOKEN }}
- uses: Swatinem/rust-cache@v2
with:
key: ${{ matrix.name }}
@@ -66,6 +71,11 @@ jobs:
uses: actions/checkout@v6
- name: Install Rust
uses: dtolnay/rust-toolchain@stable
- name: Install protoc (for lancedb)
if: matrix.name == 'all-features'
uses: arduino/setup-protoc@v3
with:
repo-token: ${{ secrets.GITHUB_TOKEN }}
- uses: Swatinem/rust-cache@v2
with:
key: windows-${{ matrix.name }}
@@ -82,6 +92,10 @@ jobs:
uses: dtolnay/rust-toolchain@stable
with:
targets: wasm32-wasip2
- name: Install protoc (for lancedb)
uses: arduino/setup-protoc@v3
with:
repo-token: ${{ secrets.GITHUB_TOKEN }}
- uses: Swatinem/rust-cache@v2
with:
key: wasm-extensions
Generated
+3014 -36
View File
File diff suppressed because it is too large Load Diff
+6
View File
@@ -125,6 +125,11 @@ 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 }
arrow-array = { version = "57", optional = true }
arrow-schema = { version = "57", optional = true }
# WASM sandbox for untrusted tool execution
wasmtime = { version = "28", features = ["component-model"] }
wasmtime-wasi = "28" # WASI support for component model
@@ -189,6 +194,7 @@ insta = "1.46.3"
[features]
default = ["postgres", "libsql", "html-to-markdown"]
lancedb = ["dep:lancedb", "dep:arrow-array", "dep:arrow-schema"]
postgres = [
"dep:deadpool-postgres",
"dep:tokio-postgres",
+1 -1
View File
@@ -340,7 +340,7 @@ This document tracks feature parity between IronClaw (Rust implementation) and O
| Gemini embeddings | ✅ | ❌ | |
| Local embeddings | ✅ | ❌ | |
| SQLite-vec backend | ✅ | ❌ | IronClaw uses PostgreSQL |
| LanceDB backend | ✅ | | Configurable auto-capture max length |
| LanceDB backend | ✅ | | VectorStore trait + LanceDbVectorStore (configured via VECTOR_BACKEND=lancedb) |
| QMD backend | ✅ | ❌ | |
| Atomic reindexing | ✅ | ✅ | |
| Embeddings batching | ✅ | ✅ | `embed_batch` on EmbeddingProvider trait |
+41
View File
@@ -386,12 +386,53 @@ impl AppBuilder {
.embeddings
.create_provider(&self.config.llm.nearai.base_url, self.session.clone());
// Create optional external vector store for workspace semantic search
let vector_store: Option<Arc<dyn crate::workspace::VectorStore>> = {
#[cfg(feature = "lancedb")]
{
if self.config.database.vector_backend == crate::config::VectorBackend::LanceDb {
let path = self
.config
.database
.lancedb_path
.clone()
.unwrap_or_else(crate::config::default_lancedb_path);
let dim = embeddings.as_ref().map(|p| p.dimension());
match crate::workspace::LanceDbVectorStore::new(path, dim).await {
Ok(store) => {
tracing::info!("LanceDB vector store connected for workspace search");
Some(Arc::new(store) as Arc<dyn crate::workspace::VectorStore>)
}
Err(e) => {
tracing::warn!("Failed to initialize LanceDB: {}", e);
None
}
}
} else {
None
}
}
#[cfg(not(feature = "lancedb"))]
{
if self.config.database.vector_backend == crate::config::VectorBackend::LanceDb {
tracing::warn!(
"VECTOR_BACKEND=lancedb but 'lancedb' feature not enabled; \
falling back to built-in vector search"
);
}
None
}
};
// Register memory tools if database is available
let workspace = if let Some(ref db) = self.db {
let mut ws = Workspace::new_with_db("default", db.clone());
if let Some(ref emb) = embeddings {
ws = ws.with_embeddings(emb.clone());
}
if let Some(ref vs) = vector_store {
ws = ws.with_vector_store(vs.clone());
}
let ws = Arc::new(ws);
tools.register_memory_tools(Arc::clone(&ws));
Some(ws)
+92
View File
@@ -82,6 +82,31 @@ impl std::str::FromStr for SslMode {
}
}
/// Which vector store backend to use for workspace semantic search.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum VectorBackend {
/// Use the database's built-in vector support (pgvector or libsql_vector_idx).
#[default]
Builtin,
/// Use LanceDB as an external vector store (requires `lancedb` feature).
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 {
@@ -101,6 +126,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 fields --
/// Which vector store to use for workspace semantic search (default: Builtin).
pub vector_backend: VectorBackend,
/// Path to LanceDB directory (default: ~/.ironclaw/lancedb when vector_backend is LanceDb).
pub lancedb_path: Option<PathBuf>,
}
impl DatabaseConfig {
@@ -159,6 +190,25 @@ impl DatabaseConfig {
});
}
let vector_backend: VectorBackend = if let Some(s) = optional_env("VECTOR_BACKEND")? {
s.parse().map_err(|e| ConfigError::InvalidValue {
key: "VECTOR_BACKEND".to_string(),
message: e,
})?
} else {
VectorBackend::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),
@@ -167,6 +217,8 @@ impl DatabaseConfig {
libsql_path,
libsql_url,
libsql_auth_token,
vector_backend,
lancedb_path,
})
}
@@ -195,6 +247,11 @@ pub fn default_libsql_path() -> PathBuf {
ironclaw_base_dir().join("ironclaw.db")
}
/// Default LanceDB directory (~/.ironclaw/lancedb).
pub fn default_lancedb_path() -> PathBuf {
ironclaw_base_dir().join("lancedb")
}
#[cfg(test)]
mod tests {
use super::*;
@@ -224,4 +281,39 @@ mod tests {
fn ssl_mode_parse_invalid() {
assert!("invalid".parse::<SslMode>().is_err());
}
#[test]
fn 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!("invalid".parse::<VectorBackend>().is_err());
}
#[test]
fn default_lancedb_path_under_ironclaw() {
let path = super::default_lancedb_path();
assert!(path.to_string_lossy().contains("ironclaw"));
assert!(path.to_string_lossy().ends_with("lancedb"));
}
}
+6 -1
View File
@@ -33,7 +33,10 @@ use crate::settings::Settings;
pub use self::agent::AgentConfig;
pub use self::builder::BuilderModeConfig;
pub use self::channels::{ChannelsConfig, CliConfig, GatewayConfig, HttpConfig, SignalConfig};
pub use self::database::{DatabaseBackend, DatabaseConfig, SslMode, default_libsql_path};
pub use self::database::{
DatabaseBackend, DatabaseConfig, SslMode, VectorBackend, default_lancedb_path,
default_libsql_path,
};
pub use self::embeddings::EmbeddingsConfig;
pub use self::heartbeat::HeartbeatConfig;
pub use self::hygiene::HygieneConfig;
@@ -107,6 +110,8 @@ impl Config {
libsql_path: Some(libsql_path),
libsql_url: None,
libsql_auth_token: None,
vector_backend: VectorBackend::default(),
lancedb_path: None,
},
llm: LlmConfig::for_testing(),
embeddings: EmbeddingsConfig::default(),
+644
View File
@@ -0,0 +1,644 @@
//! 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
/// Default embedding dimension (text-embedding-3-small).
/// Override by passing the actual provider dimension to `LanceDbVectorStore::new()`.
pub const DEFAULT_EMBEDDING_DIM: i32 = 1536;
#[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::query::{ExecutableQuery, QueryBase};
use uuid::Uuid;
use super::DEFAULT_EMBEDDING_DIM;
use crate::error::WorkspaceError;
use crate::workspace::search::RankedResult;
use crate::workspace::vector_store::VectorStore;
const TABLE_NAME: &str = "memory_chunks";
/// Escapes a string for safe use in LanceDB predicate expressions.
/// Uses SQL-style escaping: single quotes are doubled to prevent injection.
fn escape_predicate_value(s: &str) -> String {
s.replace('\'', "''")
}
/// LanceDB-backed vector store.
///
/// The `update_embedding` method uses delete-then-insert (not atomic).
/// LanceDB does not support transactions, so a crash between the two
/// operations can lose the embedding for that chunk. This is acceptable
/// for personal workspace sizes where data can be reindexed.
pub struct LanceDbVectorStore {
db: Arc<lancedb::Connection>,
table_name: String,
embedding_dim: i32,
schema: Arc<Schema>,
table: tokio::sync::OnceCell<lancedb::Table>,
}
impl LanceDbVectorStore {
/// Create a new LanceDB store at the given path.
///
/// `embedding_dim` should match `EmbeddingProvider::dimension()`.
/// Pass `None` to use the default (1536, text-embedding-3-small).
pub async fn new(
path: impl AsRef<std::path::Path>,
embedding_dim: Option<usize>,
) -> 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 dim = embedding_dim.unwrap_or(DEFAULT_EMBEDDING_DIM as usize) as i32;
let schema = Arc::new(Self::build_schema(dim));
let store = Self {
db: Arc::new(db),
table_name: TABLE_NAME.to_string(),
embedding_dim: dim,
schema,
table: tokio::sync::OnceCell::new(),
};
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(());
}
self.db
.create_empty_table(&self.table_name, self.schema.clone())
.execute()
.await
.map_err(|e| WorkspaceError::SearchFailed {
reason: format!("Failed to create table: {}", e),
})?;
// Index creation is deferred — brute-force search via
// bypass_vector_index() works without a pre-built index and is
// sufficient for personal workspace sizes.
Ok(())
}
/// Get or open the cached table handle.
async fn table(&self) -> Result<&lancedb::Table, WorkspaceError> {
self.table
.get_or_try_init(|| async {
self.db
.open_table(&self.table_name)
.execute()
.await
.map_err(|e| WorkspaceError::SearchFailed {
reason: format!("Failed to open table: {}", e),
})
})
.await
}
fn build_schema(embedding_dim: i32) -> Schema {
Schema::new(vec![
Field::new("chunk_id", DataType::Utf8, false),
Field::new("document_id", DataType::Utf8, false),
Field::new("document_path", 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)),
embedding_dim,
),
false,
),
])
}
}
#[async_trait]
impl VectorStore for LanceDbVectorStore {
async fn store_embedding(
&self,
chunk_id: Uuid,
document_id: Uuid,
document_path: &str,
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.table().await?;
let chunk_ids = StringArray::from(vec![chunk_id.to_string()]);
let document_ids = StringArray::from(vec![document_id.to_string()]);
let document_paths = StringArray::from(vec![document_path]);
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)],
self.embedding_dim,
);
let batch = RecordBatch::try_new(
self.schema.clone(),
vec![
Arc::new(chunk_ids),
Arc::new(document_ids),
Arc::new(document_paths),
Arc::new(user_ids),
Arc::new(agent_ids),
Arc::new(contents),
Arc::new(vectors),
],
)
.map_err(|e| WorkspaceError::EmbeddingFailed {
reason: format!("Failed to create record batch: {}", e),
})?;
let batches =
RecordBatchIterator::new(vec![Ok(batch)].into_iter(), self.schema.clone());
table
.add(Box::new(batches) as Box<dyn arrow_array::RecordBatchReader + Send>)
.execute()
.await
.map_err(|e| WorkspaceError::EmbeddingFailed {
reason: format!("Failed to store embedding: {}", e),
})?;
Ok(())
}
async fn update_embedding(
&self,
chunk_id: Uuid,
document_id: Uuid,
document_path: &str,
user_id: &str,
agent_id: Option<Uuid>,
content: &str,
embedding: &[f32],
) -> Result<(), WorkspaceError> {
let table = self.table().await?;
table
.delete(&format!(
"chunk_id = '{}'",
escape_predicate_value(&chunk_id.to_string())
))
.await
.map_err(|e| WorkspaceError::EmbeddingFailed {
reason: format!("Failed to delete chunk for update: {}", e),
})?;
self.store_embedding(
chunk_id,
document_id,
document_path,
user_id,
agent_id,
content,
embedding,
)
.await
}
async fn delete_embeddings(&self, document_id: Uuid) -> Result<(), WorkspaceError> {
let table = self.table().await?;
table
.delete(&format!(
"document_id = '{}'",
escape_predicate_value(&document_id.to_string())
))
.await
.map_err(|e| WorkspaceError::EmbeddingFailed {
reason: format!("Failed to delete embeddings: {}", 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.table().await?;
let filter = if let Some(aid) = agent_id {
format!(
"user_id = '{}' AND agent_id = '{}'",
escape_predicate_value(user_id),
escape_predicate_value(&aid.to_string())
)
} else {
format!(
"user_id = '{}' AND agent_id IS NULL",
escape_predicate_value(user_id)
)
};
let query = table
.query()
.nearest_to(embedding)
.map_err(|e| WorkspaceError::SearchFailed {
reason: format!("Invalid query vector: {}", e),
})?
.only_if(&filter)
.bypass_vector_index()
.limit(limit);
let mut stream = ExecutableQuery::execute(&query).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_result) = stream.next().await {
let batch = batch_result.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 document_path_col = batch.column_by_name("document_path").ok_or_else(|| {
WorkspaceError::SearchFailed {
reason: "document_path 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 document_paths = document_path_col
.as_any()
.downcast_ref::<StringArray>()
.ok_or_else(|| WorkspaceError::SearchFailed {
reason: "document_path 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 raw_chunk_id = chunk_ids.value(i);
let chunk_id =
raw_chunk_id
.parse::<Uuid>()
.map_err(|e| WorkspaceError::SearchFailed {
reason: format!("Invalid chunk_id UUID '{}': {}", raw_chunk_id, e),
})?;
let raw_document_id = document_ids.value(i);
let document_id = raw_document_id.parse::<Uuid>().map_err(|e| {
WorkspaceError::SearchFailed {
reason: format!(
"Invalid document_id UUID '{}': {}",
raw_document_id, e
),
}
})?;
let document_path = document_paths.value(i).to_string();
let content = contents.value(i).to_string();
results.push(RankedResult {
chunk_id,
document_id,
document_path,
content,
rank,
});
rank += 1;
}
}
Ok(results)
}
}
}
#[cfg(feature = "lancedb")]
pub use impl_lancedb::LanceDbVectorStore;
#[cfg(all(test, feature = "lancedb"))]
mod tests {
use tempfile::TempDir;
use uuid::Uuid;
use super::{DEFAULT_EMBEDDING_DIM, LanceDbVectorStore};
use crate::workspace::vector_store::VectorStore;
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(), None).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
.store_embedding(
chunk_id,
document_id,
"test.md",
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(), None).await.unwrap();
let doc_id = Uuid::new_v4();
let user_id = "user1";
for (i, seed) in [1.0, 2.0, 3.0].iter().enumerate() {
store
.store_embedding(
Uuid::new_v4(),
doc_id,
"test.md",
user_id,
None,
&format!("content {}", i),
&make_embedding(*seed),
)
.await
.unwrap();
}
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(), None).await.unwrap();
let doc_id = Uuid::new_v4();
let user_id = "user1";
store
.store_embedding(
Uuid::new_v4(),
doc_id,
"test.md",
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_embeddings(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(), None).await.unwrap();
let chunk_id = Uuid::new_v4();
let doc_id = Uuid::new_v4();
let user_id = "user1";
let content = "original content";
store
.store_embedding(
chunk_id,
doc_id,
"test.md",
user_id,
None,
content,
&make_embedding(1.0),
)
.await
.unwrap();
let new_embedding = make_embedding(5.0);
store
.update_embedding(
chunk_id,
doc_id,
"test.md",
user_id,
None,
content,
&new_embedding,
)
.await
.unwrap();
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(), None).await.unwrap();
let doc_id = Uuid::new_v4();
let embedding = make_embedding(1.0);
store
.store_embedding(
Uuid::new_v4(),
doc_id,
"test.md",
"user1",
None,
"user1 content",
&embedding,
)
.await
.unwrap();
store
.store_embedding(
Uuid::new_v4(),
doc_id,
"test.md",
"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(), None).await.unwrap();
let wrong_dim: Vec<f32> = vec![1.0; 100];
let err = store
.store_embedding(
Uuid::new_v4(),
Uuid::new_v4(),
"test.md",
"user1",
None,
"content",
&wrong_dim,
)
.await
.unwrap_err();
assert!(matches!(
err,
crate::error::WorkspaceError::EmbeddingFailed { .. }
));
}
}
+172 -18
View File
@@ -42,20 +42,26 @@
mod chunker;
mod document;
mod embeddings;
pub mod embeddings;
pub mod hygiene;
#[cfg(feature = "lancedb")]
pub mod lancedb_store;
#[cfg(feature = "postgres")]
mod repository;
mod search;
pub mod vector_store;
pub use chunker::{ChunkConfig, chunk_document};
pub use document::{MemoryChunk, MemoryDocument, WorkspaceEntry, paths};
pub use embeddings::{
EmbeddingProvider, MockEmbeddings, NearAiEmbeddings, OllamaEmbeddings, OpenAiEmbeddings,
};
#[cfg(feature = "lancedb")]
pub use lancedb_store::{DEFAULT_EMBEDDING_DIM, LanceDbVectorStore};
#[cfg(feature = "postgres")]
pub use repository::Repository;
pub use search::{RankedResult, SearchConfig, SearchResult, reciprocal_rank_fusion};
pub use vector_store::VectorStore;
use std::sync::Arc;
@@ -332,6 +338,12 @@ pub struct Workspace {
storage: WorkspaceStorage,
/// Embedding provider for semantic search.
embeddings: Option<Arc<dyn EmbeddingProvider>>,
/// Optional external vector store for semantic search.
///
/// When set, embeddings are stored here instead of (or in addition to)
/// the database's built-in vector support, and hybrid search uses this
/// for the vector component while FTS comes from the database.
vector_store: Option<Arc<dyn VectorStore>>,
}
impl Workspace {
@@ -343,6 +355,7 @@ impl Workspace {
agent_id: None,
storage: WorkspaceStorage::Repo(Repository::new(pool)),
embeddings: None,
vector_store: None,
}
}
@@ -355,6 +368,7 @@ impl Workspace {
agent_id: None,
storage: WorkspaceStorage::Db(db),
embeddings: None,
vector_store: None,
}
}
@@ -370,6 +384,17 @@ impl Workspace {
self
}
/// Set an external vector store for semantic search.
///
/// When set, vector operations (store/search/delete embeddings) use this
/// store instead of the database's built-in vector support. FTS continues
/// to use the database. Hybrid search combines FTS from the database with
/// vector results from this store via RRF.
pub fn with_vector_store(mut self, store: Arc<dyn VectorStore>) -> Self {
self.vector_store = Some(store);
self
}
/// Get the user ID.
pub fn user_id(&self) -> &str {
&self.user_id
@@ -458,9 +483,21 @@ impl Workspace {
/// Delete a file.
///
/// Also deletes associated chunks.
/// Also deletes associated chunks (from both DB and external vector store).
pub async fn delete(&self, path: &str) -> Result<(), WorkspaceError> {
let path = normalize_path(path);
// Clean up external vector store before DB cascade deletes chunks
if let Some(ref vs) = self.vector_store
&& let Ok(doc) = self
.storage
.get_document_by_path(&self.user_id, self.agent_id, &path)
.await
&& let Err(e) = vs.delete_embeddings(doc.id).await
{
tracing::warn!("Failed to delete embeddings from vector store: {}", e);
}
self.storage
.delete_document_by_path(&self.user_id, self.agent_id, &path)
.await
@@ -725,20 +762,67 @@ impl Workspace {
query: &str,
config: SearchConfig,
) -> Result<Vec<SearchResult>, WorkspaceError> {
// Generate embedding for semantic search if provider available
let embedding = if let Some(ref provider) = self.embeddings {
Some(
provider
.embed(query)
.await
.map_err(|e| WorkspaceError::EmbeddingFailed {
reason: e.to_string(),
})?,
)
} else {
None
};
// Generate embedding for semantic search only when vector search is enabled
let embedding =
if config.use_vector {
if let Some(ref provider) = self.embeddings {
Some(provider.embed(query).await.map_err(|e| {
WorkspaceError::EmbeddingFailed {
reason: e.to_string(),
}
})?)
} else {
None
}
} else {
None
};
// When an external vector store is configured, do FTS from the
// database and vector search from the store, then fuse with RRF.
if let Some(ref vs) = self.vector_store {
// FTS from database (disable vector to avoid double-searching)
let fts_results = if config.use_fts {
let fts_config = SearchConfig {
use_fts: true,
use_vector: false,
..config.clone()
};
let fts_search = self
.storage
.hybrid_search(&self.user_id, self.agent_id, query, None, &fts_config)
.await?;
fts_search
.into_iter()
.enumerate()
.map(|(i, r)| RankedResult {
chunk_id: r.chunk_id,
document_id: r.document_id,
document_path: r.document_path,
content: r.content,
rank: (i + 1) as u32,
})
.collect()
} else {
Vec::new()
};
// Vector search from external store
let vector_results = if config.use_vector {
if let Some(ref emb) = embedding {
vs.vector_search(&self.user_id, self.agent_id, emb, config.pre_fusion_limit)
.await?
} else {
Vec::new()
}
} else {
Vec::new()
};
return Ok(reciprocal_rank_fusion(fts_results, vector_results, &config));
}
// No external vector store — use database's built-in hybrid search
self.storage
.hybrid_search(
&self.user_id,
@@ -760,7 +844,13 @@ impl Workspace {
// Chunk the content
let chunks = chunk_document(&doc.content, ChunkConfig::default());
// Delete old chunks
// Delete old embeddings from external vector store FIRST — if this fails,
// we abort before touching DB chunks, keeping the document consistent.
if let Some(ref vs) = self.vector_store {
vs.delete_embeddings(document_id).await?;
}
// Delete old chunks from database
self.storage.delete_chunks(document_id).await?;
// Insert new chunks
@@ -778,9 +868,33 @@ impl Workspace {
None
};
self.storage
.insert_chunk(document_id, index as i32, &content, embedding.as_deref())
// When an external vector store is active, skip writing embeddings
// to the DB (they'd never be queried from there).
let db_embedding = if self.vector_store.is_some() {
None
} else {
embedding.as_deref()
};
let chunk_id = self
.storage
.insert_chunk(document_id, index as i32, &content, db_embedding)
.await?;
// Sync embedding to external vector store (propagate errors to
// avoid leaving a document with deleted-then-missing embeddings).
if let (Some(vs), Some(emb)) = (&self.vector_store, &embedding) {
vs.store_embedding(
chunk_id,
document_id,
&doc.path,
&doc.user_id,
doc.agent_id,
&content,
emb,
)
.await?;
}
}
Ok(())
@@ -1025,6 +1139,23 @@ impl Workspace {
.get_chunks_without_embeddings(&self.user_id, self.agent_id, 100)
.await?;
// Prefetch document metadata to avoid N+1 queries when syncing to vector store
let doc_map: std::collections::HashMap<Uuid, crate::workspace::document::MemoryDocument> =
if self.vector_store.is_some() {
let mut map = std::collections::HashMap::new();
for chunk in &chunks {
if !map.contains_key(&chunk.document_id)
&& let Ok(doc) =
self.storage.get_document_by_id(chunk.document_id).await
{
map.insert(doc.id, doc);
}
}
map
} else {
std::collections::HashMap::new()
};
let mut count = 0;
for chunk in chunks {
match provider.embed(&chunk.content).await {
@@ -1032,6 +1163,29 @@ impl Workspace {
self.storage
.update_chunk_embedding(chunk.id, &embedding)
.await?;
// Sync to external vector store
if let Some(ref vs) = self.vector_store
&& let Some(doc) = doc_map.get(&chunk.document_id)
&& let Err(e) = vs
.update_embedding(
chunk.id,
chunk.document_id,
&doc.path,
&doc.user_id,
doc.agent_id,
&chunk.content,
&embedding,
)
.await
{
tracing::warn!(
"Failed to sync embedding to vector store for chunk {}: {}",
chunk.id,
e
);
}
count += 1;
}
Err(e) => {
+71
View File
@@ -0,0 +1,71 @@
//! Vector store abstraction for workspace semantic search.
//!
//! Separates vector search from the main `Database` trait so that
//! third-party vector backends (LanceDB, Qdrant, Pinecone, etc.) can
//! be added by implementing a 4-method trait instead of wrapping the
//! entire ~80-method `Database` trait.
//!
//! When no external vector store is configured, the built-in database
//! vector support (pgvector / libsql_vector_idx) is used via the
//! `Database::hybrid_search` method directly.
use async_trait::async_trait;
use uuid::Uuid;
use crate::error::WorkspaceError;
use crate::workspace::search::RankedResult;
/// External vector store for semantic search.
///
/// Implementations hold chunk embeddings and perform vector similarity
/// queries. Document/chunk metadata and FTS stay in the main database;
/// only embeddings live here.
///
/// # Adding a new backend
///
/// 1. Implement this trait for your backend (4 methods).
/// 2. Feature-gate the module (`#[cfg(feature = "mybackend")]`).
/// 3. Pass `Arc<dyn VectorStore>` to `Workspace::with_vector_store()`.
///
/// That's it — no Database wrapper, no delegation boilerplate.
#[async_trait]
#[allow(clippy::too_many_arguments)]
pub trait VectorStore: Send + Sync {
/// Store an embedding for a chunk.
async fn store_embedding(
&self,
chunk_id: Uuid,
document_id: Uuid,
document_path: &str,
user_id: &str,
agent_id: Option<Uuid>,
content: &str,
embedding: &[f32],
) -> Result<(), WorkspaceError>;
/// Update an existing chunk's embedding (delete + re-insert is fine).
async fn update_embedding(
&self,
chunk_id: Uuid,
document_id: Uuid,
document_path: &str,
user_id: &str,
agent_id: Option<Uuid>,
content: &str,
embedding: &[f32],
) -> Result<(), WorkspaceError>;
/// Delete all embeddings for a document.
async fn delete_embeddings(&self, document_id: Uuid) -> Result<(), WorkspaceError>;
/// Vector similarity search, filtered by user and optional agent.
///
/// Returns results ranked by similarity (rank 1 = most similar).
async fn vector_search(
&self,
user_id: &str,
agent_id: Option<Uuid>,
embedding: &[f32],
limit: usize,
) -> Result<Vec<RankedResult>, WorkspaceError>;
}
+123
View File
@@ -0,0 +1,123 @@
//! Integration tests for LanceDB vector store with Workspace composition.
//!
//! Requires: cargo test --features "libsql,lancedb"
//!
//! Verifies that Workspace correctly composes FTS from libSQL with vector
//! search from LanceDB via the VectorStore trait.
#![cfg(all(feature = "libsql", feature = "lancedb"))]
use std::sync::Arc;
use ironclaw::db::Database;
use ironclaw::db::libsql::LibSqlBackend;
use ironclaw::workspace::{LanceDbVectorStore, SearchConfig, Workspace};
use tempfile::TempDir;
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()
}
/// Mock embedding provider that returns deterministic embeddings.
struct FixedEmbeddings {
embedding: Vec<f32>,
}
#[async_trait::async_trait]
impl ironclaw::workspace::EmbeddingProvider for FixedEmbeddings {
fn dimension(&self) -> usize {
EMBEDDING_DIM
}
fn model_name(&self) -> &str {
"fixed-test"
}
fn max_input_length(&self) -> usize {
8192
}
async fn embed(
&self,
_text: &str,
) -> Result<Vec<f32>, ironclaw::workspace::embeddings::EmbeddingError> {
Ok(self.embedding.clone())
}
}
async fn setup_workspace() -> (Workspace, TempDir, TempDir) {
// Use a temp file (not :memory:) because libSQL in-memory DBs are connection-local
let db_dir = TempDir::new().unwrap();
let db_path = db_dir.path().join("test.db");
let libsql = LibSqlBackend::new_local(&db_path).await.unwrap();
libsql.run_migrations().await.unwrap();
let lancedb_dir = TempDir::new().unwrap();
let store = LanceDbVectorStore::new(lancedb_dir.path(), None).await.unwrap();
let embedding = make_embedding(1.0);
let ws = Workspace::new_with_db("test_user", Arc::new(libsql) as Arc<dyn Database>)
.with_vector_store(Arc::new(store))
.with_embeddings(Arc::new(FixedEmbeddings { embedding }));
(ws, lancedb_dir, db_dir)
}
#[tokio::test]
async fn test_workspace_hybrid_search_with_lancedb() {
let (ws, _keep_lance, _keep_db) = setup_workspace().await;
// Write a document — this triggers chunking + embedding + LanceDB sync
ws.write(
"context/rust.md",
"Rust is a systems programming language focused on safety.",
)
.await
.unwrap();
// Hybrid search: FTS for "Rust" + vector from LanceDB
let results = ws.search("Rust", 5).await.unwrap();
assert!(!results.is_empty(), "hybrid search should return results");
assert!(results[0].content.contains("Rust"));
}
#[tokio::test]
async fn test_workspace_delete_removes_from_lancedb() {
let (ws, _keep_lance, _keep_db) = setup_workspace().await;
ws.write("notes/deleted.md", "Content to be deleted.")
.await
.unwrap();
let before = ws.search("deleted", 5).await.unwrap();
assert_eq!(before.len(), 1);
ws.delete("notes/deleted.md").await.unwrap();
let after = ws.search("deleted", 5).await.unwrap();
assert!(after.is_empty());
}
#[tokio::test]
async fn test_workspace_vector_only_search_uses_lancedb() {
let (ws, _keep_lance, _keep_db) = setup_workspace().await;
ws.write("sync/test.md", "Semantic content for vector search")
.await
.unwrap();
// Vector-only search should find via LanceDB even with non-matching FTS query
let config = SearchConfig::default().vector_only().with_limit(5);
let results = ws
.search_with_config("nonexistent_fts_term", config)
.await
.unwrap();
assert_eq!(results.len(), 1);
assert!(results[0].content.contains("Semantic content"));
}