diff --git a/README.md b/README.md index b06cd82..cfae5c1 100644 --- a/README.md +++ b/README.md @@ -251,6 +251,7 @@ ctx.checkout(first_version) print("Entries after checkout:", ctx.entries()) +<<<<<<< HEAD # Curate stored records into a trainable dataset and export it as JSONL plus a # reproducible manifest. Curation (lifecycle-correct filtering, semantic dedup, # decontamination against a holdout set, reward thresholding) runs before @@ -282,6 +283,29 @@ ctx.export_training( split={"eval_fraction": 0.1, "by": "session_id", "seed": 42}, ) +# Measure retrieval quality against a labeled query set. Each query lists the +# relevant records by stable external_id (with optional graded relevance), and +# the report carries recall@k / precision@k / MRR / nDCG@k / hit-rate plus a +# manifest (version, k, mode, distance_metric) for reproducibility. +report = ctx.evaluate( + [ + { + "query_id": "q1", + "vector": query_embedding, # vector channel + "relevant": [{"external_id": "doc-77#chunk-1", "grade": 1.0}], + }, + ], + k=10, + mode="vector", # or "hybrid" for text+vector +) +print("recall@10:", report["aggregate"]["recall"]) + +# A/B the same query set across two dataset versions (regression detection that +# a stateless vector DB can't do) and read per-metric deltas. +ab = ctx.evaluate_versions(query_set, baseline_version, candidate_version, k=10) +print("nDCG delta:", ab["deltas"]["ndcg"]) + + # Remote persistence on any object_store backend uses a generic `storage_options` # dict, matching the conventions used by `lance` and `lance-graph`. # diff --git a/crates/lance-context-core/src/eval.rs b/crates/lance-context-core/src/eval.rs new file mode 100644 index 0000000..5152798 --- /dev/null +++ b/crates/lance-context-core/src/eval.rs @@ -0,0 +1,791 @@ +//! Retrieval-quality evaluation harness. +//! +//! Measures retrieval quality (recall@k / precision@k / MRR / nDCG@k / +//! hit-rate) of [`ContextStore::search_filtered_with_options`] (vector) and +//! [`ContextStore::retrieve_filtered_with_options`] (hybrid) against a labeled +//! query set, and compares quality across dataset versions. +//! +//! Eval targets are referenced by stable `external_id`, so a query set stays +//! valid across the append-only supersession that changes internal `id`s. + +use std::collections::HashMap; + +use lance::{Error as LanceError, Result as LanceResult}; +use serde::{Deserialize, Serialize}; + +use crate::record::{LifecycleQueryOptions, RecordFilters}; +use crate::store::ContextStore; + +fn default_grade() -> f32 { + 1.0 +} + +/// A single relevance judgment: a relevant `external_id` and its grade. +/// +/// `grade` defaults to `1.0` (binary relevance). Graded relevance (e.g. `2.0` +/// for "highly relevant") is used by nDCG; any grade `> 0` counts as relevant +/// for recall/precision/MRR/hit-rate. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct RelevanceLabel { + pub external_id: String, + #[serde(default = "default_grade")] + pub grade: f32, +} + +/// One labeled query: a vector and/or text channel plus its relevant records. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct EvalQuery { + pub query_id: String, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub text: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub vector: Option>, + #[serde(default)] + pub relevant: Vec, +} + +impl EvalQuery { + fn relevance_map(&self) -> HashMap<&str, f32> { + self.relevant + .iter() + .map(|label| (label.external_id.as_str(), label.grade)) + .collect() + } +} + +/// A labeled query set, referenced by a stable `id` for reproducible reports. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct EvalQuerySet { + pub id: String, + pub queries: Vec, +} + +impl EvalQuerySet { + #[must_use] + pub fn new(id: impl Into, queries: Vec) -> Self { + Self { + id: id.into(), + queries, + } + } + + /// Parse a query set from JSONL, one [`EvalQuery`] object per line. Blank + /// lines are ignored. The set `id` is supplied by the caller (e.g. derived + /// from the file name) so the same labels can be reused across runs. + pub fn from_jsonl(id: impl Into, contents: &str) -> LanceResult { + let mut queries = Vec::new(); + for (index, line) in contents.lines().enumerate() { + let line = line.trim(); + if line.is_empty() { + continue; + } + let query: EvalQuery = serde_json::from_str(line).map_err(|err| { + LanceError::invalid_input(format!( + "invalid eval query on line {}: {err}", + index + 1 + )) + })?; + queries.push(query); + } + Ok(Self::new(id, queries)) + } + + /// Serialize the query set to JSONL (one query per line). + pub fn to_jsonl(&self) -> LanceResult { + let mut out = String::new(); + for query in &self.queries { + let line = serde_json::to_string(query) + .map_err(|err| LanceError::invalid_input(err.to_string()))?; + out.push_str(&line); + out.push('\n'); + } + Ok(out) + } +} + +/// Which retrieval API to evaluate. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)] +#[serde(rename_all = "lowercase")] +pub enum RetrievalMode { + /// Pure vector search ([`ContextStore::search_filtered_with_options`]). + #[default] + Vector, + /// Hybrid text + vector retrieval + /// ([`ContextStore::retrieve_filtered_with_options`]). + Hybrid, +} + +impl RetrievalMode { + #[must_use] + pub fn as_str(self) -> &'static str { + match self { + Self::Vector => "vector", + Self::Hybrid => "hybrid", + } + } +} + +/// Runtime configuration for an evaluation run. +#[derive(Clone)] +pub struct EvalConfig { + /// Rank cutoff `k` for all @k metrics and the retrieval limit. + pub k: usize, + pub mode: RetrievalMode, + pub filters: Option, + pub lifecycle: LifecycleQueryOptions, +} + +impl Default for EvalConfig { + fn default() -> Self { + Self { + k: 10, + mode: RetrievalMode::Vector, + filters: None, + lifecycle: LifecycleQueryOptions::default(), + } + } +} + +/// Aggregate or per-query retrieval-quality metrics, all in `0.0..=1.0`. +#[derive(Debug, Clone, Copy, Default, PartialEq, Serialize, Deserialize)] +pub struct MetricScores { + pub recall: f64, + pub precision: f64, + pub mrr: f64, + pub ndcg: f64, + pub hit_rate: f64, +} + +impl MetricScores { + /// Per-metric difference `self - baseline`, for A/B deltas. + #[must_use] + pub fn delta(&self, baseline: &MetricScores) -> MetricScores { + MetricScores { + recall: self.recall - baseline.recall, + precision: self.precision - baseline.precision, + mrr: self.mrr - baseline.mrr, + ndcg: self.ndcg - baseline.ndcg, + hit_rate: self.hit_rate - baseline.hit_rate, + } + } +} + +/// Metrics and retrieved ids for a single query. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct QueryEval { + pub query_id: String, + /// External ids retrieved in rank order (top-k). Records retrieved without + /// an `external_id` appear as an empty string and never match a label. + pub retrieved: Vec, + pub scores: MetricScores, +} + +/// A reproducible evaluation report: a manifest (query-set id, version, config) +/// plus aggregate and per-query scores. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct EvalReport { + pub query_set_id: String, + /// Dataset version the run was pinned to. + pub version: u64, + pub k: usize, + pub mode: String, + /// Distance metric the context is configured with (part of the manifest). + pub distance_metric: String, + pub num_queries: usize, + pub aggregate: MetricScores, + pub per_query: Vec, +} + +/// Result of comparing the same query set across two dataset versions. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct AbReport { + pub query_set_id: String, + pub baseline: EvalReport, + pub candidate: EvalReport, + /// `candidate.aggregate - baseline.aggregate`, per metric. + pub deltas: MetricScores, +} + +/// Compute @k metrics for one ranked result list against a relevance map. +fn compute_scores(retrieved: &[String], relevant: &HashMap<&str, f32>, k: usize) -> MetricScores { + let k = k.max(1); + let num_relevant = relevant.values().filter(|grade| **grade > 0.0).count(); + + let mut hits = 0usize; + let mut first_relevant_rank: Option = None; + let mut dcg = 0.0f64; + for (index, external_id) in retrieved.iter().take(k).enumerate() { + let grade = relevant.get(external_id.as_str()).copied().unwrap_or(0.0); + if grade > 0.0 { + hits += 1; + if first_relevant_rank.is_none() { + first_relevant_rank = Some(index + 1); + } + // rank = index + 1, discount = log2(rank + 1) = log2(index + 2) + dcg += f64::from(grade) / ((index + 2) as f64).log2(); + } + } + + // Ideal DCG: best achievable ordering of the graded labels, truncated at k. + let mut ideal_grades: Vec = relevant + .values() + .filter(|grade| **grade > 0.0) + .map(|grade| f64::from(*grade)) + .collect(); + ideal_grades.sort_by(|a, b| b.total_cmp(a)); + let idcg: f64 = ideal_grades + .iter() + .take(k) + .enumerate() + .map(|(index, grade)| grade / ((index + 2) as f64).log2()) + .sum(); + + MetricScores { + recall: if num_relevant > 0 { + hits as f64 / num_relevant as f64 + } else { + 0.0 + }, + precision: hits as f64 / k as f64, + mrr: first_relevant_rank.map_or(0.0, |rank| 1.0 / rank as f64), + ndcg: if idcg > 0.0 { dcg / idcg } else { 0.0 }, + hit_rate: if hits > 0 { 1.0 } else { 0.0 }, + } +} + +fn mean_scores(per_query: &[QueryEval]) -> MetricScores { + let n = per_query.len(); + if n == 0 { + return MetricScores::default(); + } + let mut agg = MetricScores::default(); + for query in per_query { + agg.recall += query.scores.recall; + agg.precision += query.scores.precision; + agg.mrr += query.scores.mrr; + agg.ndcg += query.scores.ndcg; + agg.hit_rate += query.scores.hit_rate; + } + let n = n as f64; + MetricScores { + recall: agg.recall / n, + precision: agg.precision / n, + mrr: agg.mrr / n, + ndcg: agg.ndcg / n, + hit_rate: agg.hit_rate / n, + } +} + +impl ContextStore { + /// Run a labeled query set against this context at its current version and + /// return a reproducible [`EvalReport`]. + /// + /// Each query is retrieved with `config.mode` (vector or hybrid) at the + /// `config.k` cutoff, with `config.filters` / `config.lifecycle` applied, + /// then scored against its `relevant` labels by `external_id`. + pub async fn evaluate( + &self, + query_set: &EvalQuerySet, + config: &EvalConfig, + ) -> LanceResult { + let mut per_query = Vec::with_capacity(query_set.queries.len()); + for query in &query_set.queries { + let retrieved = self.run_eval_query(query, config).await?; + let relevant = query.relevance_map(); + let scores = compute_scores(&retrieved, &relevant, config.k); + per_query.push(QueryEval { + query_id: query.query_id.clone(), + retrieved, + scores, + }); + } + + Ok(EvalReport { + query_set_id: query_set.id.clone(), + version: self.version(), + k: config.k, + mode: config.mode.as_str().to_string(), + distance_metric: self.distance_metric().as_str().to_string(), + num_queries: per_query.len(), + aggregate: mean_scores(&per_query), + per_query, + }) + } + + /// A/B the same query set across two dataset versions and report per-metric + /// deltas (`candidate - baseline`). The store is restored to its current + /// version before returning. + pub async fn evaluate_versions( + &mut self, + query_set: &EvalQuerySet, + config: &EvalConfig, + baseline_version: u64, + candidate_version: u64, + ) -> LanceResult { + let original_version = self.version(); + + self.checkout(baseline_version).await?; + let baseline = self.evaluate(query_set, config).await?; + self.checkout(candidate_version).await?; + let candidate = self.evaluate(query_set, config).await?; + self.checkout(original_version).await?; + + let deltas = candidate.aggregate.delta(&baseline.aggregate); + Ok(AbReport { + query_set_id: query_set.id.clone(), + baseline, + candidate, + deltas, + }) + } + + /// Retrieve the top-k `external_id`s for one query under `config`. + async fn run_eval_query( + &self, + query: &EvalQuery, + config: &EvalConfig, + ) -> LanceResult> { + let limit = Some(config.k); + let records = match config.mode { + RetrievalMode::Vector => { + let vector = query.vector.as_deref().ok_or_else(|| { + LanceError::invalid_input(format!( + "query '{}' has no vector for vector-mode eval", + query.query_id + )) + })?; + self.search_filtered_with_options( + vector, + limit, + config.filters.as_ref(), + config.lifecycle.clone(), + ) + .await? + .into_iter() + .map(|hit| hit.record) + .collect::>() + } + RetrievalMode::Hybrid => self + .retrieve_filtered_with_options( + query.text.as_deref(), + query.vector.as_deref(), + limit, + config.filters.as_ref(), + config.lifecycle.clone(), + ) + .await? + .into_iter() + .map(|hit| hit.record) + .collect::>(), + }; + + Ok(records + .into_iter() + .map(|record| record.external_id.unwrap_or_default()) + .collect()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::record::{ContextRecord, LIFECYCLE_ACTIVE}; + use crate::store::ContextStore; + use chrono::Utc; + use serde_json::json; + use tempfile::TempDir; + use uuid::Uuid; + + // ----- metric fixtures (pure) ------------------------------------------- + + fn scores(retrieved: &[&str], relevant: &[(&str, f32)], k: usize) -> MetricScores { + let retrieved: Vec = retrieved.iter().map(|s| s.to_string()).collect(); + let relevant: HashMap<&str, f32> = relevant.iter().copied().collect(); + compute_scores(&retrieved, &relevant, k) + } + + fn approx(actual: f64, expected: f64) { + assert!( + (actual - expected).abs() < 1e-4, + "expected {expected}, got {actual}" + ); + } + + #[test] + fn metrics_perfect_ranking() { + let s = scores(&["a", "b"], &[("a", 1.0), ("b", 1.0)], 2); + approx(s.recall, 1.0); + approx(s.precision, 1.0); + approx(s.mrr, 1.0); + approx(s.ndcg, 1.0); + approx(s.hit_rate, 1.0); + } + + #[test] + fn metrics_single_relevant_at_rank_two() { + // retrieved [x, a]; only `a` relevant; k=2. + let s = scores(&["x", "a"], &[("a", 1.0)], 2); + approx(s.recall, 1.0); // found the 1 relevant + approx(s.precision, 0.5); // 1 hit / k=2 + approx(s.mrr, 0.5); // first relevant at rank 2 + approx(s.hit_rate, 1.0); + // dcg = 1/log2(3); idcg = 1/log2(2) = 1.0 + approx(s.ndcg, 1.0 / 3.0_f64.log2()); + } + + #[test] + fn metrics_no_relevant_in_topk() { + let s = scores(&["x", "y"], &[("a", 1.0)], 2); + approx(s.recall, 0.0); + approx(s.precision, 0.0); + approx(s.mrr, 0.0); + approx(s.ndcg, 0.0); + approx(s.hit_rate, 0.0); + } + + #[test] + fn metrics_graded_ndcg() { + // retrieved [a(grade1), b(grade3)]; ideal order is [b, a]. + let s = scores(&["a", "b"], &[("a", 1.0), ("b", 3.0)], 2); + let dcg = 1.0 / 2.0_f64.log2() + 3.0 / 3.0_f64.log2(); + let idcg = 3.0 / 2.0_f64.log2() + 1.0 / 3.0_f64.log2(); + approx(s.ndcg, dcg / idcg); + approx(s.recall, 1.0); + } + + #[test] + fn metrics_precision_is_over_k() { + // Only one item retrieved but k=2 -> precision = 1/2. + let s = scores(&["a"], &[("a", 1.0)], 2); + approx(s.precision, 0.5); + approx(s.recall, 1.0); + approx(s.hit_rate, 1.0); + } + + #[test] + fn query_set_jsonl_round_trip() { + let jsonl = concat!( + "{\"query_id\":\"q1\",\"vector\":[1.0,0.0],\"relevant\":[{\"external_id\":\"a\"}]}\n", + "\n", + "{\"query_id\":\"q2\",\"text\":\"hi\",\"relevant\":[{\"external_id\":\"b\",\"grade\":2.0}]}\n", + ); + let set = EvalQuerySet::from_jsonl("set-1", jsonl).unwrap(); + assert_eq!(set.queries.len(), 2); + assert_eq!(set.queries[0].query_id, "q1"); + assert_eq!(set.queries[1].relevant[0].grade, 2.0); + // default grade applies when omitted + assert_eq!(set.queries[0].relevant[0].grade, 1.0); + + let reparsed = EvalQuerySet::from_jsonl("set-1", &set.to_jsonl().unwrap()).unwrap(); + assert_eq!(reparsed.queries.len(), 2); + assert_eq!(reparsed.queries[1].relevant[0].external_id, "b"); + } + + // ----- runner integration ----------------------------------------------- + + fn embedding(store: &ContextStore, lead: &[f32]) -> Vec { + let dim = store.embedding_dim() as usize; + let mut v = vec![0.0f32; dim]; + for (i, x) in lead.iter().enumerate() { + v[i] = *x; + } + v + } + + fn record(external_id: &str, text: &str, embedding: Vec) -> ContextRecord { + ContextRecord { + id: Uuid::new_v4().to_string(), + external_id: Some(external_id.to_string()), + run_id: "run".to_string(), + bot_id: None, + session_id: None, + tenant: None, + source: None, + created_at: Utc::now(), + role: "user".to_string(), + state_metadata: None, + metadata: None, + relationships: Vec::new(), + expires_at: None, + retention_policy: None, + lifecycle_status: LIFECYCLE_ACTIVE.to_string(), + retired_at: None, + retired_reason: None, + supersedes_id: None, + superseded_by_id: None, + content_type: "text/plain".to_string(), + text_payload: Some(text.to_string()), + binary_payload: None, + embedding: Some(embedding), + } + } + + #[test] + fn evaluate_vector_mode_scores_query_set() { + let dir = TempDir::new().unwrap(); + let uri = dir.path().to_string_lossy().to_string(); + let runtime = tokio::runtime::Runtime::new().unwrap(); + runtime.block_on(async { + let mut store = ContextStore::open(&uri).await.unwrap(); + let a = embedding(&store, &[1.0]); + let b = embedding(&store, &[0.5]); + let c = embedding(&store, &[0.0, 1.0]); + store + .add(&[ + record("doc-a", "alpha", a.clone()), + record("doc-b", "beta", b), + record("doc-c", "gamma", c), + ]) + .await + .unwrap(); + + // query closest to doc-a (then doc-b), only doc-a relevant. + let query_set = EvalQuerySet::new( + "qs", + vec![EvalQuery { + query_id: "q1".to_string(), + text: None, + vector: Some(a), + relevant: vec![RelevanceLabel { + external_id: "doc-a".to_string(), + grade: 1.0, + }], + }], + ); + let config = EvalConfig { + k: 2, + mode: RetrievalMode::Vector, + ..Default::default() + }; + let report = store.evaluate(&query_set, &config).await.unwrap(); + + assert_eq!(report.num_queries, 1); + assert_eq!(report.mode, "vector"); + assert_eq!(report.k, 2); + assert_eq!(report.per_query[0].retrieved.first().unwrap(), "doc-a"); + approx(report.aggregate.recall, 1.0); + approx(report.aggregate.precision, 0.5); + approx(report.aggregate.mrr, 1.0); + approx(report.aggregate.hit_rate, 1.0); + }); + } + + #[test] + fn evaluate_respects_lifecycle_visibility() { + let dir = TempDir::new().unwrap(); + let uri = dir.path().to_string_lossy().to_string(); + let runtime = tokio::runtime::Runtime::new().unwrap(); + runtime.block_on(async { + let mut store = ContextStore::open(&uri).await.unwrap(); + let q = embedding(&store, &[1.0]); + // the only relevant doc is retired -> hidden by default. + let mut retired = record("doc-a", "alpha", q.clone()); + retired.retired_at = Some(Utc::now()); + store.add(&[retired]).await.unwrap(); + + let query_set = EvalQuerySet::new( + "qs", + vec![EvalQuery { + query_id: "q1".to_string(), + text: None, + vector: Some(q), + relevant: vec![RelevanceLabel { + external_id: "doc-a".to_string(), + grade: 1.0, + }], + }], + ); + + let default_cfg = EvalConfig { + k: 5, + mode: RetrievalMode::Vector, + ..Default::default() + }; + let hidden = store.evaluate(&query_set, &default_cfg).await.unwrap(); + approx(hidden.aggregate.recall, 0.0); // retired doc excluded + + let include_retired = EvalConfig { + k: 5, + mode: RetrievalMode::Vector, + lifecycle: LifecycleQueryOptions::new(true, true), + ..Default::default() + }; + let visible = store.evaluate(&query_set, &include_retired).await.unwrap(); + approx(visible.aggregate.recall, 1.0); // surfaced with include_retired + }); + } + + #[test] + fn evaluate_respects_filters() { + let dir = TempDir::new().unwrap(); + let uri = dir.path().to_string_lossy().to_string(); + let runtime = tokio::runtime::Runtime::new().unwrap(); + runtime.block_on(async { + let mut store = ContextStore::open(&uri).await.unwrap(); + let shared = embedding(&store, &[1.0]); + let mut a = record("doc-a", "alpha", shared.clone()); + a.tenant = Some("x".to_string()); + let mut b = record("doc-b", "beta", shared.clone()); + b.tenant = Some("y".to_string()); + store.add(&[a, b]).await.unwrap(); + + // doc-b is relevant but filtered out by tenant=x. + let query_set = EvalQuerySet::new( + "qs", + vec![EvalQuery { + query_id: "q1".to_string(), + text: None, + vector: Some(shared), + relevant: vec![RelevanceLabel { + external_id: "doc-b".to_string(), + grade: 1.0, + }], + }], + ); + let config = EvalConfig { + k: 5, + mode: RetrievalMode::Vector, + filters: Some(RecordFilters::from_json_value(json!({"tenant": "x"})).unwrap()), + ..Default::default() + }; + let report = store.evaluate(&query_set, &config).await.unwrap(); + approx(report.aggregate.recall, 0.0); // doc-b excluded by filter + }); + } + + #[test] + fn evaluate_hybrid_mode_finds_relevant() { + let dir = TempDir::new().unwrap(); + let uri = dir.path().to_string_lossy().to_string(); + let runtime = tokio::runtime::Runtime::new().unwrap(); + runtime.block_on(async { + let mut store = ContextStore::open(&uri).await.unwrap(); + let a = embedding(&store, &[1.0]); + let b = embedding(&store, &[0.0, 1.0]); + store + .add(&[ + record("doc-a", "alpha unique", a.clone()), + record("doc-b", "beta other", b), + ]) + .await + .unwrap(); + + let query_set = EvalQuerySet::new( + "qs", + vec![EvalQuery { + query_id: "q1".to_string(), + text: Some("alpha".to_string()), + vector: Some(a), + relevant: vec![RelevanceLabel { + external_id: "doc-a".to_string(), + grade: 1.0, + }], + }], + ); + let config = EvalConfig { + k: 2, + mode: RetrievalMode::Hybrid, + ..Default::default() + }; + let report = store.evaluate(&query_set, &config).await.unwrap(); + approx(report.aggregate.hit_rate, 1.0); + }); + } + + #[test] + fn config_ab_delta_detects_k_sensitivity() { + let dir = TempDir::new().unwrap(); + let uri = dir.path().to_string_lossy().to_string(); + let runtime = tokio::runtime::Runtime::new().unwrap(); + runtime.block_on(async { + let mut store = ContextStore::open(&uri).await.unwrap(); + let a = embedding(&store, &[1.0]); + let b = embedding(&store, &[0.5]); + store + .add(&[ + record("doc-a", "alpha", a.clone()), + record("doc-b", "beta", b), + ]) + .await + .unwrap(); + + // relevant doc-b ranks second behind doc-a. + let query_set = EvalQuerySet::new( + "qs", + vec![EvalQuery { + query_id: "q1".to_string(), + text: None, + vector: Some(a), + relevant: vec![RelevanceLabel { + external_id: "doc-b".to_string(), + grade: 1.0, + }], + }], + ); + let k1 = EvalConfig { + k: 1, + mode: RetrievalMode::Vector, + ..Default::default() + }; + let k2 = EvalConfig { + k: 2, + mode: RetrievalMode::Vector, + ..Default::default() + }; + let at_1 = store.evaluate(&query_set, &k1).await.unwrap(); + let at_2 = store.evaluate(&query_set, &k2).await.unwrap(); + approx(at_1.aggregate.recall, 0.0); // doc-b not in top-1 + approx(at_2.aggregate.recall, 1.0); // doc-b in top-2 + let delta = at_2.aggregate.delta(&at_1.aggregate); + approx(delta.recall, 1.0); + }); + } + + #[test] + fn evaluate_versions_same_version_is_zero_delta_and_restores() { + let dir = TempDir::new().unwrap(); + let uri = dir.path().to_string_lossy().to_string(); + let runtime = tokio::runtime::Runtime::new().unwrap(); + runtime.block_on(async { + let mut store = ContextStore::open(&uri).await.unwrap(); + let a = embedding(&store, &[1.0]); + store + .add(&[record("doc-a", "alpha", a.clone())]) + .await + .unwrap(); + let version = store.version(); + + let query_set = EvalQuerySet::new( + "qs", + vec![EvalQuery { + query_id: "q1".to_string(), + text: None, + vector: Some(a), + relevant: vec![RelevanceLabel { + external_id: "doc-a".to_string(), + grade: 1.0, + }], + }], + ); + let config = EvalConfig { + k: 1, + mode: RetrievalMode::Vector, + ..Default::default() + }; + let ab = store + .evaluate_versions(&query_set, &config, version, version) + .await + .unwrap(); + + approx(ab.deltas.recall, 0.0); + approx(ab.deltas.ndcg, 0.0); + assert_eq!(ab.baseline.version, version); + assert_eq!(ab.candidate.version, version); + assert_eq!( + store.version(), + version, + "store restored to original version" + ); + }); + } +} diff --git a/crates/lance-context-core/src/lib.rs b/crates/lance-context-core/src/lib.rs index ae79d3c..2875f37 100644 --- a/crates/lance-context-core/src/lib.rs +++ b/crates/lance-context-core/src/lib.rs @@ -3,6 +3,7 @@ mod api_impl; mod context; +mod eval; mod export; mod namespace; mod record; @@ -10,6 +11,10 @@ pub mod serde; mod store; pub use context::{Context, ContextEntry, Snapshot}; +pub use eval::{ + AbReport, EvalConfig, EvalQuery, EvalQuerySet, EvalReport, MetricScores, QueryEval, + RelevanceLabel, RetrievalMode, +}; pub use export::{ ExportConfig, ExportCounts, ExportManifest, ExportTask, GroupBy, Message, PreferenceExample, PreferenceForm, Provenance, RankedCandidate, RolloutExample, RolloutResponse, SftExample, diff --git a/crates/lance-context-core/src/store.rs b/crates/lance-context-core/src/store.rs index 3ce9f13..2f64519 100644 --- a/crates/lance-context-core/src/store.rs +++ b/crates/lance-context-core/src/store.rs @@ -377,6 +377,12 @@ impl ContextStore { self.dataset.uri() } + /// Distance metric this context ranks vector-search results with. + #[must_use] + pub fn distance_metric(&self) -> DistanceMetric { + self.distance_metric + } + /// Append context records to the store and return the new dataset version. pub async fn add(&mut self, entries: &[ContextRecord]) -> LanceResult { if entries.is_empty() { diff --git a/python/python/lance_context/api.py b/python/python/lance_context/api.py index f3718b9..9b188a7 100644 --- a/python/python/lance_context/api.py +++ b/python/python/lance_context/api.py @@ -1315,6 +1315,77 @@ def export_training( ) return json.loads(manifest_json) + def evaluate( + self, + queries: Iterable[Mapping[str, Any]], + *, + query_set_id: str = "eval", + k: int = 10, + mode: str = "vector", + filters: dict[str, Any] | None = None, + include_expired: bool = False, + include_retired: bool = False, + ) -> dict[str, Any]: + """Evaluate retrieval quality against a labeled query set. + + Each query is a mapping with ``query_id``, an optional ``text`` and/or + ``vector`` channel, and ``relevant`` — a list of ``{"external_id": ..., + "grade": }`` labels. ``mode`` selects ``"vector"`` + (search) or ``"hybrid"`` (retrieve). + + Returns a report dict with ``aggregate`` metrics (``recall``, + ``precision``, ``mrr``, ``ndcg``, ``hit_rate``), a ``per_query`` + breakdown, and a manifest (``query_set_id``, ``version``, ``k``, + ``mode``, ``distance_metric``) for reproducibility. + """ + query_set = json.dumps( + {"id": query_set_id, "queries": [dict(query) for query in queries]} + ) + report_json = self._inner.evaluate( + query_set, + k, + mode, + _json_dumps(filters, "filters"), + include_expired, + include_retired, + ) + return json.loads(report_json) + + def evaluate_versions( + self, + queries: Iterable[Mapping[str, Any]], + baseline_version: int, + candidate_version: int, + *, + query_set_id: str = "eval", + k: int = 10, + mode: str = "vector", + filters: dict[str, Any] | None = None, + include_expired: bool = False, + include_retired: bool = False, + ) -> dict[str, Any]: + """A/B the same query set across two dataset versions. + + Runs :meth:`evaluate` at ``baseline_version`` and ``candidate_version`` + (via time-travel checkout) and returns ``{"baseline", "candidate", + "deltas"}`` where ``deltas`` is ``candidate - baseline`` per metric. The + context is restored to its current version before returning. + """ + query_set = json.dumps( + {"id": query_set_id, "queries": [dict(query) for query in queries]} + ) + report_json = self._inner.evaluate_versions( + query_set, + int(baseline_version), + int(candidate_version), + k, + mode, + _json_dumps(filters, "filters"), + include_expired, + include_retired, + ) + return json.loads(report_json) + def list( self, limit: int | None = None, diff --git a/python/src/lib.rs b/python/src/lib.rs index 6bedae4..484c24b 100644 --- a/python/src/lib.rs +++ b/python/src/lib.rs @@ -21,9 +21,10 @@ use lance_context_core::serde::CONTENT_TYPE_TEXT; use lance_context_core::{ CompactionConfig, CompactionMetrics, CompactionStats, Context as RustContext, ContextNamespace as RustContextNamespace, ContextRecord, ContextStore, ContextStoreOptions, - DistanceMetric, ExportConfig, ExportTask, GroupBy, IdIndexType, LifecycleQueryOptions, - PartitionInfo, PartitionSelector, PartitionSpec, PreferenceForm, RecordFilters, RecordPatch, - Relationship, RetrieveResult, SearchResult, SplitConfig, StateMetadata, LIFECYCLE_ACTIVE, + DistanceMetric, EvalConfig, EvalQuerySet, ExportConfig, ExportTask, GroupBy, IdIndexType, + LifecycleQueryOptions, PartitionInfo, PartitionSelector, PartitionSpec, PreferenceForm, + RecordFilters, RecordPatch, Relationship, RetrievalMode, RetrieveResult, SearchResult, + SplitConfig, StateMetadata, LIFECYCLE_ACTIVE, }; const DEFAULT_BINARY_CONTENT_TYPE: &str = "application/octet-stream"; @@ -317,6 +318,30 @@ fn parse_group_by(group_by: &str) -> PyResult { }) } +fn eval_config( + k: usize, + mode: &str, + filters_json: Option, + include_expired: bool, + include_retired: bool, +) -> PyResult { + let mode = match mode { + "vector" => RetrievalMode::Vector, + "hybrid" => RetrievalMode::Hybrid, + other => { + return Err(PyRuntimeError::new_err(format!( + "invalid eval mode '{other}'; use 'vector' or 'hybrid'" + ))); + } + }; + Ok(EvalConfig { + k, + mode, + filters: filters_from_json(filters_json)?, + lifecycle: LifecycleQueryOptions::new(include_expired, include_retired), + }) +} + fn selector_from_dict(dict: &Bound<'_, PyDict>) -> PyResult { let mut selector = BTreeMap::new(); for (key, value) in dict.iter() { @@ -839,6 +864,58 @@ impl Context { serde_json::to_string(&manifest).map_err(to_py_err) } + #[allow(clippy::too_many_arguments)] + #[pyo3(signature = (query_set_json, k = 10, mode = "vector", filters_json = None, include_expired = false, include_retired = false))] + fn evaluate( + &self, + py: Python<'_>, + query_set_json: &str, + k: usize, + mode: &str, + filters_json: Option, + include_expired: bool, + include_retired: bool, + ) -> PyResult { + let query_set: EvalQuerySet = serde_json::from_str(query_set_json).map_err(to_py_err)?; + let config = eval_config(k, mode, filters_json, include_expired, include_retired)?; + let report = py + .allow_threads(|| { + self.runtime + .block_on(self.store.evaluate(&query_set, &config)) + }) + .map_err(to_py_err)?; + serde_json::to_string(&report).map_err(to_py_err) + } + + #[allow(clippy::too_many_arguments)] + #[pyo3(signature = (query_set_json, baseline_version, candidate_version, k = 10, mode = "vector", filters_json = None, include_expired = false, include_retired = false))] + fn evaluate_versions( + &mut self, + py: Python<'_>, + query_set_json: &str, + baseline_version: u64, + candidate_version: u64, + k: usize, + mode: &str, + filters_json: Option, + include_expired: bool, + include_retired: bool, + ) -> PyResult { + let query_set: EvalQuerySet = serde_json::from_str(query_set_json).map_err(to_py_err)?; + let config = eval_config(k, mode, filters_json, include_expired, include_retired)?; + let report = py + .allow_threads(|| { + self.runtime.block_on(self.store.evaluate_versions( + &query_set, + &config, + baseline_version, + candidate_version, + )) + }) + .map_err(to_py_err)?; + serde_json::to_string(&report).map_err(to_py_err) + } + #[pyo3(signature = (limit = None, offset = None, filters_json = None, include_expired = false, include_retired = false))] fn list( &self, diff --git a/python/tests/test_eval.py b/python/tests/test_eval.py new file mode 100644 index 0000000..2f75064 --- /dev/null +++ b/python/tests/test_eval.py @@ -0,0 +1,137 @@ +"""Tests for the retrieval-quality evaluation harness.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import pytest + +if TYPE_CHECKING: + from pathlib import Path + +from lance_context.api import Context + +DIM = 8 + + +def _emb(*lead: float) -> list[float]: + vec = [0.0] * DIM + for i, value in enumerate(lead): + vec[i] = value + return vec + + +def _make_context(tmp_path: Path) -> Context: + ctx = Context.create(str(tmp_path / "context.lance"), embedding_dim=DIM) + ctx.add_many( + [ + { + "role": "user", + "content": "alpha", + "external_id": "doc-a", + "embedding": _emb(1.0), + }, + { + "role": "user", + "content": "beta", + "external_id": "doc-b", + "embedding": _emb(0.5), + }, + { + "role": "user", + "content": "gamma", + "external_id": "doc-c", + "embedding": _emb(0.0, 1.0), + }, + ] + ) + return ctx + + +def test_evaluate_vector_mode(tmp_path: Path) -> None: + ctx = _make_context(tmp_path) + report = ctx.evaluate( + [ + { + "query_id": "q1", + "vector": _emb(1.0), + "relevant": [{"external_id": "doc-a"}], + } + ], + k=2, + mode="vector", + ) + + assert report["mode"] == "vector" + assert report["k"] == 2 + assert report["num_queries"] == 1 + assert report["per_query"][0]["retrieved"][0] == "doc-a" + agg = report["aggregate"] + assert agg["recall"] == pytest.approx(1.0) + assert agg["precision"] == pytest.approx(0.5) + assert agg["mrr"] == pytest.approx(1.0) + assert agg["hit_rate"] == pytest.approx(1.0) + + +def test_evaluate_graded_ndcg(tmp_path: Path) -> None: + ctx = _make_context(tmp_path) + report = ctx.evaluate( + [ + { + "query_id": "q1", + "vector": _emb(1.0), + "relevant": [ + {"external_id": "doc-a", "grade": 1.0}, + {"external_id": "doc-b", "grade": 3.0}, + ], + } + ], + k=2, + mode="vector", + ) + # retrieved order is [doc-a(grade1), doc-b(grade3)]; ideal is [b, a]. + assert 0.0 < report["aggregate"]["ndcg"] < 1.0 + + +def test_evaluate_config_ab_delta(tmp_path: Path) -> None: + ctx = _make_context(tmp_path) + queries = [ + { + "query_id": "q1", + "vector": _emb(1.0), + "relevant": [{"external_id": "doc-b"}], # ranks 2nd + } + ] + at_1 = ctx.evaluate(queries, k=1) + at_2 = ctx.evaluate(queries, k=2) + assert at_1["aggregate"]["recall"] == pytest.approx(0.0) + assert at_2["aggregate"]["recall"] == pytest.approx(1.0) + + +def test_evaluate_versions_same_version_zero_delta(tmp_path: Path) -> None: + ctx = _make_context(tmp_path) + version = ctx.version() + ab = ctx.evaluate_versions( + [ + { + "query_id": "q1", + "vector": _emb(1.0), + "relevant": [{"external_id": "doc-a"}], + } + ], + version, + version, + k=1, + ) + assert ab["deltas"]["recall"] == pytest.approx(0.0) + assert ab["baseline"]["version"] == version + assert ctx.version() == version + + +def test_evaluate_invalid_mode_raises(tmp_path: Path) -> None: + ctx = _make_context(tmp_path) + with pytest.raises(RuntimeError, match="invalid eval mode"): + ctx.evaluate( + [{"query_id": "q1", "vector": _emb(1.0), "relevant": []}], + mode="bogus", + )