« prev ^ index » next coverage.py v7.15.4, created at 2026-08-25 07:19 +0800
1"""RAG benchmarking harness for comparing pipeline configurations."""
3from __future__ import annotations
5import asyncio
6from dataclasses import dataclass, field
7from datetime import UTC, datetime
8import time
9from typing import Any
11from lexigram.ai.rag.evaluation.evaluator import RAGEvaluator
12from lexigram.ai.rag.evaluation.retrieval import (
13 RetrievalPrecisionEvaluator,
14 RetrievalRecallEvaluator,
15)
16from lexigram.ai.rag.evaluation.types import MetricType, RAGEvaluationReport
17from lexigram.contracts.ai.rag import RAGContext, RAGPipelineProtocol
18from lexigram.logging import (
19 get_logger,
20)
22logger = get_logger(__name__)
25@dataclass
26class EvalExample:
27 """A single benchmark example with ground truth.
29 Attributes:
30 query: The query to evaluate.
31 relevant_doc_ids: Ground-truth relevant document identifiers.
32 reference_answer: Optional reference answer for answer-quality metrics.
33 metadata: Arbitrary metadata attached to this example.
34 """
36 query: str
37 relevant_doc_ids: list[str]
38 reference_answer: str | None = None
39 metadata: dict[str, Any] = field(default_factory=dict)
42@dataclass
43class PipelineResult:
44 """Result of running a single pipeline on a single example.
46 Attributes:
47 pipeline_name: Identifier of the pipeline that produced this result.
48 query: The query that was evaluated.
49 retrieved_doc_ids: Identifiers of documents returned by the pipeline.
50 answer: Generated answer.
51 latency_ms: End-to-end latency in milliseconds.
52 error: Error message if the pipeline failed.
53 """
55 pipeline_name: str
56 query: str
57 retrieved_doc_ids: list[str]
58 answer: str
59 latency_ms: float
60 error: str | None = None
63@dataclass
64class BenchmarkReport:
65 """Comparative benchmark report across multiple RAG pipelines.
67 Attributes:
68 pipeline_names: Names of all benchmarked pipelines.
69 total_examples: Number of evaluation examples used.
70 per_pipeline: Mapping of pipeline name → metric name → mean score.
71 best_pipeline: Name of the pipeline with the highest overall score.
72 generated_at: When this report was produced.
73 raw_results: Per-pipeline per-example evaluation reports,
74 keyed by pipeline name.
75 """
77 pipeline_names: list[str]
78 total_examples: int
79 per_pipeline: dict[str, dict[str, float]]
80 best_pipeline: str | None = None
81 generated_at: datetime = field(default_factory=lambda: datetime.now(UTC))
82 raw_results: dict[str, list[RAGEvaluationReport]] = field(default_factory=dict)
84 def summary_table(self) -> list[dict[str, Any]]:
85 """Return per-pipeline metric scores as a list of row dicts.
87 Each row has a ``pipeline`` key plus one key per metric.
89 Returns:
90 List of dicts suitable for tabular display.
91 """
92 rows = []
93 for name in self.pipeline_names:
94 row: dict[str, Any] = {"pipeline": name}
95 row.update(self.per_pipeline.get(name, {}))
96 rows.append(row)
97 return rows
100class RAGBenchmark:
101 """Compare RAG pipeline configurations on a labelled dataset.
103 Runs every pipeline against every example in the dataset, evaluates
104 retrieval and (optionally) answer quality, then aggregates scores
105 per pipeline.
107 Example:
108 >>> dataset = [EvalExample(query="...", relevant_doc_ids=["doc1"])]
109 >>> report = await RAGBenchmark().run({"pipe_a": pipeline_a}, dataset)
110 >>> print(report.best_pipeline)
111 """
113 def __init__(
114 self,
115 evaluator: RAGEvaluator | None = None,
116 *,
117 max_concurrency: int = 5,
118 ) -> None:
119 """Initialise the benchmark harness.
121 Args:
122 evaluator: Optional pre-configured :class:`RAGEvaluator`. When
123 ``None`` a default evaluator with retrieval-precision and
124 retrieval-recall is used.
125 max_concurrency: Maximum number of pipeline-example pairs that are
126 evaluated simultaneously.
127 """
128 self._evaluator = evaluator or RAGEvaluator(
129 evaluators=[
130 RetrievalPrecisionEvaluator(),
131 RetrievalRecallEvaluator(),
132 ]
133 )
134 self._max_concurrency = max_concurrency
136 async def run(
137 self,
138 pipelines: dict[str, RAGPipelineProtocol],
139 dataset: list[EvalExample],
140 metrics: list[MetricType] | None = None,
141 ) -> BenchmarkReport:
142 """Run the benchmark and return a comparative report.
144 For each (pipeline, example) pair the pipeline is executed, latency
145 is recorded, retrieval metrics are computed, and (when a
146 reference answer is provided) answer-quality metrics are computed too.
148 Args:
149 pipelines: Mapping of display-name → pipeline instance.
150 dataset: List of evaluation examples with ground-truth labels.
151 metrics: Subset of :class:`MetricType` to include in the report.
152 Defaults to all metrics produced by the evaluator.
154 Returns:
155 :class:`BenchmarkReport` with per-pipeline mean scores.
156 """
157 if not pipelines:
158 return BenchmarkReport(
159 pipeline_names=[],
160 total_examples=len(dataset),
161 per_pipeline={},
162 )
164 semaphore = asyncio.Semaphore(self._max_concurrency)
165 pipeline_names = list(pipelines.keys())
167 # raw_eval_results[pipeline_name] = list of RAGEvaluationReport
168 raw_eval_results: dict[str, list[RAGEvaluationReport]] = {
169 name: [] for name in pipeline_names
170 }
172 # Build all (pipeline_name, pipeline, example) tasks
173 tasks = [
174 (name, pipelines[name], example)
175 for name in pipeline_names
176 for example in dataset
177 ]
179 async def _run_one(
180 pipeline_name: str,
181 pipeline: RAGPipelineProtocol,
182 example: EvalExample,
183 ) -> tuple[str, RAGEvaluationReport | None]:
184 async with semaphore:
185 return await self._evaluate_one(pipeline_name, pipeline, example)
187 coros = [_run_one(name, pipeline, example) for name, pipeline, example in tasks]
188 results = await asyncio.gather(*coros, return_exceptions=False)
190 for pipeline_name, eval_report in results:
191 if eval_report is not None:
192 raw_eval_results[pipeline_name].append(eval_report)
194 # Aggregate per pipeline
195 per_pipeline = self._aggregate(raw_eval_results, metrics)
197 # Determine best pipeline by overall score
198 best = self._best_pipeline(per_pipeline)
200 return BenchmarkReport(
201 pipeline_names=pipeline_names,
202 total_examples=len(dataset),
203 per_pipeline=per_pipeline,
204 best_pipeline=best,
205 raw_results=raw_eval_results,
206 )
208 async def _evaluate_one(
209 self,
210 pipeline_name: str,
211 pipeline: RAGPipelineProtocol,
212 example: EvalExample,
213 ) -> tuple[str, RAGEvaluationReport | None]:
214 """Execute the pipeline and evaluate a single example.
216 Args:
217 pipeline_name: Display name of the pipeline.
218 pipeline: Pipeline instance to execute.
219 example: Evaluation example.
221 Returns:
222 Tuple of (pipeline_name, evaluation_report_or_None).
223 """
224 t0 = time.monotonic()
225 context = RAGContext(query=example.query)
227 try:
228 pipeline_result = await pipeline.execute(context)
229 except (OSError, RuntimeError, ValueError, TypeError) as exc:
230 logger.warning(
231 "benchmark: pipeline error",
232 pipeline=pipeline_name,
233 query=example.query,
234 error=str(exc),
235 )
236 return pipeline_name, None
238 latency_ms = (time.monotonic() - t0) * 1000.0
240 if pipeline_result.is_err():
241 logger.warning(
242 "benchmark: pipeline returned error",
243 pipeline=pipeline_name,
244 query=example.query,
245 error=str(pipeline_result.unwrap_err()),
246 )
247 return pipeline_name, None
249 rag_response = pipeline_result.unwrap()
251 # Extract retrieved doc IDs from sources
252 retrieved_doc_ids: list[str] = []
253 for source in rag_response.sources:
254 doc_id = getattr(source, "id", None) or getattr(
255 getattr(source, "document", None), "id", None
256 )
257 if doc_id is not None:
258 retrieved_doc_ids.append(str(doc_id))
260 eval_report = await self._evaluator.evaluate(
261 query=example.query,
262 retrieved_docs=retrieved_doc_ids,
263 generated_answer=rag_response.answer,
264 reference_answer=example.reference_answer,
265 relevant_doc_ids=example.relevant_doc_ids,
266 metadata={
267 "pipeline": pipeline_name,
268 "latency_ms": latency_ms,
269 **example.metadata,
270 },
271 )
273 # Inject latency as a metric result if not already present
274 if eval_report.get_metric(MetricType.LATENCY) is None:
275 from lexigram.ai.rag.evaluation.types import EvaluationResult
277 eval_report.results.append(
278 EvaluationResult(
279 metric_type=MetricType.LATENCY,
280 score=latency_ms,
281 details={"latency_ms": latency_ms},
282 )
283 )
285 return pipeline_name, eval_report
287 def _aggregate(
288 self,
289 raw: dict[str, list[RAGEvaluationReport]],
290 metrics: list[MetricType] | None,
291 ) -> dict[str, dict[str, float]]:
292 """Compute mean score per metric per pipeline.
294 Args:
295 raw: Per-pipeline list of evaluation reports.
296 metrics: Optional metric subset filter.
298 Returns:
299 Mapping of pipeline_name → metric_name → mean_score.
300 """
301 per_pipeline: dict[str, dict[str, float]] = {}
303 for pipeline_name, reports in raw.items():
304 if not reports:
305 per_pipeline[pipeline_name] = {}
306 continue
308 # Accumulate scores per metric
309 sums: dict[str, float] = {}
310 counts: dict[str, int] = {}
312 for report in reports:
313 for result in report.results:
314 if metrics and result.metric_type not in metrics:
315 continue
316 key = result.metric_type.value
317 sums[key] = sums.get(key, 0.0) + result.score
318 counts[key] = counts.get(key, 0) + 1
320 per_pipeline[pipeline_name] = {
321 key: sums[key] / counts[key] for key in sums if counts[key] > 0
322 }
324 return per_pipeline
326 def _best_pipeline(
327 self,
328 per_pipeline: dict[str, dict[str, float]],
329 ) -> str | None:
330 """Identify the pipeline with the highest mean overall score.
332 Latency and cost metrics are excluded from the ranking since lower
333 values are better for those and they are on a different scale.
335 Args:
336 per_pipeline: Aggregated per-pipeline metric scores.
338 Returns:
339 Name of the best pipeline, or ``None`` if no data is available.
340 """
341 _excluded = {MetricType.LATENCY.value, MetricType.COST.value}
343 best_name: str | None = None
344 best_score = -1.0
346 for pipeline_name, metrics in per_pipeline.items():
347 quality_scores = [
348 score
349 for key, score in metrics.items()
350 if key not in _excluded and not key.startswith(MetricType.LATENCY.value)
351 ]
352 if not quality_scores:
353 continue
354 mean = sum(quality_scores) / len(quality_scores)
355 if mean > best_score:
356 best_score = mean
357 best_name = pipeline_name
359 return best_name