feat(rag): add DuckDB provider behind the RAG driver abstraction

Phase 3 of the RAG driver abstraction. Adds a `DuckDbProvider` that keeps
vectors and document content in a `.duckdb` sidecar next to the existing
YAML metadata, selected by the `driver: duckdb` field.

- `src/rag/providers/duckdb.rs` (new): vector search via the vss extension
  and keyword search via fts, an all-or-nothing hydration path (a partial
  read is an error, never a shorter map), and an anti-wipe guard that
  refuses the destructive `CREATE OR REPLACE TABLE` when `data.vectors` is
  empty while `data.files` is not and the store still holds rows.
- `src/rag/mod.rs`: `sync_documents` now refreshes `bm25`/`node_to_docs`
  BEFORE the fallible `provider.rebuild_indexes`. `self.data` is already
  mutated by that point, so propagating a provider error afterwards would
  leave the derived in-memory state describing the previous corpus while
  `data` describes the new one. Both rebuilds are pure functions of
  `self.data` and cannot fail, so running them first is always safe.
- `src/config/paths.rs`: sidecar path helpers.
- `src/rag/providers/mod.rs`, `src/config/agent.rs`: driver dispatch and
  RAG cache keying.

Also keeps `RequestContext::rag_key` in lockstep with `rag` at the two
sites that were still missing it, so that a cache insert and its matching
invalidate are structurally incapable of disagreeing:

- `use_agent` assigned `self.rag` from the agent but never set `rag_key`.
  This one was live. Agent RAGs are inserted under `RagKey::Agent(<name>)`,
  so with `rag_key == None` the invalidation guards in `rebuild_rag` and
  `edit_rag_docs` matched nothing and `.rebuild rag` left the stale cache
  entry in place. Worse, a preceding `.rag <name>` left a stale
  `Named(<name>)` key attached to the agent's RAG, pointing the
  invalidation at an unrelated RAG's cache entry. Now mirrors the insert
  key exactly, yielding `None` when the agent has no RAG.
- `exit_agent` cleared `self.rag` but left `rag_key` behind. Latent rather
  than live, since `rebuild_rag`/`edit_rag_docs` both bail on
  `rag.is_none()` before reaching the invalidate guards, but the guards
  that make it unobservable are not the kind of thing to depend on.

