mirror of
https://github.com/outbackdingo/optimclaw.git
synced 2026-08-25 14:53:34 +00:00
Compare commits
10
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
c5b1cdc2f9 | ||
|
|
fe17196367 | ||
|
|
e84448d5ea | ||
|
|
3f92b9cb24 | ||
|
|
1e0494e72d | ||
|
|
c9dbcb2627 | ||
|
|
98f44754e4 | ||
|
|
e239ea850f | ||
|
|
6489c433f6 | ||
|
|
327e009622 |
@@ -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
|
||||
|
||||
@@ -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 }}
|
||||
|
||||
@@ -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
File diff suppressed because it is too large
Load Diff
@@ -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
@@ -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
@@ -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)
|
||||
|
||||
@@ -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
@@ -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(),
|
||||
|
||||
@@ -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
@@ -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) => {
|
||||
|
||||
@@ -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>;
|
||||
}
|
||||
@@ -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"));
|
||||
}
|
||||
Reference in New Issue
Block a user