refactor(rag): extract RagProvider trait and add YamlProvider
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:
@@ -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;
|
||||
@@ -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");
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user