Covered by `use_agent_does_not_carry_stale_rag_key`, and by a new
assertion in `exit_agent_clears_all_agent_state`.
This commit is contained in:
2026-08-10 12:51:56 -06:00
parent d734276927
commit 98d3ba4a83
6 changed files with 1364 additions and 30 deletions
+13 -1
View File
@@ -171,7 +171,15 @@ impl Agent {
let rag = app_state let rag = app_state
.rag_cache .rag_cache
.load_with(key, || async move { .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?; .await?;
Some(rag) Some(rag)
@@ -983,6 +991,10 @@ async fn init_graph_rags(
extractor_model: rag_node.extractor_model.clone(), extractor_model: rag_node.extractor_model.clone(),
extractor_prompt: rag_node.extractor_prompt.clone(), extractor_prompt: rag_node.extractor_prompt.clone(),
graph_hops: rag_node.graph_hops, 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() let fully_specified = config.embedding_model.is_some()
&& config.chunk_size.is_some() && config.chunk_size.is_some()
+115 -1
View File
@@ -16,7 +16,7 @@ use anyhow::{Context, Result, anyhow, bail};
use log::LevelFilter; use log::LevelFilter;
use std::collections::HashSet; use std::collections::HashSet;
use std::env; 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}; use std::path::{Path, PathBuf};
pub fn config_dir() -> PathBuf { pub fn config_dir() -> PathBuf {
@@ -414,6 +414,11 @@ pub fn list_rags() -> Vec<String> {
for entry in rd.flatten() { for entry in rd.flatten() {
let name = entry.file_name(); let name = entry.file_name();
if let Some(name) = name.to_string_lossy().strip_suffix(".yaml") { 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 `<name>.sbx-mixin`.
if is_rag_sidecar_name(name) {
continue;
}
names.push(name.to_string()); names.push(name.to_string());
} }
} }
@@ -424,6 +429,43 @@ pub fn list_rags() -> Vec<String> {
} }
} }
/// 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<String> { pub fn list_macros() -> Vec<String> {
list_file_names(macros_dir(), ".yaml") list_file_names(macros_dir(), ".yaml")
} }
@@ -846,4 +888,76 @@ mod tests {
} }
let _ = fs::remove_dir_all(&root); 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);
}
} }
+137 -18
View File
@@ -142,6 +142,12 @@ pub struct RequestContext {
pub role: Option<Role>, pub role: Option<Role>,
pub session: Option<Session>, pub session: Option<Session>,
pub rag: Option<Arc<Rag>>, pub rag: Option<Arc<Rag>>,
/// 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<RagKey>,
pub agent: Option<Agent>, pub agent: Option<Agent>,
pub last_message: Option<LastMessage>, pub last_message: Option<LastMessage>,
@@ -176,6 +182,7 @@ impl RequestContext {
role: None, role: None,
session: None, session: None,
rag: None, rag: None,
rag_key: None,
agent: None, agent: None,
last_message: None, last_message: None,
tool_scope: ToolScope::default(), tool_scope: ToolScope::default(),
@@ -229,6 +236,7 @@ impl RequestContext {
role: None, role: None,
session: None, session: None,
rag: None, rag: None,
rag_key: None,
agent: None, agent: None,
last_message: None, last_message: None,
tool_scope: ToolScope { tool_scope: ToolScope {
@@ -277,6 +285,7 @@ impl RequestContext {
role: self.role.clone(), role: self.role.clone(),
session: self.session.clone(), session: self.session.clone(),
rag: self.rag.clone(), rag: self.rag.clone(),
rag_key: self.rag_key.clone(),
agent: self.agent.clone(), agent: self.agent.clone(),
last_message: self.last_message.clone(), last_message: self.last_message.clone(),
tool_scope: self.tool_scope.clone(), tool_scope: self.tool_scope.clone(),
@@ -315,6 +324,7 @@ impl RequestContext {
role: None, role: None,
session: None, session: None,
rag: None, rag: None,
rag_key: None,
agent: None, agent: None,
last_message: None, last_message: None,
tool_scope: ToolScope { tool_scope: ToolScope {
@@ -2554,6 +2564,14 @@ impl RequestContext {
match file_ext { match file_ext {
Some(file_ext) => { Some(file_ext) => {
if let Some(name) = name.to_string_lossy().strip_suffix(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()); names.push(name.to_string());
} }
} }
@@ -2590,6 +2608,13 @@ impl RequestContext {
match file_ext { match file_ext {
Some(ext) => { Some(ext) => {
let path = dir.join(format!("{name}{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(|| { remove_file(&path).with_context(|| {
format!("Failed to delete {kind} at '{}'", path.display()) 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)))); .then(|| Arc::new(RwLock::new(Supervisor::new(max_concurrent, max_depth))));
self.rag = agent.rag(); self.rag = agent.rag();
// Keep `rag_key` in lockstep with `rag`. Agent RAGs are cached under
// `RagKey::Agent(<agent name>)` (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.agent = Some(agent);
self.supervisor = supervisor; self.supervisor = supervisor;
self.inbox = None; self.inbox = None;
@@ -3777,6 +3810,11 @@ impl RequestContext {
self.pending_agents_guardrail_count = 0; self.pending_agents_guardrail_count = 0;
self.todo_list = TodoList::default(); self.todo_list = TodoList::default();
self.rag.take(); self.rag.take();
// Cleared alongside `rag` so the pair never disagrees: an agent RAG is
// cached under `RagKey::Agent(<agent name>)`, 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(); self.discontinuous_last_message();
} }
Ok(()) Ok(())
@@ -4087,7 +4125,9 @@ impl RequestContext {
let rag_cache = self.rag_cache(); let rag_cache = self.rag_cache();
let working_mode = self.working_mode; let working_mode = self.working_mode;
let rag: Arc<Rag> = 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<Rag>, Option<RagKey>) = match rag {
None => { None => {
let rag_path = self.rag_file(super::TEMP_RAG_NAME); let rag_path = self.rag_file(super::TEMP_RAG_NAME);
if rag_path.exists() { if rag_path.exists() {
@@ -4095,14 +4135,28 @@ impl RequestContext {
format!("Failed to cleanup previous '{}' rag", super::TEMP_RAG_NAME) 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) => { Some(name) => {
let rag_path = self.rag_file(name); let rag_path = self.rag_file(name);
let key = RagKey::Named(name.to_string()); let key = RagKey::Named(name.to_string());
rag_cache let loaded = rag_cache
.load_with(key, || { .load_with(key.clone(), || {
let app = app.clone(); let app = app.clone();
let rag_path = rag_path.clone(); let rag_path = rag_path.clone();
let abort_signal = abort_signal.clone(); let abort_signal = abort_signal.clone();
@@ -4111,16 +4165,19 @@ impl RequestContext {
if working_mode.is_cmd() { if working_mode.is_cmd() {
bail!("Unknown RAG '{name}'"); 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 { } else {
Rag::load(&app, name, &rag_path) Rag::load(&app, name, &rag_path)
} }
} }
}) })
.await? .await?;
(loaded, Some(key))
} }
}; };
self.rag = Some(rag); self.rag = Some(rag);
self.rag_key = rag_key;
Ok(()) Ok(())
} }
@@ -4161,12 +4218,9 @@ impl RequestContext {
bail!("No changes") bail!("No changes")
} }
let key = if self.agent.is_some() { if let Some(key) = self.rag_key.clone() {
RagKey::Agent(rag.name().to_string()) self.rag_cache().invalidate(&key);
} else { }
RagKey::Named(rag.name().to_string())
};
self.rag_cache().invalidate(&key);
rag.refresh_document_paths( rag.refresh_document_paths(
&new_document_paths, &new_document_paths,
@@ -4194,12 +4248,9 @@ impl RequestContext {
); );
} }
let key = if self.agent.is_some() { if let Some(key) = self.rag_key.clone() {
RagKey::Agent(rag.name().to_string()) self.rag_cache().invalidate(&key);
} else { }
RagKey::Named(rag.name().to_string())
};
self.rag_cache().invalidate(&key);
let document_paths = rag.document_paths().to_vec(); let document_paths = rag.document_paths().to_vec();
println!( println!(
@@ -4613,6 +4664,49 @@ mod tests {
assert!(ctx.agent.is_none()); assert!(ctx.agent.is_none());
assert!(ctx.rag.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] #[test]
@@ -6000,6 +6094,31 @@ mod tests {
assert!(paths::list_rags().is_empty()); assert!(paths::list_rags().is_empty());
} }
/// A `<name>.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] #[test]
#[serial] #[serial]
fn use_agent_errors_when_already_in_session() { fn use_agent_errors_when_already_in_session() {
+128 -10
View File
@@ -12,7 +12,10 @@ mod splitter;
use self::graph::{KnowledgeGraph, extract_entities}; use self::graph::{KnowledgeGraph, extract_entities};
use self::provider::RagProvider; 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 anyhow::{Context, Result, anyhow, bail};
use bm25::{Language, SearchEngine, SearchEngineBuilder}; use bm25::{Language, SearchEngine, SearchEngineBuilder};
@@ -121,6 +124,9 @@ pub struct RagInitConfig {
pub extractor_model: Option<String>, pub extractor_model: Option<String>,
pub extractor_prompt: Option<String>, pub extractor_prompt: Option<String>,
pub graph_hops: Option<usize>, pub graph_hops: Option<usize>,
/// `None` -> "yaml". No serde attribute: this struct derives only
/// `Debug, Clone, Default` and is built in Rust, never deserialized.
pub driver: Option<String>,
} }
#[derive(Debug, Clone, Default)] #[derive(Debug, Clone, Default)]
@@ -147,7 +153,8 @@ impl Rag {
bail!("Cannot build RAG knowledge base '{name}' with no documents"); bail!("Cannot build RAG knowledge base '{name}' with no documents");
} }
println!("⚙ Initializing RAG..."); 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 mut rag = Self::create(app, name, save_path, data)?;
let loaders = app.document_loaders.clone(); let loaders = app.document_loaders.clone();
let (spinner, spinner_rx) = Spinner::create(""); let (spinner, spinner_rx) = Spinner::create("");
@@ -257,12 +264,37 @@ impl Rag {
save_path: &Path, save_path: &Path,
doc_paths: &[String], doc_paths: &[String],
abort_signal: AbortSignal, abort_signal: AbortSignal,
prompt_for_driver: bool,
) -> Result<Self> { ) -> Result<Self> {
if !*IS_STDOUT_TERMINAL { if !*IS_STDOUT_TERMINAL {
bail!("Failed to init rag in non-interactive mode"); bail!("Failed to init rag in non-interactive mode");
} }
println!("⚙ Initializing RAG..."); println!("⚙ Initializing RAG...");
let (embedding_model, chunk_size, chunk_overlap) = Self::create_config(app)?; 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 reranker_model = app.rag_reranker_model.clone();
let top_k = app.rag_top_k; let top_k = app.rag_top_k;
let extractor_model = match app.rag_extractor_model.clone() { let extractor_model = match app.rag_extractor_model.clone() {
@@ -275,7 +307,7 @@ impl Rag {
app.rag_graph_hops app.rag_graph_hops
}; };
let extractor_prompt = app.rag_extractor_prompt.clone(); let extractor_prompt = app.rag_extractor_prompt.clone();
let data = RagData::new( let mut data = RagData::new(
embedding_model.id(), embedding_model.id(),
chunk_size, chunk_size,
chunk_overlap, chunk_overlap,
@@ -288,6 +320,7 @@ impl Rag {
graph_hops: Some(graph_hops), graph_hops: Some(graph_hops),
}, },
); );
data.driver = driver.to_string();
let mut rag = Self::create(app, name, save_path, data)?; let mut rag = Self::create(app, name, save_path, data)?;
let mut paths = doc_paths.to_vec(); let mut paths = doc_paths.to_vec();
if paths.is_empty() { if paths.is_empty() {
@@ -315,9 +348,50 @@ impl Rag {
Self::create(app, name, path, data) Self::create(app, name, path, data)
} }
pub fn create(app: &AppConfig, name: &str, path: &Path, data: RagData) -> Result<Self> { /// `mut data` — the duckdb arm rehydrates `data.vectors` from the sidecar.
let bm25 = data.build_bm25(); pub fn create(app: &AppConfig, name: &str, path: &Path, mut data: RagData) -> Result<Self> {
let provider: Box<dyn RagProvider> = Box::new(YamlProvider::from_data(&data)); // 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<dyn RagProvider>, _) = 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 node_to_docs = data.knowledge_graph.build_node_to_docs();
let embedding_model = let embedding_model =
Model::retrieve_model(app, &data.embedding_model, ModelType::Embedding)?; Model::retrieve_model(app, &data.embedding_model, ModelType::Embedding)?;
@@ -460,8 +534,18 @@ impl Rag {
let path = Path::new(&self.path); let path = Path::new(&self.path);
ensure_parent_exists(path)?; ensure_parent_exists(path)?;
let content = serde_yaml::to_string(&self.data) let content = if self.data.driver == "duckdb" {
.with_context(|| format!("Failed to serde rag '{}'", self.name))?; // 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(|| { fs::write(path, content).with_context(|| {
format!("Failed to save rag '{}' to '{}'", self.name, path.display()) format!("Failed to save rag '{}' to '{}'", self.name, path.display())
})?; })?;
@@ -841,13 +925,18 @@ impl Rag {
} }
progress(&spinner, "Building store".into()); 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 / // `refresh` is true for a full re-index (.rebuild rag / --rebuild-rag /
// initial build) and false for an incremental .edit rag-docs change. // initial build) and false for an incremental .edit rag-docs change.
// Passing it through is what stops a remote provider from wiping its // Passing it through is what stops a remote provider from wiping its
// collection on a one-file add. // collection on a one-file add.
self.provider.rebuild_indexes(&self.data, refresh).await?; 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(()) Ok(())
} }
@@ -1660,10 +1749,39 @@ fn reciprocal_rank_fusion(
.collect() .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)] #[cfg(test)]
mod tests { mod tests {
use super::*; 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] #[test]
fn document_id_round_trip() { fn document_id_round_trip() {
let id = DocumentId::new(5, 17); let id = DocumentId::new(5, 17);
+964
View File
@@ -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<Mutex<Connection>>,
/// 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<Self> {
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<IndexMap<DocumentId, Vec<f32>>> {
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<f32>)> = stmt
.query_map([], |row| {
let id: u64 = row.get(0)?;
// The duckdb crate does not implement `FromSql` for `Vec<f32>`, 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<f32> = 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::<duckdb::Result<Vec<_>>>()
.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<MutexGuard<'_, Connection>> {
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<Vec<(DocumentId, f32)>> {
// 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::<Vec<_>>()
.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<Vec<(DocumentId, String)>> {
if ids.is_empty() {
return Ok(vec![]);
}
let conn = self.lock_conn()?;
let placeholders = ids.iter().map(|_| "?").collect::<Vec<_>>().join(", ");
let sql =
format!("SELECT doc_id, page_content FROM documents WHERE doc_id IN ({placeholders})");
let params: Vec<duckdb::types::Value> = 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<DocumentId, usize> =
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<Vec<(DocumentId, f32)>> {
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<dyn RagProvider> {
// 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 `<path>.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"));
}
}
+7
View File
@@ -4,3 +4,10 @@ mod yaml;
// is ambiguous (E0659) — `use` paths resolve against both this module's items // is ambiguous (E0659) — `use` paths resolve against both this module's items
// and the extern prelude, and `use` declarations may not shadow. // and the extern prelude, and `use` declarations may not shadow.
pub use self::yaml::YamlProvider; 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;