Refactor QAPromptAdapter to accept QA agent directly, rewrite mock-heavy tests

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
+# Run retrieval + QA benchmarks
-evaluations repliqa
+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:

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:
+        for case in batch:
-            qa = get_qa_agent(rag, self.config, system_prompt=instructions)
+            question = case.inputs
+            expected = case.expected_output or ""
 
-            for case in batch:
+            try:
-                question = case.inputs
+                answer, _ = await qa.answer(question)
-                expected = case.expected_output or ""
+            except Exception:
+                logger.warning(
+                    "QA agent failed for question: %s", question, exc_info=True
+                )
+                answer = None
 
-                try:
+            if answer is not None:
-                    answer, _ = await qa.answer(question)
+                score, reason = await self._judge(question, answer, expected)
-                except Exception:
+            else:
-                    logger.warning(
+                score, reason = 0.0, "QA agent failed to produce an answer"
-                        "QA agent failed for question: %s", question, exc_info=True
+
-                    )
+            outputs.append(answer)
-                    answer = None
+            scores.append(score)
-
+
-                if answer is not None:
+            if capture_traces and trajectories is not None:
-                    score, reason = await self._judge(question, answer, expected)
+                trajectories.append(
-                else:
+                    EvalTrajectory(
-                    score, reason = 0.0, "QA agent failed to produce an answer"
+                        question=question,
-
+                        expected_answer=expected,
-                outputs.append(answer)
+                        actual_answer=answer,
-                scores.append(score)
+                        score=score,
-
+                        judge_reason=reason,
-                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,

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()
+        stub_qa = AsyncMock()
-        mock_qa.answer = AsyncMock(return_value=("X is a thing.", []))
+        stub_qa.answer = AsyncMock(return_value=("X is a thing.", []))
+        adapter._judge = AsyncMock(return_value=(0.85, "Good answer"))  # type: ignore[method-assign]
 
-        with (
+        result = await adapter._evaluate_async(
-            patch("evaluations.optimization.HaikuRAG") as mock_haiku_cls,
+            sample_cases, stub_qa, capture_traces=False
-            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
-            )
 
         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()
+        stub_qa = AsyncMock()
-        mock_qa.answer = AsyncMock(return_value=("An answer.", []))
+        stub_qa.answer = AsyncMock(return_value=("An answer.", []))
+        adapter._judge = AsyncMock(return_value=(0.9, "Almost perfect"))  # type: ignore[method-assign]
 
-        with (
+        result = await adapter._evaluate_async(
-            patch("evaluations.optimization.HaikuRAG") as mock_haiku_cls,
+            sample_cases, stub_qa, capture_traces=True
-            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
-            )
 
         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()
+        stub_qa = AsyncMock()
-        mock_qa.answer = AsyncMock(side_effect=RuntimeError("LLM down"))
+        stub_qa.answer = AsyncMock(side_effect=RuntimeError("LLM down"))
 
-        with (
+        result = await adapter._evaluate_async(
-            patch("evaluations.optimization.HaikuRAG") as mock_haiku_cls,
+            sample_cases, stub_qa, capture_traces=True
-            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
-            )
 
         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()
+        test_model = TestModel(custom_output_text="Reflected response")
-        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
 
+        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:
+    def test_formats_chat_messages_into_string(self) -> None:
-        mock_result = MagicMock()
+        test_model = TestModel(custom_output_text="Chat response")
-        mock_result.output = "Chat response"
+        prompts_received: list[str] = []
-
-        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
 
+        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]] = [
+        original_run_sync = lm._agent.run_sync
-                {"role": "system", "content": "You are helpful."},
+
-                {"role": "user", "content": "Hello"},
+        def capturing_run_sync(prompt: str, **kwargs: Any) -> Any:
-            ]
+            prompts_received.append(prompt)
-            result = lm(messages)
+            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 len(prompts_received) == 1
-        assert "system: You are helpful." in call_arg
+        assert "system: You are helpful." in prompts_received[0]
-        assert "user: Hello" in call_arg
+        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()
+        gepa_mock_result.val_aggregate_scores = [0.85]
-        mock_result.best_idx = 0
+        gepa_mock_result.best_candidate = {"instructions": "saved prompt"}
-        mock_result.val_aggregate_scores = [0.85]
+        gepa_mock_result.total_metric_calls = 5
-        mock_result.best_candidate = {"instructions": "saved prompt"}
+        gepa_mock_result.num_candidates = 2
-        mock_result.total_metric_calls = 5
-        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,