Dark-Alex-17/coyote
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

refactor(rag): extract RagProvider trait and add YamlProvider

Browse Source 
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).
This commit is contained in:
Dark-Alex-17
2026-08-10 11:41:37 -06:00
parent a968c3228d
commit 5049143fcc
4 changed files with 526 additions and 76 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/rag/mod.rs
+187 -74
View File
@@ -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<DocumentId>,
// Vector storage + content retrieval.
provider: Box<dyn RagProvider>,
data: RagData,
last_sources: RwLock<Option<String>>,
node_to_docs: IndexMap<u32, Vec<DocumentId>>,
@@ -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<Connection>` 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<Rag>` — 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<Self> {
let hnsw = data.build_hnsw();
let bm25 = data.build_bm25();
let provider: Box<dyn RagProvider> = 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<Vec<(DocumentId, String)>> {
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<DocumentId> =
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<DocumentId> =
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) {
// `ids` is an `IndexSet` here, not the `Vec` of the RRF branch below,
// and `&IndexSet<_>` does not coerce to `&[DocumentId]`.
let ids: Vec<DocumentId> = 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(document.page_content.to_string());
}
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
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);
}
})
.collect::<Vec<_>>()
})
.collect();
Ok(output)
Ok(merge_vector_results(results))
}
async fn keyword_search(
&self,
query: &str,
top_k: usize,
min_score: f32,
) -> Result<Vec<(DocumentId, f32)>> {
/// 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<DocumentId> {
@@ -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<Item = (DocumentId, &RagDocument)> {
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<FileId>) {
@@ -1285,13 +1323,12 @@ impl RagData {
}
pub fn build_bm25(&self) -> SearchEngine<DocumentId> {
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::<DocumentId>::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<Vec<DocumentId>>,
list_of_weights: Vec<f32>,
@@ -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<DocumentId> = 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<FileId, RagFile>, IndexMap<String, Vec<FileId>>) {
let mut files: IndexMap<FileId, RagFile> = Default::default();
files.insert(
src/rag/provider.rs
+87
View File
@@ -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<Vec<(DocumentId, f32)>>;
/// 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<Vec<(DocumentId, String)>>;
/// 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<Vec<(DocumentId, f32)>> {
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<dyn RagProvider> 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<dyn RagProvider>;
}
src/rag/providers/mod.rs
+6
View File
@@ -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;
src/rag/providers/yaml.rs
+244
View File
@@ -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<DocumentId, String>,
}
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<DocumentId, String> {
// 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<Vec<(DocumentId, f32)>> {
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<Vec<(DocumentId, String)>> {
// 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<dyn RagProvider> {
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");
}
}