From 5b3ad9bae7827f40186e47ece8060c2682c97eaf Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Thu, 12 Mar 2026 12:32:51 +0200 Subject: [PATCH] Refactor QAPromptAdapter to accept QA agent directly, rewrite mock-heavy tests --- evaluations/README.md | 51 ++++++-- evaluations/evaluations/optimization.py | 75 ++++++----- evaluations/tests/test_optimization.py | 157 ++++++++++-------------- 3 files changed, 149 insertions(+), 134 deletions(-) diff --git a/evaluations/README.md b/evaluations/README.md index 0e5a5861..37c0941c 100644 --- a/evaluations/README.md +++ b/evaluations/README.md @@ -6,31 +6,68 @@ This package is not published to PyPI and is only used for development and testi ## Overview -Contains evaluation scripts for benchmarking RAG performance using datasets like: +Contains evaluation scripts for benchmarking RAG retrieval and QA performance, plus GEPA-based prompt optimization. Available datasets: + - RepliQA - WiX +- HotpotQA +- OpenRAG Bench ## Usage After installing the package, you can run evaluations using the `evaluations` command: ```bash -# Run evaluations with default settings -evaluations repliqa +# Run retrieval + QA benchmarks +evaluations run repliqa +evaluations run wix # Use a custom config file -evaluations repliqa --config /path/to/haiku.rag.yaml +evaluations run repliqa --config /path/to/haiku.rag.yaml # Override the database path -evaluations repliqa --db /path/to/custom.lancedb +evaluations run repliqa --db /path/to/custom.lancedb # Skip database population and run only benchmarks -evaluations repliqa --skip-db +evaluations run repliqa --skip-db + +# Skip specific benchmarks +evaluations run repliqa --skip-retrieval +evaluations run repliqa --skip-qa # Limit the number of test cases -evaluations repliqa --limit 100 +evaluations run repliqa --limit 100 ``` +### Pre-built Databases + +Download pre-built evaluation databases from HuggingFace: + +```bash +evaluations download repliqa +evaluations download all +evaluations download repliqa --force +``` + +Upload databases (maintainer only): + +```bash +evaluations upload repliqa +evaluations upload all +``` + +### Prompt Optimization + +Optimize QA system prompts using GEPA (Generalized Evolutionary Prompt Algorithm): + +```bash +evaluations optimize wix +evaluations optimize repliqa --limit 40 --iterations 30 +evaluations optimize wix --output optimized_prompt.txt +``` + +See [Tuning docs](https://ggozad.github.io/haiku.rag/tuning/#prompt-optimization-gepa) for details on applying results. + ## Database Storage By default, evaluation databases are stored in the haiku.rag data directory: diff --git a/evaluations/evaluations/optimization.py b/evaluations/evaluations/optimization.py index f0851d7b..20a4cf4c 100644 --- a/evaluations/evaluations/optimization.py +++ b/evaluations/evaluations/optimization.py @@ -13,7 +13,7 @@ from gepa.core.adapter import EvaluationBatch from evaluations.benchmark import JUDGE_MODEL_CONFIG from evaluations.config import DatasetSpec -from haiku.rag.agents.qa import get_qa_agent +from haiku.rag.agents.qa import QuestionAnswerAgent, get_qa_agent from haiku.rag.agents.qa.prompts import QA_SYSTEM_PROMPT from haiku.rag.client import HaikuRAG from haiku.rag.config.models import AppConfig, ModelConfig @@ -78,51 +78,60 @@ class QAPromptAdapter: capture_traces: bool = False, ) -> EvaluationBatch[EvalTrajectory, str | None]: instructions = candidate["instructions"] - return asyncio.run(self._evaluate_async(batch, instructions, capture_traces)) + return asyncio.run( + self._evaluate_with_setup(batch, instructions, capture_traces) + ) + + async def _evaluate_with_setup( + self, + batch: list[QACase], + instructions: str, + capture_traces: bool, + ) -> EvaluationBatch[EvalTrajectory, str | None]: + async with HaikuRAG(self.db_path, config=self.config) as rag: + qa = get_qa_agent(rag, self.config, system_prompt=instructions) + return await self._evaluate_async(batch, qa, capture_traces) async def _evaluate_async( self, batch: list[QACase], - instructions: str, + qa: QuestionAnswerAgent, capture_traces: bool, ) -> EvaluationBatch[EvalTrajectory, str | None]: outputs: list[str | None] = [] scores: list[float] = [] trajectories: list[EvalTrajectory] | None = [] if capture_traces else None - async with HaikuRAG(self.db_path, config=self.config) as rag: - qa = get_qa_agent(rag, self.config, system_prompt=instructions) + for case in batch: + question = case.inputs + expected = case.expected_output or "" - for case in batch: - question = case.inputs - expected = case.expected_output or "" + try: + answer, _ = await qa.answer(question) + except Exception: + logger.warning( + "QA agent failed for question: %s", question, exc_info=True + ) + answer = None - try: - answer, _ = await qa.answer(question) - except Exception: - logger.warning( - "QA agent failed for question: %s", question, exc_info=True - ) - answer = None - - if answer is not None: - score, reason = await self._judge(question, answer, expected) - else: - score, reason = 0.0, "QA agent failed to produce an answer" - - outputs.append(answer) - scores.append(score) - - if capture_traces and trajectories is not None: - trajectories.append( - EvalTrajectory( - question=question, - expected_answer=expected, - actual_answer=answer, - score=score, - judge_reason=reason, - ) + if answer is not None: + score, reason = await self._judge(question, answer, expected) + else: + score, reason = 0.0, "QA agent failed to produce an answer" + + outputs.append(answer) + scores.append(score) + + if capture_traces and trajectories is not None: + trajectories.append( + EvalTrajectory( + question=question, + expected_answer=expected, + actual_answer=answer, + score=score, + judge_reason=reason, ) + ) return EvaluationBatch( outputs=outputs, diff --git a/evaluations/tests/test_optimization.py b/evaluations/tests/test_optimization.py index ea53dee1..5522f964 100644 --- a/evaluations/tests/test_optimization.py +++ b/evaluations/tests/test_optimization.py @@ -3,6 +3,7 @@ from typing import Any from unittest.mock import AsyncMock, MagicMock, patch import pytest +from pydantic_ai.models.test import TestModel from pydantic_evals import Case from gepa.core.adapter import EvaluationBatch @@ -128,22 +129,13 @@ class TestEvaluateAsync: adapter: QAPromptAdapter, sample_cases: list[Case[str, str, dict[str, str]]], ) -> None: - mock_qa = AsyncMock() - mock_qa.answer = AsyncMock(return_value=("X is a thing.", [])) + stub_qa = AsyncMock() + stub_qa.answer = AsyncMock(return_value=("X is a thing.", [])) + adapter._judge = AsyncMock(return_value=(0.85, "Good answer")) # type: ignore[method-assign] - with ( - patch("evaluations.optimization.HaikuRAG") as mock_haiku_cls, - patch("evaluations.optimization.get_qa_agent", return_value=mock_qa), - ): - mock_haiku = AsyncMock() - mock_haiku_cls.return_value.__aenter__ = AsyncMock(return_value=mock_haiku) - mock_haiku_cls.return_value.__aexit__ = AsyncMock(return_value=None) - - adapter._judge = AsyncMock(return_value=(0.85, "Good answer")) # type: ignore[method-assign] - - result = await adapter._evaluate_async( - sample_cases, "test prompt", capture_traces=False - ) + result = await adapter._evaluate_async( + sample_cases, stub_qa, capture_traces=False + ) assert len(result.outputs) == 2 assert len(result.scores) == 2 @@ -157,22 +149,13 @@ class TestEvaluateAsync: adapter: QAPromptAdapter, sample_cases: list[Case[str, str, dict[str, str]]], ) -> None: - mock_qa = AsyncMock() - mock_qa.answer = AsyncMock(return_value=("An answer.", [])) + stub_qa = AsyncMock() + stub_qa.answer = AsyncMock(return_value=("An answer.", [])) + adapter._judge = AsyncMock(return_value=(0.9, "Almost perfect")) # type: ignore[method-assign] - with ( - patch("evaluations.optimization.HaikuRAG") as mock_haiku_cls, - patch("evaluations.optimization.get_qa_agent", return_value=mock_qa), - ): - mock_haiku = AsyncMock() - mock_haiku_cls.return_value.__aenter__ = AsyncMock(return_value=mock_haiku) - mock_haiku_cls.return_value.__aexit__ = AsyncMock(return_value=None) - - adapter._judge = AsyncMock(return_value=(0.9, "Almost perfect")) # type: ignore[method-assign] - - result = await adapter._evaluate_async( - sample_cases, "test prompt", capture_traces=True - ) + result = await adapter._evaluate_async( + sample_cases, stub_qa, capture_traces=True + ) assert result.trajectories is not None assert len(result.trajectories) == 2 @@ -190,20 +173,12 @@ class TestEvaluateAsync: adapter: QAPromptAdapter, sample_cases: list[Case[str, str, dict[str, str]]], ) -> None: - mock_qa = AsyncMock() - mock_qa.answer = AsyncMock(side_effect=RuntimeError("LLM down")) + stub_qa = AsyncMock() + stub_qa.answer = AsyncMock(side_effect=RuntimeError("LLM down")) - with ( - patch("evaluations.optimization.HaikuRAG") as mock_haiku_cls, - patch("evaluations.optimization.get_qa_agent", return_value=mock_qa), - ): - mock_haiku = AsyncMock() - mock_haiku_cls.return_value.__aenter__ = AsyncMock(return_value=mock_haiku) - mock_haiku_cls.return_value.__aexit__ = AsyncMock(return_value=None) - - result = await adapter._evaluate_async( - sample_cases, "test prompt", capture_traces=True - ) + result = await adapter._evaluate_async( + sample_cases, stub_qa, capture_traces=True + ) assert all(o is None for o in result.outputs) assert all(s == 0.0 for s in result.scores) @@ -217,50 +192,40 @@ class TestEvaluateAsync: class TestReflectionLM: def test_handles_string_prompt(self) -> None: - mock_result = MagicMock() - mock_result.output = "Reflected response" - - with ( - patch("evaluations.optimization.get_model"), - patch("pydantic_ai.Agent") as mock_agent_cls, - ): - mock_agent = MagicMock() - mock_agent.run_sync.return_value = mock_result - mock_agent_cls.return_value = mock_agent + test_model = TestModel(custom_output_text="Reflected response") + with patch("evaluations.optimization.get_model", return_value=test_model): lm = ReflectionLM(model_config=AppConfig().qa.model, config=AppConfig()) - lm._agent = mock_agent - result = lm("test prompt") + result = lm("test prompt") assert result == "Reflected response" - mock_agent.run_sync.assert_called_once_with("test prompt") - def test_handles_chat_messages(self) -> None: - mock_result = MagicMock() - mock_result.output = "Chat response" - - with ( - patch("evaluations.optimization.get_model"), - patch("pydantic_ai.Agent") as mock_agent_cls, - ): - mock_agent = MagicMock() - mock_agent.run_sync.return_value = mock_result - mock_agent_cls.return_value = mock_agent + def test_formats_chat_messages_into_string(self) -> None: + test_model = TestModel(custom_output_text="Chat response") + prompts_received: list[str] = [] + with patch("evaluations.optimization.get_model", return_value=test_model): lm = ReflectionLM(model_config=AppConfig().qa.model, config=AppConfig()) - lm._agent = mock_agent - messages: list[dict[str, Any]] = [ - {"role": "system", "content": "You are helpful."}, - {"role": "user", "content": "Hello"}, - ] - result = lm(messages) + original_run_sync = lm._agent.run_sync + + def capturing_run_sync(prompt: str, **kwargs: Any) -> Any: + prompts_received.append(prompt) + return original_run_sync(prompt, **kwargs) + + lm._agent.run_sync = capturing_run_sync # type: ignore[method-assign] + + messages: list[dict[str, Any]] = [ + {"role": "system", "content": "You are helpful."}, + {"role": "user", "content": "Hello"}, + ] + result = lm(messages) assert result == "Chat response" - call_arg = mock_agent.run_sync.call_args[0][0] - assert "system: You are helpful." in call_arg - assert "user: Hello" in call_arg + assert len(prompts_received) == 1 + assert "system: You are helpful." in prompts_received[0] + assert "user: Hello" in prompts_received[0] class TestEvaluateSync: @@ -277,7 +242,7 @@ class TestEvaluateSync: with patch.object( adapter, - "_evaluate_async", + "_evaluate_with_setup", new_callable=AsyncMock, return_value=expected_batch, ) as mock_eval: @@ -306,6 +271,17 @@ def _make_cases(n: int) -> list[Case[str, str, dict[str, str]]]: ] +@pytest.fixture +def gepa_mock_result() -> MagicMock: + mock_result = MagicMock() + mock_result.best_idx = 0 + mock_result.val_aggregate_scores = [0.95] + mock_result.best_candidate = {"instructions": "optimized prompt"} + mock_result.total_metric_calls = 10 + mock_result.num_candidates = 3 + return mock_result + + class TestRunOptimization: def _make_spec(self, db_path: Path) -> DatasetSpec: return DatasetSpec( @@ -318,21 +294,14 @@ class TestRunOptimization: system_prompt="You are a test assistant.", ) - def test_returns_results(self, tmp_path: Path) -> None: + def test_returns_results(self, tmp_path: Path, gepa_mock_result: MagicMock) -> None: spec = self._make_spec(tmp_path / "test.lancedb") cases = _make_cases(4) - mock_result = MagicMock() - mock_result.best_idx = 0 - mock_result.val_aggregate_scores = [0.95] - mock_result.best_candidate = {"instructions": "optimized prompt"} - mock_result.total_metric_calls = 10 - mock_result.num_candidates = 3 - with ( patch("evaluations.optimization.get_model"), patch("evaluations.optimization.ReflectionLM"), - patch("gepa.optimize", return_value=mock_result), + patch("gepa.optimize", return_value=gepa_mock_result), ): result = run_optimization( spec=spec, @@ -347,22 +316,22 @@ class TestRunOptimization: assert result["total_calls"] == 10 assert result["num_candidates"] == 3 - def test_saves_output_file(self, tmp_path: Path) -> None: + def test_saves_output_file( + self, tmp_path: Path, gepa_mock_result: MagicMock + ) -> None: spec = self._make_spec(tmp_path / "test.lancedb") cases = _make_cases(4) output_path = tmp_path / "prompt.txt" - mock_result = MagicMock() - mock_result.best_idx = 0 - mock_result.val_aggregate_scores = [0.85] - mock_result.best_candidate = {"instructions": "saved prompt"} - mock_result.total_metric_calls = 5 - mock_result.num_candidates = 2 + gepa_mock_result.val_aggregate_scores = [0.85] + gepa_mock_result.best_candidate = {"instructions": "saved prompt"} + gepa_mock_result.total_metric_calls = 5 + gepa_mock_result.num_candidates = 2 with ( patch("evaluations.optimization.get_model"), patch("evaluations.optimization.ReflectionLM"), - patch("gepa.optimize", return_value=mock_result), + patch("gepa.optimize", return_value=gepa_mock_result), ): run_optimization( spec=spec,