From 5049143fcc6e4684ea29cb62b6c78e47ddc4893e Mon Sep 17 00:00:00 2001 From: Alex Clarke Date: Mon, 10 Aug 2026 11:41:37 -0600 Subject: [PATCH] refactor(rag): extract RagProvider trait and add YamlProvider MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Introduce a narrow `RagProvider` trait covering vector search and content retrieval, and make `Rag` delegate to a boxed provider instead of owning an HNSW index directly. `YamlProvider` is the sole implementation for now. The trait deliberately stays narrow: embeddings, chunking, BM25 keyword search, graph RAG, entity extraction, RRF merging and persistence all remain on `Rag`/`RagData`, so a new storage backend does not have to reimplement Coyote's indexing logic. Notable points: - `fetch_content`'s ordering contract is part of the trait, not an accident. Implementations must return results in input-`ids` order; `hybrid_search` passes an RRF-ranked list straight to the prompt builder, so a provider returning storage order would silently discard the ranking. - The content store is keyed on `data.files`, never `data.vectors`. Both the content map and BM25 now route through the new `RagData::iter_documents()` so the two key spaces match by construction. `RagData::add` zips document ids with embeddings and truncates silently, so ids in `files \ vectors` are genuinely reachable. - A provider keyword-search failure degrades to an empty ranker with a warning rather than failing the whole query; it is one of three RRF inputs. It deliberately does not fall back to the local BM25, which would be a silent ranking-algorithm swap once a provider with native FTS exists. - The rerank path builds its text and id vectors from a single `fetch_content` result in one pass, so the reranker's positional indices cannot desync. This is not a bit-for-bit no-op. `vector_search` now dedups by best score and sorts globally instead of concatenating per-chunk hit lists. Single-chunk queries (the common case) are unaffected. Multi-chunk queries get corrected rank assignment and no longer let a document that matched several query chunks accumulate multiple RRF contributions. There is no overall cap on the merged pool — truncation remains `reciprocal_rank_fusion`'s job. `RagData::get()` is removed: its only two callers were the content lookups replaced here, and an unused private-module method fails the build under `--deny warnings`. Its three tests were rewritten against `iter_documents()`, one of which now guards the files-vs-vectors keying directly. Implements Phase 2 of the RAG driver abstraction design (§6). --- src/rag/mod.rs | 265 +++++++++++++++++++++++++++----------- src/rag/provider.rs | 87 +++++++++++++ src/rag/providers/mod.rs | 6 + src/rag/providers/yaml.rs | 244 +++++++++++++++++++++++++++++++++++ 4 files changed, 526 insertions(+), 76 deletions(-) create mode 100644 src/rag/provider.rs create mode 100644 src/rag/providers/mod.rs create mode 100644 src/rag/providers/yaml.rs diff --git a/src/rag/mod.rs b/src/rag/mod.rs index 1c91f68..9011296 100644 --- a/src/rag/mod.rs +++ b/src/rag/mod.rs @@ -5,10 +5,14 @@ use crate::config::*; use crate::utils::*; mod graph; +mod provider; +mod providers; mod serde_vectors; mod splitter; use self::graph::{KnowledgeGraph, extract_entities}; +use self::provider::RagProvider; +use self::providers::YamlProvider; use anyhow::{Context, Result, anyhow, bail}; use bm25::{Language, SearchEngine, SearchEngineBuilder}; @@ -56,8 +60,11 @@ pub struct Rag { name: String, path: String, embedding_model: Model, - hnsw: Hnsw<'static, f32, DistCosine>, + // Local BM25: keyword search + graph seeding. Always built from `data.files` + // regardless of driver, and kept on `Rag` so the sync `graph_search` can use it. bm25: SearchEngine, + // Vector storage + content retrieval. + provider: Box, data: RagData, last_sources: RwLock>, node_to_docs: IndexMap>, @@ -74,6 +81,19 @@ impl Debug for Rag { } } +// CLONING A `Rag` DOES NOT SNAPSHOT ITS BACKING STORE. +// +// `provider.duplicate(&self.data)` is a true snapshot for YamlProvider only. +// DuckDbProvider Arc-clones one shared `Mutex` over one file, and +// QdrantProvider addresses the same remote collection. So for those drivers the +// clone and the original are two views of ONE store. +// +// INVARIANT: after calling `rebuild_indexes` on a cloned `Rag`, the pre-clone +// instance MUST be discarded immediately and MUST NOT serve further queries. +// Cloning to READ is always fine; cloning to REBUILD makes the original a +// half-truth (pre-rebuild `data`, post-rebuild store). Note that reassigning +// `RequestContext.rag` drops only one holder of the old `Arc` — forked +// request contexts, agents, captured inputs and the RAG cache keep theirs. impl Clone for Rag { fn clone(&self) -> Self { Self { @@ -81,8 +101,8 @@ impl Clone for Rag { name: self.name.clone(), path: self.path.clone(), embedding_model: self.embedding_model.clone(), - hnsw: self.data.build_hnsw(), bm25: self.data.build_bm25(), + provider: self.provider.duplicate(&self.data), node_to_docs: self.data.knowledge_graph.build_node_to_docs(), data: self.data.clone(), last_sources: RwLock::new(None), @@ -296,8 +316,8 @@ impl Rag { } pub fn create(app: &AppConfig, name: &str, path: &Path, data: RagData) -> Result { - let hnsw = data.build_hnsw(); let bm25 = data.build_bm25(); + let provider: Box = Box::new(YamlProvider::from_data(&data)); let node_to_docs = data.knowledge_graph.build_node_to_docs(); let embedding_model = Model::retrieve_model(app, &data.embedding_model, ModelType::Embedding)?; @@ -307,8 +327,8 @@ impl Rag { path: path.display().to_string(), data, embedding_model, - hnsw, bm25, + provider, node_to_docs, last_sources: RwLock::new(None), }; @@ -821,7 +841,11 @@ impl Rag { } progress(&spinner, "Building store".into()); - self.hnsw = self.data.build_hnsw(); + // `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(); @@ -834,17 +858,30 @@ impl Rag { top_k: usize, rerank_model: Option<&str>, ) -> Result> { - let (vector_search_results, keyword_search_results) = tokio::join!( - self.vector_search(query, top_k, 0.0), - self.keyword_search(query, top_k, 0.0), - ); - - let vector_search_results = vector_search_results?; + let vector_search_results = self.vector_search(query, top_k, 0.0).await?; debug!("vector_search_results: {vector_search_results:?}",); let vector_search_ids: Vec = vector_search_results.into_iter().map(|(v, _)| v).collect(); - let keyword_search_results = keyword_search_results?; + let keyword_search_results: Vec<(DocumentId, f32)> = + if self.provider.has_native_keyword_search() { + // Keyword is ONE of three RRF rankers (vector + keyword + graph); + // its absence is survivable and produces a slightly worse ranking, + // whereas a `?` here turns a provider FTS fault into TOTAL query + // failure — the user gets an error instead of the results the + // vector and graph rankers already retrieved. Degrade, do not + // propagate, and do not silently swap in the local BM25 either: + // that would change the ranking algorithm mid-query. + match self.provider.keyword_search(query, top_k).await { + Ok(v) => v, + Err(e) => { + warn!("native keyword search failed, dropping the keyword ranker: {e}"); + Vec::new() + } + } + } else { + self.keyword_search(query, top_k, 0.0) + }; debug!("keyword_search_results: {keyword_search_results:?}",); let keyword_search_ids: Vec = keyword_search_results.into_iter().map(|(v, _)| v).collect(); @@ -857,13 +894,20 @@ impl Rag { .concat() .into_iter() .collect(); - let mut documents = vec![]; - let mut documents_ids = vec![]; - for id in ids { - if let Some(document) = self.data.get(id) { - documents_ids.push(id); - documents.push(document.page_content.to_string()); - } + // `ids` is an `IndexSet` here, not the `Vec` of the RRF branch below, + // and `&IndexSet<_>` does not coerce to `&[DocumentId]`. + let ids: Vec = ids.into_iter().collect(); + let fetched = self.provider.fetch_content(&ids).await?; + // Build both vectors from the SAME source in the SAME iteration — + // never zip two independently-built lists. The reranker returns + // positional indices into `documents`, so any drift between the two + // resolves reranked hits to the wrong document's text. A partial + // fetch simply yields a shorter pair, and both shrink together. + let mut documents_ids = Vec::with_capacity(fetched.len()); + let mut documents = Vec::with_capacity(fetched.len()); + for (id, text) in fetched { + documents_ids.push(id); + documents.push(text); } let data = RerankData::new(query.to_string(), documents, top_k); let list = client.rerank(&data).await.context("Failed to rerank")?; @@ -895,13 +939,9 @@ impl Rag { ids } }; - let output = ids - .into_iter() - .filter_map(|id| { - let document = self.data.get(id)?; - Some((id, document.page_content.clone())) - }) - .collect(); + // `ids` is the ranked list; `fetch_content` preserves that order per the + // trait's ordering contract, so the result is returned as-is. + let output = self.provider.fetch_content(&ids).await?; Ok(output) } @@ -918,35 +958,24 @@ impl Rag { ); let texts = splitter.split_text(query); let embeddings_data = EmbeddingsData::new(texts, true); - let embeddings = self.create_embeddings(embeddings_data, None).await?; - let output = self - .hnsw - .parallel_search(&embeddings, top_k, 30) - .into_iter() - .flat_map(|list| { - list.into_iter() - .filter_map(|v| { - let score = 1.0 - v.distance; - if score > min_score { - Some((DocumentId(v.d_id), score)) - } else { - None - } - }) - .collect::>() - }) - .collect(); - Ok(output) + let query_embeddings = self.create_embeddings(embeddings_data, None).await?; + + let mut results: Vec<(DocumentId, f32)> = vec![]; + for embedding in &query_embeddings { + let batch = self + .provider + .vector_search(embedding, top_k, min_score) + .await?; + results.extend(batch); + } + Ok(merge_vector_results(results)) } - async fn keyword_search( - &self, - query: &str, - top_k: usize, - min_score: f32, - ) -> Result> { + /// Local in-memory BM25 over `data.files` — empty for attached RAGs, which is + /// correct: they have no local text. + fn keyword_search(&self, query: &str, top_k: usize, min_score: f32) -> Vec<(DocumentId, f32)> { let results = self.bm25.search(query, top_k); - let output: Vec<(DocumentId, f32)> = results + results .into_iter() .filter_map(|v| { let score = v.score; @@ -956,8 +985,7 @@ impl Rag { None } }) - .collect(); - Ok(output) + .collect() } fn graph_search(&self, query: &str, top_k: usize) -> Vec { @@ -1243,11 +1271,21 @@ impl RagData { } } - pub fn get(&self, id: DocumentId) -> Option<&RagDocument> { - let (file_index, document_index) = id.split(); - let file = self.files.get(&file_index)?; - let document = file.documents.get(document_index)?; - Some(document) + /// Every (DocumentId, &RagDocument) in the corpus, in `files` order. + /// + /// This — NOT `vectors` — is the authoritative document id space. BM25, the + /// knowledge graph and content lookup all key off it; `vectors` is a subset, + /// since `add`'s zip truncates whenever fewer embeddings come back than + /// document ids were sent. + pub fn iter_documents(&self) -> impl Iterator { + self.files.iter().flat_map(|(file_index, file)| { + file.documents + .iter() + .enumerate() + .map(move |(document_index, document)| { + (DocumentId::new(*file_index, document_index), document) + }) + }) } pub fn del(&mut self, file_ids: Vec) { @@ -1285,13 +1323,12 @@ impl RagData { } pub fn build_bm25(&self) -> SearchEngine { - let mut documents = vec![]; - for (file_index, file) in self.files.iter() { - for (document_index, document) in file.documents.iter().enumerate() { - let id = DocumentId::new(*file_index, document_index); - documents.push(bm25::Document::new(id, &document.page_content)) - } - } + // Shares `iter_documents` with the providers' content maps so the BM25 key + // space and the content key space are identical by construction. + let documents: Vec<_> = self + .iter_documents() + .map(|(id, doc)| bm25::Document::new(id, &doc.page_content)) + .collect(); SearchEngineBuilder::::with_documents(Language::English, documents) .k1(1.5) .b(0.75) @@ -1580,6 +1617,27 @@ fn find_hash_skip( .map(|(i, v)| (i, *v)) } +/// Global score sort + dedup keeping the best score per document. +/// +/// NO overall cap: each `provider.vector_search` call already returns <= top_k, +/// so the pool is bounded by `top_k * query_chunks`, and `reciprocal_rank_fusion` +/// truncates to `top_k` itself. Capping here would let whichever query chunk has +/// the strongest absolute scores crowd out every other chunk's hits. +/// +/// Free function (not a method) so it is unit-testable without an embeddings client. +fn merge_vector_results(mut results: Vec<(DocumentId, f32)>) -> Vec<(DocumentId, f32)> { + debug_assert!( + results.iter().all(|(_, score)| score.is_finite()), + "provider returned a non-finite score; NaN silently degrades sort order" + ); + results.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(Ordering::Equal)); + let mut seen = IndexSet::new(); + results + .into_iter() + .filter(|(id, _)| seen.insert(*id)) + .collect() +} + fn reciprocal_rank_fusion( list_of_document_ids: Vec>, list_of_weights: Vec, @@ -1695,7 +1753,7 @@ mod tests { } #[test] - fn rag_data_get_returns_document() { + fn rag_data_iter_documents_yields_all_documents_in_file_order() { let mut data = RagData::new( "m".into(), 100, @@ -1712,15 +1770,21 @@ mod tests { }; data.files.insert(0, file); - let doc = data.get(DocumentId::new(0, 0)).unwrap(); - assert_eq!(doc.page_content, "first"); - - let doc = data.get(DocumentId::new(0, 1)).unwrap(); - assert_eq!(doc.page_content, "second"); + let documents: Vec<_> = data + .iter_documents() + .map(|(id, doc)| (id, doc.page_content.as_str())) + .collect(); + assert_eq!( + documents, + vec![ + (DocumentId::new(0, 0), "first"), + (DocumentId::new(0, 1), "second"), + ] + ); } #[test] - fn rag_data_get_returns_none_for_missing_file() { + fn rag_data_iter_documents_is_empty_without_files() { let data = RagData::new( "m".into(), 100, @@ -1730,11 +1794,14 @@ mod tests { None, GraphRagConfig::default(), ); - assert!(data.get(DocumentId::new(99, 0)).is_none()); + assert_eq!(data.iter_documents().count(), 0); } + /// The document id space is `files`, never `vectors`: `add`'s zip truncates + /// silently, so a vector may exist for an id no file provides. Content lookup + /// and BM25 both key off this iterator and must agree. #[test] - fn rag_data_get_returns_none_for_missing_document() { + fn rag_data_iter_documents_ignores_vector_only_ids() { let mut data = RagData::new( "m".into(), 100, @@ -1750,7 +1817,10 @@ mod tests { documents: vec![RagDocument::new("only one")], }; data.files.insert(0, file); - assert!(data.get(DocumentId::new(0, 5)).is_none()); + data.vectors.insert(DocumentId::new(0, 5), vec![1.0]); + + let ids: Vec<_> = data.iter_documents().map(|(id, _)| id).collect(); + assert_eq!(ids, vec![DocumentId::new(0, 0)]); } #[test] @@ -1998,6 +2068,49 @@ mod tests { ); } + #[test] + fn merge_vector_results_empty_input() { + let result = super::merge_vector_results(vec![]); + assert!(result.is_empty(), "empty input should produce empty output"); + } + + #[test] + fn merge_vector_results_keeps_best_score_per_document() { + let doc = DocumentId::new(0, 0); + let result = super::merge_vector_results(vec![(doc, 0.2), (doc, 0.9)]); + assert_eq!(result.len(), 1, "a document must not be double-counted"); + assert_eq!(result[0].0, doc); + assert_eq!( + result[0].1, 0.9, + "dedup must keep the highest score, not the first seen" + ); + } + + #[test] + fn merge_vector_results_sorts_globally_by_descending_score() { + let doc_a = DocumentId::new(0, 0); + let doc_b = DocumentId::new(1, 0); + let doc_c = DocumentId::new(2, 0); + // Interleaved as two per-chunk hit lists would arrive: concatenating them + // would yield a, c, b — only a global sort produces c, a, b. + let result = super::merge_vector_results(vec![(doc_a, 0.5), (doc_c, 0.9), (doc_b, 0.1)]); + let ids: Vec = result.iter().map(|(id, _)| *id).collect(); + assert_eq!(ids, vec![doc_c, doc_a, doc_b]); + } + + #[test] + fn merge_vector_results_does_not_truncate() { + let input: Vec<(DocumentId, f32)> = (0..10) + .map(|i| (DocumentId::new(i, 0), i as f32 / 10.0)) + .collect(); + let result = super::merge_vector_results(input); + assert_eq!( + result.len(), + 10, + "merging must not cap the pool; truncation belongs to reciprocal_rank_fusion" + ); + } + fn hash_skip_fixture() -> (IndexMap, IndexMap>) { let mut files: IndexMap = Default::default(); files.insert( diff --git a/src/rag/provider.rs b/src/rag/provider.rs new file mode 100644 index 0000000..6889e81 --- /dev/null +++ b/src/rag/provider.rs @@ -0,0 +1,87 @@ +use super::{DocumentId, RagData}; +use anyhow::Result; +use async_trait::async_trait; + +/// Abstracts where RAG vector data is stored and queried. +/// +/// Implementors: +/// - YamlProvider: HNSW in-memory, state derived from RagData.vectors/files +/// - DuckDbProvider: DuckDB on-disk vector index + document store +/// - QdrantProvider: remote Qdrant collection +/// +/// The Rag orchestrator owns: embeddings, chunking, BM25 keyword search, graph RAG, +/// entity extraction, RRF merging. Providers own: vector storage and content retrieval. +#[async_trait] +pub trait RagProvider: Send + Sync { + /// Vector similarity search. Returns (DocumentId, score) sorted by score desc. + /// `embedding` is a single query vector from Coyote's embedding model. + async fn vector_search( + &self, + embedding: &[f32], + top_k: usize, + min_score: f32, + ) -> Result>; + + /// Resolve document IDs to their page content. + /// + /// **Ordering contract:** implementations MUST return results in the same + /// relative order as the input `ids` slice. `hybrid_search` passes an + /// RRF-ranked list and feeds the result straight to the LLM — a provider + /// that returns rows in storage order (e.g. Qdrant `get_points`, DuckDB + /// `WHERE id IN (...)`) would silently discard the ranking. Implementations + /// that query an unordered backend must re-sort by input position before + /// returning. + /// + /// Returns only IDs that were found; callers must handle partial returns + /// (a missing ID is skipped, not an error). + /// YamlProvider: reads from an in-memory content map built from data.files. + /// DuckDbProvider: queries the documents table by id. + /// QdrantProvider: fetches payload from the remote collection. + async fn fetch_content(&self, ids: &[DocumentId]) -> Result>; + + /// Rebuild internal indexes from freshly updated RagData. + /// Called once at the end of every sync_documents pass. + /// + /// `full_rebuild` mirrors `sync_documents`' `refresh` parameter: + /// - `true` — a full re-index (`.rebuild rag`, `--rebuild-rag`, initial build). + /// Destructive strategies (wipe-then-reindex) are permitted. + /// - `false` — an incremental change (`.edit rag-docs` adding/removing a file). + /// Implementations MUST NOT wipe existing state; upsert only. + /// + /// The parameter is part of the signature from the outset so it is fixed + /// while there is exactly one implementor. Yaml/DuckDb ignore it — + /// rebuilding their local state wholesale is fast and always correct. + /// Only a remote provider is destructive enough to care. + /// + /// YamlProvider: rebuilds HNSW + content map from data.vectors/files. + /// DuckDbProvider: writes new rows to DuckDB, deletes removed rows. + /// QdrantProvider: no-op while attach-only — remote data is unchanged. + async fn rebuild_indexes(&mut self, data: &RagData, full_rebuild: bool) -> Result<()>; + + /// Keyword / full-text search. Returns (DocumentId, BM25-style score) sorted desc. + /// + /// Default impl returns `Ok(vec![])` — callers fall back to `Rag.bm25` (local in-memory + /// BM25 built from `data.files`). DuckDbProvider overrides this with a native FTS query + /// (DuckDB's `fts` extension, installed once at schema-creation time). + /// + /// Callers check `has_native_keyword_search()` before deciding which path to take: + /// - true → call this method; skip `Rag.bm25` + /// - false → call `Rag.keyword_search()` which uses `Rag.bm25` (sync, infallible) + /// + /// YamlProvider and QdrantProvider do NOT override this (return empty). + async fn keyword_search(&self, query: &str, top_k: usize) -> Result> { + let _ = (query, top_k); + Ok(vec![]) + } + + /// Returns true if this provider implements a native keyword-search index. + /// When false, `Rag.hybrid_search` uses the local `Rag.bm25` field instead. + fn has_native_keyword_search(&self) -> bool { + false + } + + /// Deep-clone the provider with fresh indexes derived from `data`. + /// Required because Box is not Clone. + /// Called by Rag's Clone impl (which clones before mutating in rebuild_rag/edit_rag_docs). + fn duplicate(&self, data: &RagData) -> Box; +} diff --git a/src/rag/providers/mod.rs b/src/rag/providers/mod.rs new file mode 100644 index 0000000..509f02c --- /dev/null +++ b/src/rag/providers/mod.rs @@ -0,0 +1,6 @@ +mod yaml; +// Use `self::` on every re-export in this file. Once a `mod duckdb;` sits here +// alongside a dependency on the `duckdb` CRATE, a bare `pub use duckdb::...` +// 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; diff --git a/src/rag/providers/yaml.rs b/src/rag/providers/yaml.rs new file mode 100644 index 0000000..df1f6f5 --- /dev/null +++ b/src/rag/providers/yaml.rs @@ -0,0 +1,244 @@ +use crate::rag::provider::RagProvider; +use crate::rag::{DocumentId, RagData}; + +use anyhow::Result; +use async_trait::async_trait; +use hnsw_rs::prelude::*; +use indexmap::IndexMap; + +pub struct YamlProvider { + hnsw: Hnsw<'static, f32, DistCosine>, + content_map: IndexMap, +} + +impl YamlProvider { + pub fn from_data(data: &RagData) -> Self { + Self { + hnsw: data.build_hnsw(), + content_map: Self::build_content_map(data), + } + } + + fn build_content_map(data: &RagData) -> IndexMap { + // Keyed on `files`, NOT `vectors`: this is the exact replacement for the + // per-id document lookup it supersedes, and it must resolve every id that + // BM25 or graph_search can produce — both of which enumerate `files`. + data.iter_documents() + .map(|(id, doc)| (id, doc.page_content.clone())) + .collect() + } +} + +#[async_trait] +impl RagProvider for YamlProvider { + async fn vector_search( + &self, + embedding: &[f32], + top_k: usize, + min_score: f32, + ) -> Result> { + let results = self + .hnsw + .parallel_search(&[embedding.to_vec()], top_k, 30) + .into_iter() + .flat_map(|list| { + list.into_iter().filter_map(|v| { + let score = 1.0 - v.distance; + if score > min_score { + Some((DocumentId(v.d_id), score)) + } else { + None + } + }) + }) + .collect(); + Ok(results) + } + + async fn fetch_content(&self, ids: &[DocumentId]) -> Result> { + // Iterating `ids` (not `content_map`) satisfies the trait's ordering + // contract for free — output order mirrors input order. + Ok(ids + .iter() + .filter_map(|id| self.content_map.get(id).map(|text| (*id, text.clone()))) + .collect()) + } + + async fn rebuild_indexes(&mut self, data: &RagData, _full_rebuild: bool) -> Result<()> { + // Local in-memory state — a wholesale rebuild is fast and always correct, + // so the incremental/full distinction is irrelevant here. + self.hnsw = data.build_hnsw(); + self.content_map = Self::build_content_map(data); + Ok(()) + } + + fn duplicate(&self, data: &RagData) -> Box { + Box::new(YamlProvider::from_data(data)) + } +} + +#[cfg(test)] +mod provider_tests { + use super::*; + // `RagFile` and `RagDocument` are not used by the impl above, so they are + // imported here rather than at module scope. Both have private fields, which + // is why these tests must live in-crate rather than under `tests/`. + use crate::rag::{RagDocument, RagFile}; + + fn minimal_rag_data() -> RagData { + // `..Default::default()` rather than an exhaustive struct literal so that + // later additions to `RagData` do not break this helper. + RagData { + embedding_model: "text-embedding-3-small".to_string(), + chunk_size: 1024, + chunk_overlap: 50, + top_k: 5, + driver: "yaml".to_string(), + attached: false, + ..Default::default() + } + } + + /// Two files, one chunk each, with vectors — the minimum needed to exercise + /// `build_content_map` and the `fetch_content` ordering contract. + /// `DocumentId::new(f, d)` packs (file_index, document_index); `RagData::add` + /// is the real insertion path but a direct literal is sufficient and avoids + /// the embedding pipeline. + fn populated_rag_data() -> RagData { + let mut data = minimal_rag_data(); + // `files` must be populated: build_content_map iterates data.iter_documents(), + // which enumerates `files`. Populating `vectors` alone would produce an EMPTY + // content map, and every assertion below would vacuously pass on a broken impl. + // The vectors inserted at the end are for the HNSW side only. + data.files.insert( + 0, + RagFile { + hash: "h0".to_string(), + path: "/tmp/a.md".to_string(), + documents: vec![RagDocument { + page_content: "alpha".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".to_string(), + metadata: Default::default(), + }], + }, + ); + data.vectors + .insert(DocumentId::new(0, 0), vec![1.0, 0.0, 0.0]); + data.vectors + .insert(DocumentId::new(1, 0), vec![0.0, 1.0, 0.0]); + data + } + + #[tokio::test] + async fn yaml_provider_empty_data_returns_nothing() { + let data = minimal_rag_data(); + let provider = YamlProvider::from_data(&data); + let results = provider.fetch_content(&[]).await.unwrap(); + assert!(results.is_empty()); + } + + /// `fetch_content` MUST return results in the same relative order as the input + /// ids. The reversed-input case is the one that fails if an implementation ever + /// iterates its own map instead of `ids`. + #[tokio::test] + async fn yaml_provider_fetch_content_preserves_input_order() { + let data = populated_rag_data(); + let provider = YamlProvider::from_data(&data); + + let a = DocumentId::new(0, 0); + let b = DocumentId::new(1, 0); + + let forward = provider.fetch_content(&[a, b]).await.unwrap(); + assert_eq!(forward.len(), 2, "both documents must resolve"); + assert_eq!(forward[0].1, "alpha"); + assert_eq!(forward[1].1, "beta"); + + // Reversed input must produce reversed output — NOT storage order. + let reversed = provider.fetch_content(&[b, a]).await.unwrap(); + assert_eq!( + reversed[0].1, "beta", + "fetch_content must honor input order" + ); + assert_eq!(reversed[1].1, "alpha"); + } + + /// A missing id is skipped, not an error, and does not disturb the order of + /// the ids that DO resolve. + #[tokio::test] + async fn yaml_provider_fetch_content_skips_missing_ids() { + let data = populated_rag_data(); + let provider = YamlProvider::from_data(&data); + + let a = DocumentId::new(0, 0); + let missing = DocumentId::new(99, 0); + let b = DocumentId::new(1, 0); + + let out = provider.fetch_content(&[a, missing, b]).await.unwrap(); + assert_eq!(out.len(), 2, "missing id is skipped, not an error"); + assert_eq!(out[0].1, "alpha"); + assert_eq!(out[1].1, "beta"); + } + + /// `YamlProvider::duplicate()` rebuilds from `data`, so the clone is a genuine + /// independent snapshot. Providers backed by a shared store deliberately are not. + #[tokio::test] + async fn yaml_provider_duplicate_returns_equivalent_content() { + // MUST be populated_rag_data(): on minimal_rag_data() both providers hold an + // EMPTY content map, so `assert_eq!(r1, r2)` compares two empty vectors and + // passes against a duplicate() that returns nothing at all. + let data = populated_rag_data(); + let provider = YamlProvider::from_data(&data); + let dup = provider.duplicate(&data); + let ids = [DocumentId::new(0, 0), DocumentId::new(1, 0)]; + // Query with REAL ids, not `&[]` — an empty slice is answered without ever + // touching the content map, so it would pass against a broken duplicate(). + let r1 = provider.fetch_content(&ids).await.unwrap(); + let r2 = dup.fetch_content(&ids).await.unwrap(); + // Guard against the vacuous case: if both sides resolved nothing, the equality + // below proves nothing. Assert the fixture actually produced content first. + assert_eq!(r1.len(), 2, "fixture must resolve both documents"); + assert_eq!( + r1, r2, + "duplicate must resolve the same content as the original" + ); + } + + /// The content store is keyed on `files`, never on `vectors`. A vector may exist + /// for an id with no backing file (a stale entry, or a file dropped mid-sync); + /// keying on `vectors` would surface such an id with empty text instead of + /// dropping it. The shared fixture only ever inserts vectors for ids that also + /// have files, so this case has to be constructed here. + #[tokio::test] + async fn yaml_provider_content_is_keyed_on_files_not_vectors() { + let mut data = populated_rag_data(); + let orphan = DocumentId::new(9, 0); + data.vectors.insert(orphan, vec![0.0, 0.0, 1.0]); + + let provider = YamlProvider::from_data(&data); + + let out = provider.fetch_content(&[orphan]).await.unwrap(); + assert!( + out.is_empty(), + "an id present only in `vectors` must not resolve to content" + ); + + // The file-backed ids still resolve, so the assertion above is not vacuous. + let real = provider + .fetch_content(&[DocumentId::new(0, 0), DocumentId::new(1, 0)]) + .await + .unwrap(); + assert_eq!(real.len(), 2, "file-backed documents must still resolve"); + assert_eq!(real[0].1, "alpha"); + assert_eq!(real[1].1, "beta"); + } +}