diff --git a/src/config/agent.rs b/src/config/agent.rs index d9d95ea..16081ff 100644 --- a/src/config/agent.rs +++ b/src/config/agent.rs @@ -171,7 +171,15 @@ impl Agent { let rag = app_state .rag_cache .load_with(key, || async move { - Rag::init(&app_clone, "rag", &rag_path_clone, &document_paths, abort).await + Rag::init( + &app_clone, + "rag", + &rag_path_clone, + &document_paths, + abort, + false, + ) + .await }) .await?; Some(rag) @@ -983,6 +991,10 @@ async fn init_graph_rags( extractor_model: rag_node.extractor_model.clone(), extractor_prompt: rag_node.extractor_prompt.clone(), graph_hops: rag_node.graph_hops, + // Graph-node RAGs are yaml-only: `RagNode` has no `driver` field, so + // there is nothing to forward. The rest-pattern also keeps this literal + // from breaking on future `RagInitConfig` additions. + ..Default::default() }; let fully_specified = config.embedding_model.is_some() && config.chunk_size.is_some() diff --git a/src/config/paths.rs b/src/config/paths.rs index 49f2135..268e497 100644 --- a/src/config/paths.rs +++ b/src/config/paths.rs @@ -16,7 +16,7 @@ use anyhow::{Context, Result, anyhow, bail}; use log::LevelFilter; use std::collections::HashSet; use std::env; -use std::fs::{read_dir, read_to_string}; +use std::fs::{read_dir, read_to_string, remove_file}; use std::path::{Path, PathBuf}; pub fn config_dir() -> PathBuf { @@ -414,6 +414,11 @@ pub fn list_rags() -> Vec { for entry in rd.flatten() { let name = entry.file_name(); if let Some(name) = name.to_string_lossy().strip_suffix(".yaml") { + // Sidecars are not RAGs. `.duckdb` files are already excluded by + // the `.yaml` suffix check above; this rejects `.sbx-mixin`. + if is_rag_sidecar_name(name) { + continue; + } names.push(name.to_string()); } } @@ -424,6 +429,43 @@ pub fn list_rags() -> Vec { } } +/// True for the sidecar YAML files that must never be listed or deleted as RAGs. +/// `name` is the already-stripped stem (i.e. after `strip_suffix(".yaml")`). +/// Uses `ends_with`, not `contains('.')`, so a RAG legitimately named "v2.docs" is +/// not rejected. +pub(crate) fn is_rag_sidecar_name(name: &str) -> bool { + name.ends_with(".sbx-mixin") +} + +/// Remove every sidecar belonging to RAG `name` in `dir`. Missing files are NOT an +/// error. A failure to remove an EXISTING mixin IS an error and must propagate — a +/// silently-orphaned mixin keeps a sandbox network permission alive after the user +/// believes it is gone. The `.duckdb` orphan is only wasted disk, so its removal +/// failure is ignorable; the asymmetry is deliberate. +/// +/// Callers must run this BEFORE unlinking the primary `.yaml`. If the YAML goes first +/// and this then fails, the RAG disappears from `list_rags()` — so the user can no +/// longer select it to retry — while its `allowedDomains` entry keeps being injected +/// into every sandbox launch. +pub(crate) fn remove_rag_sidecars(dir: &Path, name: &str) -> Result<()> { + let duckdb_path = dir.join(format!("{name}.duckdb")); + if duckdb_path.exists() { + let _ = remove_file(&duckdb_path); + } + let mixin_path = dir.join(format!("{name}.sbx-mixin.yaml")); + if mixin_path.exists() { + remove_file(&mixin_path).with_context(|| { + format!( + "Failed to remove the sandbox mixin for RAG '{name}' at '{}'. \ + The RAG was NOT deleted so you can retry; this host remains \ + whitelisted in the sandbox until the file is removed.", + mixin_path.display() + ) + })?; + } + Ok(()) +} + pub fn list_macros() -> Vec { list_file_names(macros_dir(), ".yaml") } @@ -846,4 +888,76 @@ mod tests { } let _ = fs::remove_dir_all(&root); } + + /// Unique temp dir for the sidecar helper tests. These take `dir: &Path` directly, + /// so no env-var mutation and therefore no `#[serial]` is needed. + fn sidecar_temp_dir(label: &str) -> PathBuf { + let unique = time::SystemTime::now() + .duration_since(time::UNIX_EPOCH) + .unwrap() + .as_nanos(); + let root = env::temp_dir().join(format!("coyote-{label}-test-{unique}")); + fs::create_dir_all(&root).unwrap(); + root + } + + #[test] + fn is_rag_sidecar_name_accepts_dotted_rag_names() { + // A RAG legitimately named "v2.docs" must not be mistaken for a sidecar. + assert!(!is_rag_sidecar_name("v2.docs")); + assert!(!is_rag_sidecar_name("myrag")); + assert!(is_rag_sidecar_name("myrag.sbx-mixin")); + assert!(is_rag_sidecar_name("v2.docs.sbx-mixin")); + } + + #[test] + fn remove_rag_sidecars_removes_both() { + let root = sidecar_temp_dir("rag-sidecars-both"); + let duckdb = root.join("docs.duckdb"); + let mixin = root.join("docs.sbx-mixin.yaml"); + fs::write(&duckdb, "db").unwrap(); + fs::write(&mixin, "mixin").unwrap(); + + remove_rag_sidecars(&root, "docs").unwrap(); + + assert!(!duckdb.exists(), "the .duckdb sidecar must be removed"); + assert!( + !mixin.exists(), + "the .sbx-mixin.yaml sidecar must be removed" + ); + let _ = fs::remove_dir_all(&root); + } + + #[test] + fn remove_rag_sidecars_is_ok_when_absent() { + let root = sidecar_temp_dir("rag-sidecars-absent"); + assert!(remove_rag_sidecars(&root, "docs").is_ok()); + let _ = fs::remove_dir_all(&root); + } + + #[test] + fn remove_rag_sidecars_runs_before_yaml_unlink() { + let root = sidecar_temp_dir("rag-sidecars-order"); + let yaml = root.join("docs.yaml"); + fs::write(&yaml, "rag").unwrap(); + // A non-empty DIRECTORY at the mixin path makes remove_file fail, standing in + // for any real removal failure (permissions, a busy mount). + let mixin = root.join("docs.sbx-mixin.yaml"); + fs::create_dir_all(&mixin).unwrap(); + fs::write(mixin.join("blocker"), "x").unwrap(); + + let err = remove_rag_sidecars(&root, "docs").unwrap_err(); + assert!( + err.to_string() + .contains("Failed to remove the sandbox mixin"), + "got: {err}" + ); + // The whole point of removing sidecars first: the RAG is still on disk, still + // listed, and the deletion is retryable. + assert!( + yaml.exists(), + "the .yaml must survive a sidecar-removal failure so the delete is retryable" + ); + let _ = fs::remove_dir_all(&root); + } } diff --git a/src/config/request_context.rs b/src/config/request_context.rs index f01850d..e2f3883 100644 --- a/src/config/request_context.rs +++ b/src/config/request_context.rs @@ -142,6 +142,12 @@ pub struct RequestContext { pub role: Option, pub session: Option, pub rag: Option>, + /// The cache key `self.rag` was actually inserted under, carried rather than + /// reconstructed. Reconstruction was the bug: the invalidation sites do not have + /// the information needed to rebuild the key (agent RAGs are inserted under the + /// AGENT's name but `rag.name()` is the constant "rag"), so insert and invalidate + /// silently disagreed. `None` for the temp RAG, which bypasses the cache entirely. + pub rag_key: Option, pub agent: Option, pub last_message: Option, @@ -176,6 +182,7 @@ impl RequestContext { role: None, session: None, rag: None, + rag_key: None, agent: None, last_message: None, tool_scope: ToolScope::default(), @@ -229,6 +236,7 @@ impl RequestContext { role: None, session: None, rag: None, + rag_key: None, agent: None, last_message: None, tool_scope: ToolScope { @@ -277,6 +285,7 @@ impl RequestContext { role: self.role.clone(), session: self.session.clone(), rag: self.rag.clone(), + rag_key: self.rag_key.clone(), agent: self.agent.clone(), last_message: self.last_message.clone(), tool_scope: self.tool_scope.clone(), @@ -315,6 +324,7 @@ impl RequestContext { role: None, session: None, rag: None, + rag_key: None, agent: None, last_message: None, tool_scope: ToolScope { @@ -2554,6 +2564,14 @@ impl RequestContext { match file_ext { Some(file_ext) => { if let Some(name) = name.to_string_lossy().strip_suffix(file_ext) { + // Sidecars are not independently deletable assets. + // Guarded on `kind == "rag"` because this scan is shared + // by all six kinds, and `session`/`macro` also use + // `.yaml`. The helper lives in paths.rs beside + // list_rags() so both filters cannot drift apart. + if kind == "rag" && paths::is_rag_sidecar_name(name) { + continue; + } names.push(name.to_string()); } } @@ -2590,6 +2608,13 @@ impl RequestContext { match file_ext { Some(ext) => { let path = dir.join(format!("{name}{ext}")); + // Sidecars FIRST. If this fails, the .yaml is still on disk, the + // RAG is still listed, and the user can retry. Unlinking the .yaml + // first would make the deletion unretryable while leaving an + // orphaned mixin whitelisting a host in every sandbox launch. + if kind == "rag" { + paths::remove_rag_sidecars(&dir, &name)?; + } remove_file(&path).with_context(|| { format!("Failed to delete {kind} at '{}'", path.display()) })?; @@ -3729,6 +3754,14 @@ impl RequestContext { .then(|| Arc::new(RwLock::new(Supervisor::new(max_concurrent, max_depth)))); self.rag = agent.rag(); + // Keep `rag_key` in lockstep with `rag`. Agent RAGs are cached under + // `RagKey::Agent()` (see `Agent::init`), so mirror that key exactly; + // leaving the previous key in place would let `.rebuild rag` invalidate an + // unrelated RAG's cache entry, and leaving it `None` would invalidate nothing. + self.rag_key = self + .rag + .is_some() + .then(|| RagKey::Agent(agent.name().to_string())); self.agent = Some(agent); self.supervisor = supervisor; self.inbox = None; @@ -3777,6 +3810,11 @@ impl RequestContext { self.pending_agents_guardrail_count = 0; self.todo_list = TodoList::default(); self.rag.take(); + // Cleared alongside `rag` so the pair never disagrees: an agent RAG is + // cached under `RagKey::Agent()`, and leaving that key behind + // would outlive the RAG it names. Latent rather than live today only + // because `rebuild_rag`/`edit_rag_docs` bail on `rag.is_none()` first. + self.rag_key = None; self.discontinuous_last_message(); } Ok(()) @@ -4087,7 +4125,9 @@ impl RequestContext { let rag_cache = self.rag_cache(); let working_mode = self.working_mode; - let rag: Arc = match rag { + // The key is returned alongside the Rag rather than assigned inside the match: + // `rag_cache` borrows `self`, so writing `self.rag_key` there is E0506. + let (rag, rag_key): (Arc, Option) = match rag { None => { let rag_path = self.rag_file(super::TEMP_RAG_NAME); if rag_path.exists() { @@ -4095,14 +4135,28 @@ impl RequestContext { format!("Failed to cleanup previous '{}' rag", super::TEMP_RAG_NAME) })?; } - Arc::new(Rag::init(&app, super::TEMP_RAG_NAME, &rag_path, &[], abort_signal).await?) + // The temp RAG is never inserted into the cache, so it has no key. + ( + Arc::new( + Rag::init( + &app, + super::TEMP_RAG_NAME, + &rag_path, + &[], + abort_signal, + false, + ) + .await?, + ), + None, + ) } Some(name) => { let rag_path = self.rag_file(name); let key = RagKey::Named(name.to_string()); - rag_cache - .load_with(key, || { + let loaded = rag_cache + .load_with(key.clone(), || { let app = app.clone(); let rag_path = rag_path.clone(); let abort_signal = abort_signal.clone(); @@ -4111,16 +4165,19 @@ impl RequestContext { if working_mode.is_cmd() { bail!("Unknown RAG '{name}'"); } - Rag::init(&app, name, &rag_path, &[], abort_signal.clone()).await + Rag::init(&app, name, &rag_path, &[], abort_signal.clone(), true) + .await } else { Rag::load(&app, name, &rag_path) } } }) - .await? + .await?; + (loaded, Some(key)) } }; self.rag = Some(rag); + self.rag_key = rag_key; Ok(()) } @@ -4161,12 +4218,9 @@ impl RequestContext { bail!("No changes") } - let key = if self.agent.is_some() { - RagKey::Agent(rag.name().to_string()) - } else { - RagKey::Named(rag.name().to_string()) - }; - self.rag_cache().invalidate(&key); + if let Some(key) = self.rag_key.clone() { + self.rag_cache().invalidate(&key); + } rag.refresh_document_paths( &new_document_paths, @@ -4194,12 +4248,9 @@ impl RequestContext { ); } - let key = if self.agent.is_some() { - RagKey::Agent(rag.name().to_string()) - } else { - RagKey::Named(rag.name().to_string()) - }; - self.rag_cache().invalidate(&key); + if let Some(key) = self.rag_key.clone() { + self.rag_cache().invalidate(&key); + } let document_paths = rag.document_paths().to_vec(); println!( @@ -4613,6 +4664,49 @@ mod tests { assert!(ctx.agent.is_none()); assert!(ctx.rag.is_none()); + assert_eq!(ctx.rag_key, None); + } + + #[test] + #[serial] + fn use_agent_does_not_carry_stale_rag_key() { + let _guard = TestConfigDirGuard::new(); + let mut ctx = create_test_ctx(); + let app = ctx.app.config.clone(); + let agent_name = format!( + "test_agent_{}", + SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap() + .as_nanos() + ); + let agent_dir = paths::agent_data_dir(&agent_name); + create_dir_all(&agent_dir).unwrap(); + write( + agent_dir.join("config.yaml"), + format!("name: {agent_name}\ninstructions: hi\n"), + ) + .unwrap(); + + // Stand in for the state `.rag docs` leaves behind: `use_rag` sets `rag` and + // `rag_key` together, so a named key is live when the agent is entered. + ctx.rag_key = Some(RagKey::Named("docs".to_string())); + + tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap() + .block_on(async { + ctx.use_agent(&app, &agent_name, None, utils::create_abort_signal()) + .await + .unwrap(); + }); + + // This agent has no RAG, so `rag` is None and `rag_key` must be None as well. + // Carrying `Named("docs")` across the transition would point `.rebuild rag` + // at an unrelated RAG's cache entry. + assert!(ctx.rag.is_none()); + assert_eq!(ctx.rag_key, None); } #[test] @@ -6000,6 +6094,31 @@ mod tests { assert!(paths::list_rags().is_empty()); } + /// A `.sbx-mixin.yaml` sidecar must not appear as a phantom RAG in TAB + /// completion or `.list rag`. A RAG whose name legitimately contains a dot must + /// still be listed — the filter uses `ends_with`, not `contains('.')`. + #[test] + #[serial] + fn list_rags_skips_sbx_mixin_sidecars() { + let _guard = TestConfigDirGuard::new(); + let rags_dir = paths::rags_dir(); + create_dir_all(&rags_dir).unwrap(); + write(rags_dir.join("docs.yaml"), "embedding_model: test").unwrap(); + write(rags_dir.join("docs.sbx-mixin.yaml"), "kind: mixin").unwrap(); + write(rags_dir.join("v2.docs.yaml"), "embedding_model: test").unwrap(); + + let names = paths::list_rags(); + assert!(names.contains(&"docs".to_string())); + assert!( + names.contains(&"v2.docs".to_string()), + "a dotted RAG name must still be listed: {names:?}" + ); + assert!( + !names.contains(&"docs.sbx-mixin".to_string()), + "the sandbox mixin sidecar must not appear as a RAG: {names:?}" + ); + } + #[test] #[serial] fn use_agent_errors_when_already_in_session() { diff --git a/src/rag/mod.rs b/src/rag/mod.rs index dd2ba91..6c8d41a 100644 --- a/src/rag/mod.rs +++ b/src/rag/mod.rs @@ -12,7 +12,10 @@ mod splitter; use self::graph::{KnowledgeGraph, extract_entities}; use self::provider::RagProvider; -use self::providers::YamlProvider; +// `providers::duckdb_path_from_yaml(path)` is called through the module path in +// `create()`, so `providers` itself must stay in scope — do not collapse it into +// the `use` below. +use self::providers::{DuckDbProvider, YamlProvider}; use anyhow::{Context, Result, anyhow, bail}; use bm25::{Language, SearchEngine, SearchEngineBuilder}; @@ -121,6 +124,9 @@ pub struct RagInitConfig { pub extractor_model: Option, pub extractor_prompt: Option, pub graph_hops: Option, + /// `None` -> "yaml". No serde attribute: this struct derives only + /// `Debug, Clone, Default` and is built in Rust, never deserialized. + pub driver: Option, } #[derive(Debug, Clone, Default)] @@ -147,7 +153,8 @@ impl Rag { bail!("Cannot build RAG knowledge base '{name}' with no documents"); } println!("⚙ Initializing RAG..."); - let data = Self::resolve_init_data(app, config)?; + let mut data = Self::resolve_init_data(app, config)?; + data.driver = config.driver.clone().unwrap_or_else(|| "yaml".to_string()); let mut rag = Self::create(app, name, save_path, data)?; let loaders = app.document_loaders.clone(); let (spinner, spinner_rx) = Spinner::create(""); @@ -257,12 +264,37 @@ impl Rag { save_path: &Path, doc_paths: &[String], abort_signal: AbortSignal, + prompt_for_driver: bool, ) -> Result { if !*IS_STDOUT_TERMINAL { bail!("Failed to init rag in non-interactive mode"); } println!("⚙ Initializing RAG..."); let (embedding_model, chunk_size, chunk_overlap) = Self::create_config(app)?; + // Only interactive named-RAG creation offers a driver choice. Temp RAGs and + // agent startup pass `false`; an explicit flag is used rather than inferring + // from the name because the agent path passes the literal name "rag", which is + // indistinguishable from a user creating a RAG genuinely named `rag`. + let driver = if prompt_for_driver { + let options = vec![ + "yaml — portable, in-memory HNSW; usable from several Coyote processes at once (default)", + "duckdb — persistent on-disk store; vectors and content survive restarts; HNSW approximate search. Can only be open in ONE Coyote process at a time, and its driver cannot be changed later without recreating the RAG", + ]; + let sel = Select::new("RAG storage driver:", options) + .with_starting_cursor(0) + .prompt()?; + if sel.starts_with("duckdb") { + println!( + "Note: a duckdb RAG can only be open in one Coyote process at a time, \ + and changing its driver later means deleting and recreating the RAG." + ); + "duckdb" + } else { + "yaml" + } + } else { + "yaml" + }; let reranker_model = app.rag_reranker_model.clone(); let top_k = app.rag_top_k; let extractor_model = match app.rag_extractor_model.clone() { @@ -275,7 +307,7 @@ impl Rag { app.rag_graph_hops }; let extractor_prompt = app.rag_extractor_prompt.clone(); - let data = RagData::new( + let mut data = RagData::new( embedding_model.id(), chunk_size, chunk_overlap, @@ -288,6 +320,7 @@ impl Rag { graph_hops: Some(graph_hops), }, ); + data.driver = driver.to_string(); let mut rag = Self::create(app, name, save_path, data)?; let mut paths = doc_paths.to_vec(); if paths.is_empty() { @@ -315,9 +348,50 @@ impl Rag { Self::create(app, name, path, data) } - pub fn create(app: &AppConfig, name: &str, path: &Path, data: RagData) -> Result { - let bm25 = data.build_bm25(); - let provider: Box = Box::new(YamlProvider::from_data(&data)); + /// `mut data` — the duckdb arm rehydrates `data.vectors` from the sidecar. + pub fn create(app: &AppConfig, name: &str, path: &Path, mut data: RagData) -> Result { + // Deliberately does NOT call rebuild_indexes: both callers construct the Rag + // before any documents are added, so rebuilding empty data would be a no-op. + // Actual population happens later via sync_documents. + let (provider, bm25): (Box, _) = match data.driver.as_str() { + "duckdb" => { + let db_path = providers::duckdb_path_from_yaml(path); + let dim = embedding_dim_for_model(&data.embedding_model); + let duck = DuckDbProvider::open(&db_path, dim)?; + // HYDRATE — mandatory, not an optimization. The YAML file for a duckdb + // RAG deliberately omits `vectors`, so `data.vectors` arrives empty from + // disk. Refilling it from the sidecar is what makes the NEXT incremental + // sync non-destructive: rebuild_indexes does CREATE OR REPLACE TABLE and + // writes exactly what data.vectors holds. Skip this and the first + // `.edit rag-docs` after a restart wipes every previously indexed vector. + // + // Guarded on is_empty() so a caller that already has vectors in memory + // is never overwritten by an empty table. + // + // 🔴 `?`, NOT `unwrap_or_default()`. A hydration failure must propagate. + // Degrading to an empty map here loads a RAG that looks healthy, answers + // every query with nothing, and then loses the store permanently on the + // first `.edit rag-docs`. The legitimate "nothing indexed yet" case is + // already Ok(empty) — open() runs CREATE TABLE IF NOT EXISTS — so `?` + // costs a new RAG nothing. + if data.vectors.is_empty() { + data.vectors = duck.read_all_vectors()?; + } + // data.files is always populated for duckdb, so build_bm25() is the only + // path; there is no from-DuckDB fallback. + let bm25 = data.build_bm25(); + (Box::new(duck), bm25) + } + "qdrant" => bail!( + "Qdrant RAGs cannot be constructed via Rag::create(); \ + use Rag::attach() or Rag::load_async() instead" + ), + _ => { + // "yaml" and any unknown driver — in-memory HNSW. + let bm25 = data.build_bm25(); + (Box::new(YamlProvider::from_data(&data)), bm25) + } + }; let node_to_docs = data.knowledge_graph.build_node_to_docs(); let embedding_model = Model::retrieve_model(app, &data.embedding_model, ModelType::Embedding)?; @@ -460,8 +534,18 @@ impl Rag { let path = Path::new(&self.path); ensure_parent_exists(path)?; - let content = serde_yaml::to_string(&self.data) - .with_context(|| format!("Failed to serde rag '{}'", self.name))?; + let content = if self.data.driver == "duckdb" { + // Embeddings live in the .duckdb sidecar; keep them out of the YAML file. + // Clone-and-empty rather than mutating self.data — the live map must stay + // complete for the next incremental sync, and save() takes &self, so any + // clear-then-restore would leave the object corrupted on an early return. + let mut on_disk = self.data.clone(); + on_disk.vectors.clear(); + serde_yaml::to_string(&on_disk) + } else { + serde_yaml::to_string(&self.data) + } + .with_context(|| format!("Failed to serde rag '{}'", self.name))?; fs::write(path, content).with_context(|| { format!("Failed to save rag '{}' to '{}'", self.name, path.display()) })?; @@ -841,13 +925,18 @@ impl Rag { } progress(&spinner, "Building store".into()); + // Derived in-memory state is refreshed BEFORE the fallible provider rebuild. + // `self.data` has already been mutated at this point, so returning early on a + // provider error while `bm25`/`node_to_docs` still describe the previous corpus + // would leave this Rag internally inconsistent. Both are pure functions of + // `self.data` and cannot fail, so doing them first is always safe. + self.bm25 = self.data.build_bm25(); + self.node_to_docs = self.data.knowledge_graph.build_node_to_docs(); // `refresh` is true for a full re-index (.rebuild rag / --rebuild-rag / // initial build) and false for an incremental .edit rag-docs change. // Passing it through is what stops a remote provider from wiping its // collection on a one-file add. self.provider.rebuild_indexes(&self.data, refresh).await?; - self.bm25 = self.data.build_bm25(); - self.node_to_docs = self.data.knowledge_graph.build_node_to_docs(); Ok(()) } @@ -1660,10 +1749,39 @@ fn reciprocal_rank_fusion( .collect() } +/// Map an embedding model id to its vector dimension. +/// +/// The DuckDB `FLOAT[N]` column type and its HNSW index are fixed at schema-creation +/// time, so this value must be decided before the first insert. An unrecognized model +/// falls back to 1536; if that is wrong, DuckDB raises a dimension-mismatch error on +/// the first insert rather than silently corrupting the schema, and the recovery is to +/// delete the sidecar and re-ingest from source. +fn embedding_dim_for_model(model_id: &str) -> usize { + match model_id { + m if m.contains("3-large") => 3072, + m if m.contains("3-small") || m.contains("ada-002") => 1536, + m if m.contains("nomic-embed-text") || m.contains("all-minilm") => 768, + m if m.contains("jina-embeddings-v2") => 1024, + _ => 1536, + } +} + #[cfg(test)] mod tests { use super::*; + #[test] + fn embedding_dim_for_model_maps_known_models() { + assert_eq!(embedding_dim_for_model("text-embedding-3-large"), 3072); + assert_eq!(embedding_dim_for_model("text-embedding-3-small"), 1536); + assert_eq!(embedding_dim_for_model("text-embedding-ada-002"), 1536); + assert_eq!(embedding_dim_for_model("nomic-embed-text"), 768); + assert_eq!(embedding_dim_for_model("all-minilm"), 768); + assert_eq!(embedding_dim_for_model("jina-embeddings-v2-base-en"), 1024); + // Unknown models fall back to the OpenAI-compatible default. + assert_eq!(embedding_dim_for_model("some-unknown-model"), 1536); + } + #[test] fn document_id_round_trip() { let id = DocumentId::new(5, 17); diff --git a/src/rag/providers/duckdb.rs b/src/rag/providers/duckdb.rs new file mode 100644 index 0000000..5c877fe --- /dev/null +++ b/src/rag/providers/duckdb.rs @@ -0,0 +1,964 @@ +use crate::rag::provider::RagProvider; +use crate::rag::{DocumentId, RagData}; + +use anyhow::{Context, Result, anyhow, bail}; +use async_trait::async_trait; +use duckdb::Connection; +use indexmap::IndexMap; +use log::warn; +use std::path::{Path, PathBuf}; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::{Arc, Mutex, MutexGuard}; + +/// Derive the DuckDB sidecar path from a RAG's YAML path: `docs.yaml` -> `docs.duckdb`. +pub(crate) fn duckdb_path_from_yaml(yaml_path: &Path) -> PathBuf { + yaml_path.with_extension("duckdb") +} + +pub struct DuckDbProvider { + path: PathBuf, + conn: Arc>, + /// Embedding dimension; fixed at open time because the `FLOAT[N]` column depends on it. + dim: usize, + /// True once an FTS index has been built on `documents`. Until then + /// `fts_main_documents.match_bm25` does not exist and any keyword query would + /// fail with a DuckDB catalog error. Backs `has_native_keyword_search`. + fts_ready: AtomicBool, +} + +impl DuckDbProvider { + /// Open (or create) the DuckDB file. `dim` is the embedding vector dimension, + /// supplied by the caller who knows the model. + pub fn open(db_path: &Path, dim: usize) -> Result { + let conn = Connection::open(db_path).with_context(|| { + format!( + "Failed to open the DuckDB store at '{}'. If another Coyote process (or \ + another window) has this RAG open, close it and retry — a duckdb RAG can \ + only be open in ONE process at a time. Unlike the yaml driver, its data \ + lives in a single file with an exclusive lock.", + db_path.display() + ) + })?; + // Statement order is load-bearing. `hnsw_enable_experimental_persistence` is + // registered BY the vss extension, so setting it before `LOAD vss` fails with + // "Setting with name ... is not in the catalog, but it exists in the vss + // extension". Omitting it entirely makes CREATE INDEX ... USING HNSW on a + // file-backed database fail with "HNSW index persistence is not yet supported + // by default". LOAD vss -> LOAD fts -> SET -> CREATE INDEX. + conn.execute_batch(&format!( + "LOAD vss; + LOAD fts; + SET hnsw_enable_experimental_persistence = true; + CREATE TABLE IF NOT EXISTS vectors ( + doc_id UBIGINT PRIMARY KEY, + embedding FLOAT[{dim}] + ); + CREATE INDEX IF NOT EXISTS hnsw_idx + ON vectors USING HNSW (embedding) + WITH (metric = 'cosine'); + CREATE TABLE IF NOT EXISTS documents ( + doc_id UBIGINT PRIMARY KEY, + page_content TEXT NOT NULL + );" + )) + .context("Failed to initialize DuckDB schema")?; + // A reopened file may already carry a live FTS index from a previous session, + // in which case keyword search works immediately. + let fts_exists = Self::probe_fts_index(&conn); + Ok(Self { + path: db_path.to_path_buf(), + conn: Arc::new(Mutex::new(conn)), + dim, + fts_ready: AtomicBool::new(fts_exists), + }) + } + + /// Read all `(doc_id, embedding)` pairs so `create()` can hydrate `data.vectors` + /// from disk. This is what makes the next incremental sync non-destructive, and is + /// mandatory rather than an optimization. + /// + /// 🔴 THIS PATTERN IS DUCKDB-ONLY. NEVER write the qdrant equivalent. Qdrant Cosine + /// collections L2-normalize stored vectors on write, so hydrating `data.vectors` + /// from a Qdrant read fills it with NORMALIZED vectors and the next save() destroys + /// the originals permanently. The symmetry with this function is exactly why that + /// bug is easy to introduce. + /// + /// ALL-OR-NOTHING. An empty `vectors` table is `Ok(empty)`; ANY decode failure or + /// non-finite embedding is an `Err`. A partially-hydrated map is worse than no map: + /// `rebuild_indexes` does `CREATE OR REPLACE TABLE` and writes exactly what it is + /// given, so a thinned map is committed as the new truth on the next sync. + pub(crate) fn read_all_vectors(&self) -> Result>> { + let conn = self.lock_conn()?; + let mut stmt = conn.prepare("SELECT doc_id, embedding FROM vectors")?; + // Collect into a Result, NOT a filter_map. `.filter_map(|r| r.ok())` here would + // turn a systematic decode failure (e.g. a schema written by a different duckdb + // version) into a silently short map, indistinguishable from an empty store. + let raw: Vec<(u64, Vec)> = stmt + .query_map([], |row| { + let id: u64 = row.get(0)?; + // The duckdb crate does not implement `FromSql` for `Vec`, so the + // column is read as a `Value` and destructured. A FLOAT[N] column yields + // `Value::Array`, a FLOAT[] column yields `Value::List`; match both so + // the reader survives a file written under either schema. + let embedding: Vec = match row.get::<_, duckdb::types::Value>(1)? { + duckdb::types::Value::Array(vals) | duckdb::types::Value::List(vals) => vals + .into_iter() + .map(|v| match v { + duckdb::types::Value::Float(f) => f, + duckdb::types::Value::Double(d) => d as f32, + _ => f32::NAN, + }) + .collect(), + _ => Vec::new(), + }; + Ok((id, embedding)) + })? + .collect::>>() + .with_context(|| { + format!( + "Failed to decode the `vectors` table in '{}'. The DuckDB sidecar is \ + unreadable; delete it and run `.rebuild rag` to re-ingest from source.", + self.path.display() + ) + })?; + + // Validate OUTSIDE the closure so the error can name the offending row. A + // non-finite embedding is NOT droppable: dropping it thins the map, and the + // thinned map is what the next CREATE OR REPLACE commits. + let mut out = IndexMap::with_capacity(raw.len()); + for (id, embedding) in raw { + if embedding.is_empty() || embedding.iter().any(|f| !f.is_finite()) { + bail!( + "Vector for doc_id {id} in '{}' is empty or contains a non-finite \ + value. Refusing to hydrate a partial vector map — that would erase \ + the remaining vectors on the next sync. Delete the sidecar and run \ + `.rebuild rag` to re-ingest from source.", + self.path.display() + ); + } + out.insert(DocumentId(id as usize), embedding); + } + Ok(out) + } + + /// Does a LIVE FTS index exist on `documents`? Probed at open time so a reopened + /// database reports native keyword search accurately. + /// + /// A schema-existence check alone is NOT sufficient. After + /// `CREATE OR REPLACE TABLE documents`, the `fts_main_documents` schema still + /// exists and `match_bm25` still SUCCEEDS — but returns zero rows for every term. + /// The index is silently dead. The probe therefore asserts that a KNOWN row comes + /// back rather than merely that the call did not error. + fn probe_fts_index(conn: &Connection) -> bool { + // 1. Structural check — cheap, and short-circuits a never-built index. + let schema_exists = conn + .query_row( + "SELECT COUNT(*) FROM duckdb_schemas() WHERE schema_name = 'fts_main_documents'", + [], + |row| row.get::<_, i64>(0), + ) + .map(|n| n > 0) + .unwrap_or(false); + if !schema_exists { + return false; + } + // 2. Liveness check — pull a real token out of a real row and confirm the index + // scores that same row. A live index returns >= 1; a stale one returns 0 + // without erroring. An empty `documents` table cannot be probed, and + // reporting false is correct: there is nothing to keyword-search, and the + // next rebuild_indexes sets the flag directly. + conn.query_row( + "SELECT COUNT(*) FROM documents d + WHERE fts_main_documents.match_bm25( + d.doc_id, + (SELECT string_split(trim(page_content), ' ')[1] + FROM documents WHERE length(trim(page_content)) > 0 LIMIT 1) + ) IS NOT NULL", + [], + |row| row.get::<_, i64>(0), + ) + .map(|n| n > 0) + .unwrap_or(false) + } + + /// Lock the shared connection, converting mutex poisoning into an `anyhow` error. + /// + /// Never `.lock().unwrap()` here: a panic anywhere inside a locked scope poisons the + /// mutex permanently, and an unwrap would then turn every subsequent RAG query into + /// a panic for the remaining life of the process. + fn lock_conn(&self) -> Result> { + self.conn + .lock() + .map_err(|e| anyhow!("DuckDB connection mutex was poisoned: {e}")) + } +} + +#[async_trait] +impl RagProvider for DuckDbProvider { + // NOTE: the MutexGuard is held across the body of these `async fn`s. That is sound + // ONLY because no `.await` appears inside a locked scope. Adding one would make the + // generated future non-`Send` and break the `RagProvider: Send + Sync` bound. + async fn vector_search( + &self, + embedding: &[f32], + top_k: usize, + min_score: f32, + ) -> Result> { + // Validate before building the SQL literal — a NaN would produce malformed SQL. + if embedding.iter().any(|f| !f.is_finite()) { + bail!("Query embedding contains a non-finite value (NaN or infinity)"); + } + let vals: String = embedding + .iter() + .map(|f| f.to_string()) + .collect::>() + .join(", "); + let dim = self.dim; + let conn = self.lock_conn()?; + // array_cosine_distance requires a FLOAT[N] ARRAY, not the LIST type FLOAT[]. + // ORDER BY distance ASC is required for the planner to use hnsw_idx; the + // similarity form (DESC) does NOT trigger the ANN index. Distance is converted + // back to a similarity score on return. + let sql = format!( + "SELECT doc_id, \ + array_cosine_distance(embedding, [{vals}]::FLOAT[{dim}]) AS distance \ + FROM vectors ORDER BY distance ASC LIMIT {top_k}" + ); + let mut stmt = conn.prepare(&sql)?; + let results = stmt + .query_map([], |row| { + let id: u64 = row.get(0)?; + // `array_cosine_distance` on a FLOAT[N] column returns FLOAT (f32), NOT + // DOUBLE. Reading it as f64 raises InvalidColumnType INSIDE the closure, + // which a bare `.filter_map(|r| r.ok())` would silently discard — + // yielding ZERO results with no error and no log line. + let distance: f32 = match row.get::<_, duckdb::types::Value>(1)? { + duckdb::types::Value::Float(f) => f, + duckdb::types::Value::Double(d) => d as f32, + other => { + warn!("unexpected distance type from DuckDB: {other:?}"); + return Err(duckdb::Error::InvalidQuery); + } + }; + Ok((DocumentId(id as usize), 1.0_f32 - distance)) + })? + // Log-and-drop rather than a bare `.ok()`: a systematic decode failure here + // is otherwise indistinguishable from "no matches". + .filter_map(|r| match r { + Ok(v) => Some(v), + Err(e) => { + warn!("vector_search row decode failed: {e}"); + None + } + }) + .filter(|(_, score)| *score > min_score) + .collect(); + Ok(results) + } + + async fn fetch_content(&self, ids: &[DocumentId]) -> Result> { + if ids.is_empty() { + return Ok(vec![]); + } + let conn = self.lock_conn()?; + let placeholders = ids.iter().map(|_| "?").collect::>().join(", "); + let sql = + format!("SELECT doc_id, page_content FROM documents WHERE doc_id IN ({placeholders})"); + let params: Vec = ids + .iter() + .map(|id| duckdb::types::Value::UBigInt(id.0 as u64)) + .collect(); + let mut stmt = conn.prepare(&sql)?; + let mut rows: Vec<(DocumentId, String)> = stmt + .query_map(duckdb::params_from_iter(params.iter()), |row| { + let id: u64 = row.get(0)?; + let text: String = row.get(1)?; + Ok((DocumentId(id as usize), text)) + })? + .filter_map(|r| match r { + Ok(v) => Some(v), + Err(e) => { + warn!("fetch_content row decode failed: {e}"); + None + } + }) + .collect(); + // Ordering contract: `WHERE doc_id IN (...)` returns rows in storage order, NOT + // in the order of `ids`. Since `ids` is the RRF-ranked list, returning storage + // order would silently discard the ranking. Re-sort by input position. + let position: std::collections::HashMap = + ids.iter().enumerate().map(|(i, id)| (*id, i)).collect(); + rows.sort_by_key(|(id, _)| position.get(id).copied().unwrap_or(usize::MAX)); + Ok(rows) + } + + async fn rebuild_indexes(&mut self, data: &RagData, _full_rebuild: bool) -> Result<()> { + // Local on-disk state — a wholesale table reset plus re-INSERT is always + // correct, so the incremental/full distinction is ignored. + // + // 🔴 ANTI-WIPE GUARD. The CREATE OR REPLACE TABLE below writes exactly what + // `data.vectors` holds. An empty map on a RAG that HAS indexed files, against a + // store that ALREADY holds vectors, is always a bug — a failed/skipped + // hydration, or a caller that emptied the live map. Refuse rather than commit + // the loss. Every legitimate empty-vector rebuild also has an empty `files` + // (a fresh RAG; the zero-document tests), so this cannot fire on a correct call. + // The `existing > 0` conjunct is load-bearing, not belt-and-braces: a populated + // fixture rebuilt against a fresh temp database has count 0 and must proceed. + if data.vectors.is_empty() && !data.files.is_empty() { + let existing: i64 = { + // Scoped: the guard MUST be dropped before `lock_conn()` is taken again + // below. `Mutex` is not reentrant — holding both self-deadlocks at + // runtime, with no compile error. + let conn = self.lock_conn()?; + conn.query_row("SELECT count(*) FROM vectors", [], |r| r.get(0)) + .context("Failed to count existing vectors before rebuild")? + }; + if existing > 0 { + bail!( + "Refusing to rebuild the DuckDB store at '{}': the in-memory vector \ + map is empty, but this RAG has {} indexed file(s) and the store \ + already holds {existing} vector(s). Rebuilding would erase them. \ + Vector hydration failed, or `data.vectors` was cleared on a live \ + Rag. Nothing was written; the store is intact.", + self.path.display(), + data.files.len() + ); + } + } + // Validate BEFORE opening the transaction so a bad embedding aborts before any + // write, not mid-write with the guard held. + for (doc_id, embedding) in &data.vectors { + if embedding.iter().any(|f| !f.is_finite()) { + bail!( + "Embedding for document {} contains a non-finite value", + doc_id.0 + ); + } + } + let dim = self.dim; + // `Connection::transaction()` takes `&mut self`, so this binding must be `mut`. + let mut conn = self.lock_conn()?; + let tx = conn + .transaction() + .context("Failed to begin DuckDB transaction")?; + + // Use `CREATE OR REPLACE TABLE`, not `DELETE FROM`. Deleting rows and + // re-inserting the SAME primary keys inside ONE transaction violates DuckDB's PK + // constraint: the index holds deleted keys until commit, so the second rebuild + // fails with a duplicate-key error. `CREATE OR REPLACE` drops the table and its + // indexes atomically, leaving no stale keys. It also drops the HNSW index, which + // is recreated after commit. + tx.execute_batch(&format!( + "CREATE OR REPLACE TABLE vectors ( + doc_id UBIGINT PRIMARY KEY, + embedding FLOAT[{dim}] + ); + CREATE OR REPLACE TABLE documents ( + doc_id UBIGINT PRIMARY KEY, + page_content TEXT NOT NULL + );" + )) + .context("Failed to reset DuckDB tables")?; + + { + // Plain INSERT: tables are empty after CREATE OR REPLACE, so no PK conflict. + // The embedding is bound as TEXT and cast in SQL — the duckdb crate's + // `bind_parameter` has no arm for List/Array, so binding a vector directly + // fails at runtime with "binding List parameters is not yet supported". + let mut vstmt = tx.prepare(&format!( + "INSERT INTO vectors (doc_id, embedding) VALUES (?, CAST(? AS FLOAT[{dim}]))" + ))?; + let mut dstmt = + tx.prepare("INSERT INTO documents (doc_id, page_content) VALUES (?, ?)")?; + + // 🔴 TWO INDEPENDENT LOOPS. `documents` is keyed on `files`, `vectors` on + // `vectors` — they are DIFFERENT key sets and neither is a subset of the + // other. Do not merge these into one loop over `&data.vectors` gated on a + // lookup into `files`: that produces keys(documents) ⊆ keys(vectors), and + // every id in `files \ vectors` (which `RagData::add`'s zip truncation + // really does produce) becomes a live, rankable id from BM25 and + // graph_search that resolves to nothing — no error, no log line. + for (id, doc) in data.iter_documents() { + dstmt.execute(duckdb::params![id.0 as u64, doc.page_content.as_str()])?; + } + + for (doc_id, embedding) in &data.vectors { + // Serialize as a DuckDB array literal: "[0.1,0.2,...]". Finiteness was + // validated above, so `to_string()` cannot emit NaN/inf here. + let embedding_text = { + let mut s = String::with_capacity(embedding.len() * 12 + 2); + s.push('['); + for (i, f) in embedding.iter().enumerate() { + if i > 0 { + s.push(','); + } + s.push_str(&f.to_string()); + } + s.push(']'); + s + }; + vstmt.execute(duckdb::params![doc_id.0 as u64, embedding_text.as_str()])?; + } + } // drop prepared statements before commit + + tx.commit().context("Failed to commit DuckDB transaction")?; + + // Recreate the HNSW index, dropped by CREATE OR REPLACE TABLE above. When + // CREATE INDEX USING HNSW fails, the preceding COMMIT still returns Ok, so the + // failure lands here; swallowing it would leave the database populated but + // UNINDEXED, with every vector_search silently falling back to a full scan. + conn.execute_batch( + "CREATE INDEX IF NOT EXISTS hnsw_idx \ + ON vectors USING HNSW (embedding) WITH (metric = 'cosine');", + ) + .context("Failed to recreate HNSW index")?; + + // Rebuild the FTS index (must be outside the transaction). This rebuild is + // MANDATORY on every pass, not an optimization: CREATE OR REPLACE TABLE above + // leaves the old fts_main_documents schema in place but DEAD — match_bm25 keeps + // succeeding while returning zero rows for every term. + // GUARDED on a non-empty table: FTS cannot index an empty table, and + // rebuild_indexes is legitimately called with zero documents (a fresh RAG, and + // three of the tests below). + let doc_count: i64 = conn + .query_row("SELECT count(*) FROM documents", [], |r| r.get(0)) + .context("Failed to count documents before FTS rebuild")?; + + if doc_count > 0 { + conn.execute_batch( + "DROP FUNCTION IF EXISTS fts_main_documents_match_bm25; + PRAGMA create_fts_index('documents', 'doc_id', 'page_content', overwrite=1);", + ) + .context("Failed to rebuild DuckDB FTS index")?; + } + + // Live only if an index was actually built. With zero documents there is no FTS + // index, so `has_native_keyword_search()` must stay false and hybrid_search must + // fall back to local BM25 — which is also empty, and therefore correct. + self.fts_ready.store(doc_count > 0, Ordering::Relaxed); + + Ok(()) + } + + async fn keyword_search(&self, query: &str, top_k: usize) -> Result> { + let conn = self.lock_conn()?; + // match_bm25 returns NULL for non-matching rows; WHERE filters them out. + let mut stmt = conn.prepare( + "SELECT doc_id, fts_main_documents.match_bm25(doc_id, ?) AS score + FROM documents + WHERE score IS NOT NULL + ORDER BY score DESC + LIMIT ?", + )?; + let results = stmt + .query_map(duckdb::params![query, top_k as u64], |row| { + let id: u64 = row.get(0)?; + // Same hazard as vector_search's distance column: guessing the width + // wrong raises InvalidColumnType INSIDE the closure, which a bare + // `.filter_map(|r| r.ok())` silently discards — yielding ZERO keyword + // hits, indistinguishable from "the query matched nothing". + let score: f32 = match row.get::<_, duckdb::types::Value>(1)? { + duckdb::types::Value::Double(d) => d as f32, + duckdb::types::Value::Float(f) => f, + other => { + warn!("unexpected match_bm25 score type from DuckDB: {other:?}"); + return Err(duckdb::Error::InvalidQuery); + } + }; + Ok((DocumentId(id as usize), score)) + })? + .filter_map(|r| match r { + Ok(v) => Some(v), + Err(e) => { + warn!("keyword_search row decode failed: {e}"); + None + } + }) + .collect(); + Ok(results) + } + + fn has_native_keyword_search(&self) -> bool { + // NOT an unconditional `true`. The FTS schema only exists after + // rebuild_indexes has run the pragma (or after reopening a database where a + // previous session did). `create()` deliberately does NOT call rebuild_indexes, + // so a freshly created RAG reaches this point with no FTS index at all. + // Returning true there would route every query into keyword_search, whose `?` + // would propagate a DuckDB catalog error and fail the WHOLE search. + self.fts_ready.load(Ordering::Relaxed) + } + + fn duplicate(&self, _data: &RagData) -> Box { + // Do NOT call DuckDbProvider::open() here. `Connection::open()` instantiates a + // NEW DuckDB *database* handle on the same file; the first handle still holds + // the file lock, so the second open fails with a locking error. Cloning the Arc + // shares the already-open connection, serialized by the Mutex. + // + // This means DuckDbProvider clones are NOT independent snapshots: state lives on + // disk and a rebuild through one handle is immediately visible to the other. + // That is unavoidable for any on-disk store and is handled by the discipline + // documented on `Rag`'s Clone impl — the pre-clone instance must be discarded. + Box::new(DuckDbProvider { + path: self.path.clone(), + conn: Arc::clone(&self.conn), + dim: self.dim, + fts_ready: AtomicBool::new(self.fts_ready.load(Ordering::Relaxed)), + }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + // Trait methods are only callable with the trait in scope. + use crate::rag::provider::RagProvider; + // `RagFile` / `RagDocument` have private fields; struct-literal construction + // compiles only because this module is a descendant of `crate::rag`. They are + // deliberately not in the non-test `use` block — unused there, and `--deny warnings` + // rejects that. + use crate::rag::{RagDocument, RagFile}; + + /// Unique temp path per test. `tempfile` is not a dev-dependency; this mirrors the + /// house pattern used elsewhere in the tree. + struct TempDb { + path: PathBuf, + } + + impl TempDb { + fn new(tag: &str) -> Self { + let unique = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_nanos(); + let path = std::env::temp_dir().join(format!("coyote-duckdb-{tag}-{unique}.duckdb")); + Self { path } + } + } + + impl Drop for TempDb { + fn drop(&mut self) { + let _ = std::fs::remove_file(&self.path); + // DuckDB writes a `.wal` sidecar; remove it too or /tmp accumulates + // one per test run. + let _ = std::fs::remove_file(self.path.with_extension("duckdb.wal")); + } + } + + /// ⚠️ NOTE THE EMPTY `files`. This fixture describes a RAG with ZERO documents. + /// Since `rebuild_indexes` populates the `documents` table from + /// `data.iter_documents()` (keyed on `files`), any test built on this helper alone + /// exercises the `vectors` table and NOTHING ELSE — no `documents` rows, no FTS + /// index. That is correct for the rebuild tests below, and is exactly why the + /// FTS-flag test and the `fetch_content` tests use `populated_rag_data()` instead. + /// Do not "simplify" them back onto this helper. + fn minimal_rag_data() -> RagData { + RagData { + embedding_model: "text-embedding-3-small".to_string(), + chunk_size: 1024, + chunk_overlap: 50, + top_k: 5, + driver: "duckdb".to_string(), + attached: false, + ..Default::default() + } + } + + /// Two files, one document each — the minimum fixture that produces `documents` + /// rows. + fn populated_rag_data() -> RagData { + let mut data = minimal_rag_data(); + debug_assert_eq!(DocumentId::new(0, 0), DocumentId(0)); + data.files.insert( + 0, + RagFile { + hash: "h0".to_string(), + path: "/tmp/a.md".to_string(), + documents: vec![RagDocument { + page_content: "alpha keyword".to_string(), + metadata: Default::default(), + }], + }, + ); + data.files.insert( + 1, + RagFile { + hash: "h1".to_string(), + path: "/tmp/b.md".to_string(), + documents: vec![RagDocument { + page_content: "beta keyword".to_string(), + metadata: Default::default(), + }], + }, + ); + data + } + + #[tokio::test] + async fn open_creates_schema() { + let db = TempDb::new("schema"); + let provider = DuckDbProvider::open(&db.path, 3).unwrap(); + let conn = provider.conn.lock().unwrap(); + let v: i64 = conn + .query_row("SELECT COUNT(*) FROM vectors", [], |r| r.get(0)) + .unwrap(); + let d: i64 = conn + .query_row("SELECT COUNT(*) FROM documents", [], |r| r.get(0)) + .unwrap(); + assert_eq!(v, 0); + assert_eq!(d, 0); + } + + #[tokio::test] + async fn vector_search_returns_top_result() { + let db = TempDb::new("vsearch"); + let provider = DuckDbProvider::open(&db.path, 3).unwrap(); + { + let conn = provider.conn.lock().unwrap(); + // The ::FLOAT[3] cast is REQUIRED: a bare [0.1, 0.2, 0.3] literal infers + // DOUBLE[], which does not match the FLOAT[N] ARRAY column type. + conn.execute( + "INSERT INTO vectors (doc_id, embedding) VALUES (0, [0.1, 0.2, 0.3]::FLOAT[3])", + [], + ) + .unwrap(); + } + let results = provider + .vector_search(&[0.1, 0.2, 0.3], 5, 0.0) + .await + .unwrap(); + assert_eq!(results.len(), 1); + assert_eq!(results[0].0.0, 0); + assert!(results[0].1 > 0.99); + } + + #[tokio::test] + async fn fetch_content_returns_stored_text() { + let db = TempDb::new("fetch"); + let provider = DuckDbProvider::open(&db.path, 3).unwrap(); + { + let conn = provider.conn.lock().unwrap(); + conn.execute( + "INSERT INTO documents (doc_id, page_content) VALUES (42, 'hello world')", + [], + ) + .unwrap(); + } + let results = provider.fetch_content(&[DocumentId(42)]).await.unwrap(); + assert_eq!(results.len(), 1); + assert_eq!(results[0].1, "hello world"); + } + + #[tokio::test] + async fn rebuild_indexes_then_vector_search() { + let db = TempDb::new("rebuild"); + let mut provider = DuckDbProvider::open(&db.path, 3).unwrap(); + let data = minimal_rag_data(); + provider.rebuild_indexes(&data, true).await.unwrap(); + let results = provider + .vector_search(&[0.1, 0.2, 0.3], 5, 0.0) + .await + .unwrap(); + assert!(results.is_empty()); // no vectors in empty data + } + + /// Rebuilding TWICE must succeed. + /// + /// The first rebuild always passes because the tables start empty. The bug this + /// guards against — `DELETE FROM` plus re-INSERT of the same PKs inside one + /// transaction — only fires on the SECOND rebuild, with "Duplicate key ... violates + /// primary key constraint". A single-rebuild test cannot catch it. + #[tokio::test] + async fn rebuild_indexes_is_idempotent_across_two_passes() { + let db = TempDb::new("idempotent"); + let mut provider = DuckDbProvider::open(&db.path, 3).unwrap(); + let mut data = minimal_rag_data(); + data.vectors.insert(DocumentId(0), vec![0.1, 0.2, 0.3]); + data.vectors.insert(DocumentId(1), vec![0.4, 0.5, 0.6]); + + provider.rebuild_indexes(&data, true).await.unwrap(); + // Second pass over the SAME doc_ids — this is the assertion that matters. + provider + .rebuild_indexes(&data, true) + .await + .expect("second rebuild must not violate the primary key constraint"); + + let conn = provider.conn.lock().unwrap(); + let count: i64 = conn + .query_row("SELECT COUNT(*) FROM vectors", [], |r| r.get(0)) + .unwrap(); + assert_eq!(count, 2, "rebuild must replace rows, not duplicate them"); + } + + /// Guards the incremental data-loss bug: `sync_documents` hash-skips unchanged + /// files, so on an incremental pass `data.vectors` holds ONLY the newly embedded + /// chunks. If the live map is ever emptied, the `CREATE OR REPLACE TABLE` in + /// rebuild_indexes writes only those, destroying every previously indexed vector. + /// Hydration via `read_all_vectors()` is what prevents it. + /// + /// This simulates the real restart-then-edit sequence: build, drop the in-memory + /// map, rehydrate from disk, add one vector, rebuild incrementally. + #[tokio::test] + async fn rebuild_indexes_preserves_prior_vectors_on_incremental_pass() { + let db = TempDb::new("incremental"); + let mut provider = DuckDbProvider::open(&db.path, 3).unwrap(); + + // Pass 1 — initial full build with two vectors. + let mut data = minimal_rag_data(); + data.vectors.insert(DocumentId(0), vec![0.1, 0.2, 0.3]); + data.vectors.insert(DocumentId(1), vec![0.4, 0.5, 0.6]); + provider.rebuild_indexes(&data, true).await.unwrap(); + + // Simulate a restart: the YAML on disk carries NO vectors, so a fresh RagData + // starts empty. create() hydrates it back from the sidecar. + let mut reloaded = minimal_rag_data(); + assert!( + reloaded.vectors.is_empty(), + "YAML-loaded data starts with no vectors" + ); + reloaded.vectors = provider.read_all_vectors().unwrap(); + assert_eq!( + reloaded.vectors.len(), + 2, + "hydration must restore both vectors" + ); + + // Pass 2 — incremental add of ONE new chunk. + reloaded.vectors.insert(DocumentId(2), vec![0.7, 0.8, 0.9]); + provider.rebuild_indexes(&reloaded, false).await.unwrap(); + + let conn = provider.conn.lock().unwrap(); + let count: i64 = conn + .query_row("SELECT COUNT(*) FROM vectors", [], |r| r.get(0)) + .unwrap(); + // Without hydration this is 1 — the two originals silently vanish. + assert_eq!( + count, 3, + "an incremental pass must not destroy previously indexed vectors" + ); + } + + /// `has_native_keyword_search()` must be false before any FTS index exists, + /// otherwise hybrid_search routes into keyword_search and the `?` turns a DuckDB + /// catalog error into a total query failure. + /// + /// 🔴 THE FIXTURE MUST CARRY DOCUMENTS. `rebuild_indexes` builds the FTS index only + /// when the `documents` table is non-empty, and that table is filled from + /// `data.iter_documents()` (keyed on `files`), never from `vectors`. With `files` + /// empty the pragma is skipped and `fts_ready` stores `false`. + #[tokio::test] + async fn keyword_search_is_not_advertised_before_first_rebuild() { + let db = TempDb::new("ftsflag"); + let mut provider = DuckDbProvider::open(&db.path, 3).unwrap(); + assert!( + !provider.has_native_keyword_search(), + "a freshly opened DB has no FTS index yet" + ); + + // 2 files × 1 document → 2 `documents` rows → doc_count > 0 → pragma runs. + let mut data = populated_rag_data(); + data.vectors.insert(DocumentId(0), vec![0.1, 0.2, 0.3]); + provider.rebuild_indexes(&data, true).await.unwrap(); + { + // Pin the precondition explicitly. If this ever reads 0 the assertion below + // is vacuous and the FTS hazard is unguarded again. + let conn = provider.conn.lock().unwrap(); + let docs: i64 = conn + .query_row("SELECT COUNT(*) FROM documents", [], |r| r.get(0)) + .unwrap(); + assert_eq!(docs, 2, "the fixture must produce documents rows"); + } + assert!( + provider.has_native_keyword_search(), + "rebuild_indexes creates the FTS index and must flip the flag" + ); + } + + /// The `documents` table is keyed on `files`, NOT on `vectors`. + /// + /// `fetch_content` is the sink for BOTH the BM25 and graph_search paths, and both + /// enumerate `files` via `iter_documents()`. If `rebuild_indexes` only inserts a + /// documents row for ids that also appear in `data.vectors`, every id in + /// `files \ vectors` resolves to nothing: healthy BM25 scores, healthy graph hits, + /// zero results, no error. + /// + /// The incremental test CANNOT catch this — its fixture has an empty `files`, so the + /// documents table is empty in every assertion either way. + #[tokio::test] + async fn fetch_content_resolves_documents_without_vectors() { + let db = TempDb::new("docsnovec"); + let mut provider = DuckDbProvider::open(&db.path, 3).unwrap(); + + // Two files, one document each — but a vector for ONLY THE FIRST. This is the + // `RagData::add` zip-truncation shape. + let mut data = populated_rag_data(); + let id_a = DocumentId::new(0, 0); + let id_b = DocumentId::new(1, 0); + data.vectors.insert(id_a, vec![0.1, 0.2, 0.3]); + assert!( + !data.vectors.contains_key(&id_b), + "precondition: b has no vector" + ); + + provider.rebuild_indexes(&data, true).await.unwrap(); + + let got = provider.fetch_content(&[id_a, id_b]).await.unwrap(); + assert_eq!( + got.len(), + 2, + "fetch_content must resolve BOTH documents; the one without a vector is \ + still reachable by keyword and graph search" + ); + assert_eq!(got[0].1, "alpha keyword"); + assert_eq!(got[1].1, "beta keyword"); + } + + /// `fetch_content` must return rows in INPUT order, not storage order. + /// + /// `YamlProvider` satisfies this contract structurally — it maps over `ids` — so the + /// existing yaml test proves nothing about DuckDB. DuckDB's `WHERE doc_id IN (...)` + /// returns storage order and relies on an explicit positional re-sort, which is what + /// can actually regress. `ids` is the RRF-ranked list, so losing the order silently + /// discards the ranking while still returning the right documents. + #[tokio::test] + async fn duckdb_fetch_content_preserves_input_order() { + let db = TempDb::new("order"); + let mut provider = DuckDbProvider::open(&db.path, 3).unwrap(); + let data = populated_rag_data(); + provider.rebuild_indexes(&data, true).await.unwrap(); + + let id_a = DocumentId::new(0, 0); // inserted first → storage order 0 + let id_b = DocumentId::new(1, 0); // inserted second → storage order 1 + + // REVERSED relative to storage order. Without the positional re-sort this + // returns ["alpha keyword", "beta keyword"] and the assertion fails. + let got = provider.fetch_content(&[id_b, id_a]).await.unwrap(); + assert_eq!(got.len(), 2); + assert_eq!(got[0].0, id_b, "input order must win over storage order"); + assert_eq!(got[0].1, "beta keyword"); + assert_eq!(got[1].0, id_a); + assert_eq!(got[1].1, "alpha keyword"); + } + + /// The anti-wipe guard. + /// + /// `rebuild_indexes` does `CREATE OR REPLACE TABLE` and writes exactly what + /// `data.vectors` holds. An empty map on a RAG that HAS indexed files, against a + /// store that already holds vectors, is always a bug (failed hydration, or a caller + /// that emptied the live map) and must be refused rather than committed. The loss + /// would otherwise be permanent: nothing re-embeds it back. + #[tokio::test] + async fn rebuild_indexes_refuses_to_wipe_when_vectors_are_empty_but_files_are_not() { + let db = TempDb::new("nowipe"); + let mut provider = DuckDbProvider::open(&db.path, 3).unwrap(); + + // A healthy store: two files with documents, two vectors. + let mut data = populated_rag_data(); + data.vectors + .insert(DocumentId::new(0, 0), vec![0.1, 0.2, 0.3]); + data.vectors + .insert(DocumentId::new(1, 0), vec![0.4, 0.5, 0.6]); + provider.rebuild_indexes(&data, true).await.unwrap(); + + // Now the failure state: vectors lost in memory, files intact. + let broken = populated_rag_data(); // same files, NO vectors + assert!( + broken.vectors.is_empty() && !broken.files.is_empty(), + "precondition" + ); + + let err = provider.rebuild_indexes(&broken, true).await.unwrap_err(); + assert!( + err.to_string().contains("Refusing to rebuild"), + "got: {err}" + ); + + // The store must be untouched — this is the whole point. + let conn = provider.conn.lock().unwrap(); + let count: i64 = conn + .query_row("SELECT COUNT(*) FROM vectors", [], |r| r.get(0)) + .unwrap(); + assert_eq!( + count, 2, + "the guard must abort BEFORE the CREATE OR REPLACE" + ); + } + + /// The guard must NOT fire on a legitimately empty RAG — otherwise a fresh + /// `Rag::init()` cannot complete. Empty `vectors` AND empty `files` is fine. + #[tokio::test] + async fn rebuild_indexes_allows_empty_vectors_when_files_are_also_empty() { + let db = TempDb::new("emptyok"); + let mut provider = DuckDbProvider::open(&db.path, 3).unwrap(); + let data = minimal_rag_data(); // no files, no vectors + provider + .rebuild_indexes(&data, true) + .await + .expect("a fresh RAG with nothing indexed must rebuild cleanly"); + } + + /// `duplicate()` must clone the Arc, NOT call `open()` again. + /// + /// This asserts SHARED state, which is the actual contract. It deliberately does NOT + /// use `fetch_content(&[])` — that early-returns before ever touching the + /// connection, so it would pass even against a broken double-opening implementation. + #[tokio::test] + async fn duplicate_shares_the_same_connection() { + let db = TempDb::new("dup"); + let provider = DuckDbProvider::open(&db.path, 3).unwrap(); + let dup = provider.duplicate(&minimal_rag_data()); + + // Write through the ORIGINAL... + { + let conn = provider.conn.lock().unwrap(); + conn.execute( + "INSERT INTO documents (doc_id, page_content) VALUES (7, 'shared row')", + [], + ) + .unwrap(); + } + // ...and read it back through the DUPLICATE. Only possible if they share state. + let via_dup = dup.fetch_content(&[DocumentId(7)]).await.unwrap(); + assert_eq!( + via_dup.len(), + 1, + "duplicate() must see writes made via the original" + ); + assert_eq!(via_dup[0].1, "shared row"); + } + + /// A decode failure must abort hydration rather than yield a thinned map: a short + /// map is what the next CREATE OR REPLACE commits as the new truth. + #[tokio::test] + async fn read_all_vectors_rejects_non_finite_embeddings() { + let db = TempDb::new("nonfinite"); + let provider = DuckDbProvider::open(&db.path, 3).unwrap(); + { + let conn = provider.conn.lock().unwrap(); + conn.execute( + "INSERT INTO vectors (doc_id, embedding) VALUES (0, [0.1, 0.2, 0.3]::FLOAT[3])", + [], + ) + .unwrap(); + conn.execute( + // Each element is cast individually: a bare ['nan', 0.2, 0.3] literal + // mixes text and numerics, so DuckDB infers DECIMAL and the conversion + // fails before the row is ever stored. + "INSERT INTO vectors (doc_id, embedding) VALUES \ + (1, [CAST('nan' AS FLOAT), CAST(0.2 AS FLOAT), CAST(0.3 AS FLOAT)]::FLOAT[3])", + [], + ) + .unwrap(); + } + let err = provider.read_all_vectors().unwrap_err(); + assert!( + err.to_string().contains("non-finite"), + "hydration must abort rather than silently drop the row; got: {err}" + ); + } + + #[test] + fn duckdb_path_from_yaml_swaps_extension() { + let p = duckdb_path_from_yaml(Path::new("/tmp/rags/docs.yaml")); + assert_eq!(p, PathBuf::from("/tmp/rags/docs.duckdb")); + } +} diff --git a/src/rag/providers/mod.rs b/src/rag/providers/mod.rs index 509f02c..19d9c9a 100644 --- a/src/rag/providers/mod.rs +++ b/src/rag/providers/mod.rs @@ -4,3 +4,10 @@ mod yaml; // is ambiguous (E0659) — `use` paths resolve against both this module's items // and the extern prelude, and `use` declarations may not shadow. pub use self::yaml::YamlProvider; + +mod duckdb; +pub use self::duckdb::DuckDbProvider; +// `create()` in rag/mod.rs derives the sidecar path through this. It is `pub(crate)` +// in providers/duckdb.rs, and the re-export must be `pub(crate)` too — a `pub use` of +// a `pub(crate)` item is E0364/E0365. +pub(crate) use self::duckdb::duckdb_path_from_yaml;