danvics/haiku.rag
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

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

Browse source
This commit is contained in:
Yiorgis Gozadinos 2026-03-12 12:32:51 +02:00
parent 24275d2ef9
commit 5b3ad9bae7
No known key found for this signature in database
3 changed files with 149 additions and 134 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

51
evaluations/README.md
View file

@ -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:

75
evaluations/evaluations/optimization.py
View file

@ -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,

157
evaluations/tests/test_optimization.py
View file

@ -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,