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

Drop GEPA prompt optimization and --target qa from evaluations

Browse source
This commit is contained in:
Yiorgis Gozadinos 2026-05-19 11:00:13 +03:00
parent 947a26b391
commit ceb645564c
No known key found for this signature in database
10 changed files with 46 additions and 936 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
CHANGELOG.md
View file

@ -17,6 +17,8 @@
- `AnalysisResult.program`. The per-execution programs are still tracked on `AnalysisState.executions` (the analysis skill's `execute_code` tool populates it); consumers that need the executed code should pull it from the skill state instead of the function return value.
- `--cite` flag on `haiku-rag ask`. Citations always render after the answer now.
- `system_prompt` kwarg on `client.ask`. No production caller used it; `config.prompts.domain_preamble` already covers the preamble use case.
- `evaluations optimize` subcommand and the GEPA prompt-optimization module. Drops the `gepa` dependency. Hand-tuning SKILL.md / `prompts.qa` against `evaluations run` is the loop we actually iterate; the auto-mutation surface was unused and conflicted with the project's no-prompt-string-tests rule.
- `--target qa` from `evaluations run`. The skill path supersedes the standalone QA agent — use `--target rag-skill` (now the default) or `--target analysis-skill`.
### Changed

37
docs/tuning.md
View file

@ -34,8 +34,6 @@ Model and temperature selection affect answer quality directly — see [Provider
`domain_preamble` prepends domain context to all agent prompts — including the main agent, skill subagents, and internal agents (QA, research). Use it to describe what the knowledge base contains and clarify domain-specific terminology. For full prompt replacement, set `prompts.qa` directly. See [Prompt Customization](configuration/prompts.md).
For automated prompt optimization, see [Prompt Optimization (GEPA)](#prompt-optimization-gepa) below.
## What Requires a Rebuild
| Change | Rebuild required? |
@ -66,38 +64,3 @@ evaluations run <dataset> --limit 50
```
See [Benchmarks](benchmarks.md) for dataset details, methodology, and baseline results.
## Prompt Optimization (GEPA)
The `evaluations optimize` command uses GEPA (Generalized Evolutionary Prompt Algorithm) to evolve the QA system prompt. It evaluates candidates on minibatches scored by an LLM judge, reflects on failures, proposes mutations, and accepts improvements.
```bash
# Basic optimization
evaluations optimize wix
# Constrained run
evaluations optimize repliqa --limit 40 --num-candidates 30
# Save result
evaluations optimize wix --output optimized_prompt.txt
```
| Option | Default | Description |
|--------|---------|-------------|
| `--limit` | all cases | QA cases to use (split 50/50 train/val) |
| `--num-candidates` | `50` | Number of candidate prompts to evaluate |
| `--output` | — | Save optimized prompt to file |
| `--config` | auto | haiku.rag YAML config path |
| `--db` | auto | Database path override |
| `--judge-model` | `config.qa.model` | LLM judge as `provider:name` |
| `--reflect-model` | `config.qa.model` | Reflection LLM as `provider:name` |
Apply the result in your config:
```yaml
prompts:
qa: |
Your optimized prompt text here...
```
Or programmatically: `get_qa_agent(client, config, system_prompt=optimized_prompt)`.

28
evaluations/README.md
View file

@ -6,7 +6,7 @@ This package is not published to PyPI and is only used for development and testi
## Overview
Contains evaluation scripts for benchmarking RAG retrieval and QA performance, plus GEPA-based prompt optimization. Available datasets:
Contains evaluation scripts for benchmarking RAG retrieval and QA performance. Available datasets:
- RepliQA (`repliqa`)
- WiX (`wix`)
@ -41,10 +41,11 @@ evaluations run repliqa --skip-qa
evaluations run repliqa --limit 100
```
### Benchmarking the skills
### Choosing the target
By default `evaluations run` benchmarks the QA agent. Pass `--target` to
benchmark the RAG or analysis skill instead, against the same datasets and judge:
`evaluations run` benchmarks `--target rag-skill` by default. Use
`--target analysis-skill` to benchmark the analysis skill against the same
datasets and judge:
```bash
evaluations run wix --target rag-skill
@ -52,9 +53,10 @@ evaluations run wix --target analysis-skill --skill-model ollama:gpt-oss
```
`--skill-model "provider:name"` overrides the skill model independently from
the judge (defaults to `qa.model`). For skill targets, a citation retrieval
metric (`cited_mrr` / `cited_map`) is computed alongside QA accuracy from the
URIs the skill registered via the `cite` tool.
the judge (defaults to `qa.model`, or `analysis.model` when set for the
analysis-skill target). A citation retrieval metric (`cited_mrr` / `cited_map`)
is computed alongside QA accuracy from the URIs the skill registered via the
`cite` tool.
### Pre-built Databases
@ -73,18 +75,6 @@ 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 --num-candidates 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:

127
evaluations/evaluations/benchmark.py
View file

@ -28,7 +28,6 @@ from haiku.rag.client import HaikuRAG
from haiku.rag.config import AppConfig, find_config_file, load_yaml_config
from haiku.rag.config.models import ModelConfig
from haiku.rag.logging import configure_cli_logging
from haiku.rag.agents.qa import get_qa_agent
from haiku.rag.utils import get_model, parse_model_option
_CITATION_EVALUATORS: dict[type[Evaluator], type[Evaluator]] = {
@ -36,8 +35,8 @@ _CITATION_EVALUATORS: dict[type[Evaluator], type[Evaluator]] = {
MAPEvaluator: CitationMAPEvaluator,
}
Target = Literal["qa", "rag-skill", "analysis-skill"]
TARGETS: tuple[Target, ...] = ("qa", "rag-skill", "analysis-skill")
Target = Literal["rag-skill", "analysis-skill"]
TARGETS: tuple[Target, ...] = ("rag-skill", "analysis-skill")
# Pinned judge model. Decoupled from `config.qa.model` so a user changing
# their QA model does not inadvertently change the judge — keeps cross-run
@ -59,7 +58,7 @@ def build_experiment_metadata(
test_cases: int,
config: AppConfig,
judge_config: ModelConfig | None = None,
target: Target = "qa",
target: Target = "rag-skill",
skill_config: ModelConfig | None = None,
) -> dict[str, Any]:
"""Build experiment metadata for Logfire tracking."""
@ -350,7 +349,7 @@ async def run_qa_benchmark(
name: str | None = None,
db_path: Path | None = None,
judge_model: ModelConfig | None = None,
target: Target = "qa",
target: Target = "rag-skill",
skill_model: ModelConfig | None = None,
) -> ReportCaseFailure[str, str, dict[str, str]] | None:
corpus = spec.qa_loader()
@ -363,9 +362,7 @@ async def run_qa_benchmark(
]
judge_config = judge_model or DEFAULT_JUDGE_MODEL
if target == "qa":
skill_config = None
elif target == "analysis-skill":
if target == "analysis-skill":
# Mirror the skill-code resolver: explicit analysis.model wins,
# else fall back to qa.model.
skill_config = skill_model or config.analysis.model or config.qa.model
@ -373,10 +370,8 @@ async def run_qa_benchmark(
skill_config = skill_model or config.qa.model
db = spec.db_path(db_path)
citation_evaluator: Evaluator | None = None
if target != "qa":
_attach_relevant_uris(cases, spec, limit)
citation_evaluator = _citation_evaluator_for(spec.retrieval_evaluator)
_attach_relevant_uris(cases, spec, limit)
citation_evaluator = _citation_evaluator_for(spec.retrieval_evaluator)
evaluators: list[Evaluator] = [
LLMJudge(
@ -416,32 +411,21 @@ async def run_qa_benchmark(
metadata=experiment_metadata,
)
if target == "qa":
async with HaikuRAG(db, config=config) as rag:
qa = get_qa_agent(rag, config)
skill_factory = _skill_factory_for_target(target)
resolved_skill_model = get_model(skill_config, config)
async def answer_question(question: str) -> str:
answer, _ = await qa.answer(question)
return answer
async def answer_question(question: str) -> str:
result = await run_skill_question(
skill_factory=skill_factory,
db_path=db,
config=config,
question=question,
skill_model=resolved_skill_model,
)
set_eval_attribute("cited_uris", result.cited_uris)
return result.answer
report = await _evaluate(answer_question)
else:
skill_factory = _skill_factory_for_target(target)
assert skill_config is not None
resolved_skill_model = get_model(skill_config, config)
async def answer_question(question: str) -> str:
result = await run_skill_question(
skill_factory=skill_factory,
db_path=db,
config=config,
question=question,
skill_model=resolved_skill_model,
)
set_eval_attribute("cited_uris", result.cited_uris)
return result.answer
report = await _evaluate(answer_question)
report = await _evaluate(answer_question)
passing_cases = sum(
1
@ -505,7 +489,7 @@ async def evaluate_dataset(
vacuum_interval: int = 100,
multimodal_only: bool = False,
judge_model: ModelConfig | None = None,
target: Target = "qa",
target: Target = "rag-skill",
skill_model: ModelConfig | None = None,
) -> None:
if not skip_db:
@ -613,16 +597,16 @@ def run(
help="Judge model as 'provider:name'. Defaults to ollama:qwen3.6.",
),
target: str = typer.Option(
"qa",
"rag-skill",
"--target",
help="What to benchmark: qa | rag-skill | analysis-skill.",
help="What to benchmark: rag-skill | analysis-skill.",
),
skill_model: str | None = typer.Option(
None,
"--skill-model",
help=(
"Skill model as 'provider:name'. Used when --target is rag-skill or "
"analysis-skill. Defaults to qa.model from the config."
"Skill model as 'provider:name'. Defaults to qa.model (or "
"analysis.model when --target is analysis-skill) from the config."
),
),
) -> None:
@ -635,10 +619,6 @@ def run(
target_value = cast(Target, target)
judge_model_config = parse_model_option(judge_model) if judge_model else None
skill_model_config = parse_model_option(skill_model) if skill_model else None
if target_value == "qa" and skill_model_config is not None:
raise typer.BadParameter(
"--skill-model is only valid when --target is rag-skill or analysis-skill."
)
asyncio.run(
evaluate_dataset(
@ -659,63 +639,6 @@ def run(
)
@app.command()
def optimize(
dataset: str = typer.Argument(..., help="Dataset key to optimize prompt for."),
config: Path | None = typer.Option(
None, "--config", help="Path to haiku.rag YAML config file."
),
db: Path | None = typer.Option(None, "--db", help="Override the database path."),
limit: int | None = typer.Option(
None, "--limit", help="Limit QA cases (split 50/50 into train/val)."
),
num_candidates: int = typer.Option(
50, "--num-candidates", help="Number of candidate prompts to evaluate."
),
output: Path | None = typer.Option(
None, "--output", help="Save optimized prompt to file."
),
judge_model: str | None = typer.Option(
None,
"--judge-model",
help="Judge model as 'provider:name'. Defaults to ollama:qwen3.6.",
),
reflect_model: str | None = typer.Option(
None,
"--reflect-model",
help="Reflect model as 'provider:name' (e.g. 'anthropic:claude-sonnet-4-20250514').",
),
) -> None:
"""Optimize QA system prompt using GEPA evolutionary optimization."""
from evaluations.optimization import run_optimization
spec = _resolve_dataset(dataset)
app_config = _load_config(config)
corpus = spec.qa_loader()
if limit is not None:
corpus = corpus.select(range(min(limit, len(corpus))))
cases: list[Case[str, str, dict[str, str]]] = [
spec.qa_case_builder(index, cast(Mapping[str, Any], doc))
for index, doc in enumerate(corpus, start=1)
]
judge_model_config = parse_model_option(judge_model) if judge_model else None
reflect_model_config = parse_model_option(reflect_model) if reflect_model else None
run_optimization(
spec=spec,
config=app_config,
cases=cases,
num_candidates=num_candidates,
db_path=db,
output=output,
judge_model=judge_model_config,
reflect_model=reflect_model_config,
)
@app.command()
def download(
dataset: str = typer.Argument(..., help="Dataset key or 'all' to download all."),

292
evaluations/evaluations/optimization.py
View file

@ -1,292 +0,0 @@
import asyncio
import logging
from collections.abc import Mapping, Sequence
from dataclasses import dataclass
from pathlib import Path
from typing import Any
from pydantic_ai.models import Model
from pydantic_evals import Case
from pydantic_evals.evaluators.llm_as_a_judge import judge_input_output_expected
from gepa.core.adapter import EvaluationBatch
from evaluations.config import DatasetSpec
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
from haiku.rag.utils import get_model
logger = logging.getLogger(__name__)
OPTIMIZATION_SCORING_RUBRIC = """You are evaluating the quality of an answer to a question,
comparing it against a reference answer.
Score on a scale of 0.0 to 1.0:
- 1.0: The answer is factually correct, complete, and concise. It covers all key points
from the reference answer without contradictions or significant omissions.
- 0.7-0.9: The answer is mostly correct and addresses the core question, but may miss
some secondary details or include minor inaccuracies.
- 0.4-0.6: The answer is partially correct — it addresses some aspects of the question
but misses key information or contains notable inaccuracies.
- 0.1-0.3: The answer is mostly incorrect or fails to address the core question,
though it may contain some tangentially relevant information.
- 0.0: The answer is completely wrong, irrelevant, or empty.
GUIDELINES:
- Focus on factual correctness relative to the reference answer
- Ignore differences in phrasing, style, or formatting
- A concise correct answer scores higher than a verbose partially correct one
- "I cannot find enough information" when the reference has an answer scores 0.0
"""
@dataclass
class EvalTrajectory:
"""Per-case evaluation result for GEPA reflection."""
question: str
expected_answer: str
actual_answer: str | None
score: float
judge_reason: str | None = None
QACase = Case[str, str, dict[str, str]]
@dataclass
class QAPromptAdapter:
"""GEPA adapter that evaluates QA prompt candidates against a dataset.
Implements the GEPAAdapter protocol:
- evaluate(): Run QA agent with candidate prompt, score with LLMJudge
- make_reflective_dataset(): Build failure records for the GEPA proposer
"""
config: AppConfig
db_path: Path
judge_model: Model
def evaluate(
self,
batch: list[QACase],
candidate: dict[str, str],
capture_traces: bool = False,
) -> EvaluationBatch[EvalTrajectory, str | None]:
instructions = candidate["instructions"]
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],
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
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
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,
scores=scores,
trajectories=trajectories,
)
async def _judge(
self, question: str, answer: str, expected: str
) -> tuple[float, str | None]:
"""Score an answer using pydantic-evals LLMJudge with float scoring."""
result = await judge_input_output_expected(
inputs=question,
output=answer,
expected_output=expected,
rubric=OPTIMIZATION_SCORING_RUBRIC,
model=self.judge_model,
)
return result.score, result.reason
def make_reflective_dataset(
self,
candidate: dict[str, str],
eval_batch: EvaluationBatch[EvalTrajectory, str | None],
components_to_update: list[str],
) -> Mapping[str, Sequence[Mapping[str, Any]]]:
if eval_batch.trajectories is None:
return {}
records: list[dict[str, Any]] = []
for traj in eval_batch.trajectories:
records.append(
{
"Inputs": {"question": traj.question},
"Generated Outputs": {
"answer": traj.actual_answer or "(no answer)"
},
"Feedback": (
f"Expected answer: {traj.expected_answer}\n"
f"Score: {traj.score:.2f}\n"
f"Judge reasoning: {traj.judge_reason or 'N/A'}"
),
}
)
return {"instructions": records}
propose_new_texts = None
class ReflectionLM:
"""LanguageModel implementation for GEPA's ReflectiveMutationProposer.
Wraps a pydantic-ai Agent to satisfy GEPA's LanguageModel protocol.
"""
def __init__(self, model_config: ModelConfig, config: AppConfig) -> None:
from pydantic_ai import Agent
model = get_model(model_config, config)
self._agent: Agent[None, str] = Agent(model=model, output_type=str)
def __call__(self, prompt: str | list[dict[str, Any]]) -> str:
if isinstance(prompt, list):
text = "\n".join(
f"{msg.get('role', 'user')}: {msg.get('content', '')}" for msg in prompt
)
else:
text = prompt
result = self._agent.run_sync(text)
return result.output
# Cases per GEPA reflection minibatch (used for budget calculation)
REFLECTION_MINIBATCH_SIZE = 3
def run_optimization(
spec: DatasetSpec,
config: AppConfig,
cases: list[QACase],
num_candidates: int,
db_path: Path | None = None,
output: Path | None = None,
judge_model: ModelConfig | None = None,
reflect_model: ModelConfig | None = None,
) -> dict[str, Any]:
"""Run GEPA optimization and return results summary."""
from rich.console import Console
console = Console()
judge_config = judge_model or config.qa.model
judge = get_model(judge_config, config)
db = spec.db_path(db_path)
adapter = QAPromptAdapter(
config=config,
db_path=db,
judge_model=judge,
)
reflect_config = reflect_model or config.qa.model
reflection_lm = ReflectionLM(reflect_config, config)
seed_prompt = config.prompts.qa or QA_SYSTEM_PROMPT
seed_candidate = {"instructions": seed_prompt}
mid = len(cases) // 2
trainset = cases[:mid]
valset = cases[mid:]
# Budget: initial valset eval + per-candidate worst case
# (each candidate: 2 minibatch evals + full valset if accepted)
max_metric_calls = len(valset) + num_candidates * (
2 * REFLECTION_MINIBATCH_SIZE + len(valset)
)
console.print(f"Optimizing prompt for dataset: {spec.key}", style="bold magenta")
console.print(
f"Train: {len(trainset)}, Val: {len(valset)}, "
f"Candidates: {num_candidates}, Budget: {max_metric_calls} eval calls"
)
console.print(f"Seed prompt length: {len(seed_prompt)} chars")
from gepa import optimize as gepa_optimize
result = gepa_optimize(
seed_candidate=seed_candidate,
trainset=trainset,
valset=valset,
adapter=adapter,
reflection_lm=reflection_lm,
max_metric_calls=max_metric_calls,
display_progress_bar=True,
)
best_score = result.val_aggregate_scores[result.best_idx]
best_prompt = result.best_candidate
if isinstance(best_prompt, dict):
best_prompt = best_prompt["instructions"]
total_calls = result.total_metric_calls or "unknown"
console.print("\n=== Optimization Results ===", style="bold cyan")
console.print(f"Total metric calls: {total_calls}")
console.print(f"Candidates explored: {result.num_candidates}")
console.print(f"Best score: {best_score:.4f}")
console.print(f"\nOptimized prompt:\n{best_prompt}")
if output:
output.write_text(best_prompt)
console.print(f"\nSaved to: {output}", style="green")
return {
"best_score": best_score,
"best_prompt": best_prompt,
"total_calls": total_calls,
"num_candidates": result.num_candidates,
}

1
evaluations/pyproject.toml
View file

@ -14,7 +14,6 @@ dependencies = [
"huggingface_hub>=0.20.0",
"typer>=0.21.0,<0.22.0",
"python-dotenv>=1.2.2",
"gepa>=0.1.0",
]
[project.scripts]

20
evaluations/tests/test_benchmark.py
View file

@ -131,8 +131,7 @@ class TestRunQaBenchmarkJudgeModel:
with (
patch("evaluations.benchmark.get_model") as mock_get_model,
patch("evaluations.benchmark.HaikuRAG"),
patch("evaluations.benchmark.get_qa_agent"),
patch("evaluations.benchmark.run_skill_question", new_callable=AsyncMock),
):
mock_get_model.return_value = "fake-model"
await run_qa_benchmark(
@ -142,7 +141,7 @@ class TestRunQaBenchmarkJudgeModel:
judge_model=custom_judge,
)
mock_get_model.assert_called_once_with(custom_judge, AppConfig())
mock_get_model.assert_any_call(custom_judge, AppConfig())
@pytest.mark.asyncio
async def test_defaults_to_pinned_judge_model(self, tmp_path: Path) -> None:
@ -150,8 +149,7 @@ class TestRunQaBenchmarkJudgeModel:
with (
patch("evaluations.benchmark.get_model") as mock_get_model,
patch("evaluations.benchmark.HaikuRAG"),
patch("evaluations.benchmark.get_qa_agent"),
patch("evaluations.benchmark.run_skill_question", new_callable=AsyncMock),
):
mock_get_model.return_value = "fake-model"
await run_qa_benchmark(
@ -160,7 +158,7 @@ class TestRunQaBenchmarkJudgeModel:
db_path=tmp_path / "test.lancedb",
)
mock_get_model.assert_called_once_with(DEFAULT_JUDGE_MODEL, AppConfig())
mock_get_model.assert_any_call(DEFAULT_JUDGE_MODEL, AppConfig())
class TestEvaluateDatasetJudgeModel:
@ -197,11 +195,11 @@ class TestEvaluateDatasetJudgeModel:
class TestExperimentMetadataTargets:
def test_default_target_is_qa(self) -> None:
def test_default_target_is_rag_skill(self) -> None:
result = build_experiment_metadata(
dataset_key="test", test_cases=1, config=AppConfig()
)
assert result["target"] == "qa"
assert result["target"] == "rag-skill"
assert "skill_provider" not in result
assert "skill_model" not in result
@ -255,7 +253,7 @@ class TestEvaluateDatasetTarget:
assert mock_qa.call_args[1]["skill_model"] is skill
@pytest.mark.asyncio
async def test_default_target_is_qa(self) -> None:
async def test_default_target_is_rag_skill(self) -> None:
with patch(
"evaluations.benchmark.run_qa_benchmark", new_callable=AsyncMock
) as mock_qa:
@ -269,7 +267,7 @@ class TestEvaluateDatasetTarget:
name=None,
db_path=None,
)
assert mock_qa.call_args[1]["target"] == "qa"
assert mock_qa.call_args[1]["target"] == "rag-skill"
assert mock_qa.call_args[1]["skill_model"] is None
@ -323,7 +321,7 @@ class TestRunQaBenchmarkSkillTarget:
assert _skill_factory_for_target("rag-skill") is rag_factory
assert _skill_factory_for_target("analysis-skill") is analysis_factory
with pytest.raises(ValueError, match="not a skill target"):
_skill_factory_for_target("qa") # type: ignore[arg-type]
_skill_factory_for_target("unknown") # type: ignore[arg-type] # ty: ignore[invalid-argument-type]
class TestCitationEvaluatorWiring:

462
evaluations/tests/test_optimization.py
View file

@ -1,462 +0,0 @@
from pathlib import Path
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
from evaluations.config import DatasetSpec
from evaluations.optimization import (
EvalTrajectory,
QAPromptAdapter,
ReflectionLM,
run_optimization,
)
from haiku.rag.config.models import AppConfig, ModelConfig
@pytest.fixture
def sample_cases() -> list[Case[str, str, dict[str, str]]]:
return [
Case(
name="q1",
inputs="What is X?",
expected_output="X is a thing.",
metadata={"case_index": "1"},
),
Case(
name="q2",
inputs="How does Y work?",
expected_output="Y works by Z.",
metadata={"case_index": "2"},
),
]
@pytest.fixture
def adapter(tmp_path: Path) -> QAPromptAdapter:
return QAPromptAdapter(
config=AppConfig(),
db_path=tmp_path / "test.lancedb",
judge_model=MagicMock(),
)
class TestMakeReflectiveDataset:
def test_builds_records_from_trajectories(self, adapter: QAPromptAdapter) -> None:
trajectories = [
EvalTrajectory(
question="What is X?",
expected_answer="X is a thing.",
actual_answer="X is wrong.",
score=0.2,
judge_reason="Factually incorrect",
),
EvalTrajectory(
question="How does Y?",
expected_answer="Y works by Z.",
actual_answer="Y works by Z.",
score=1.0,
judge_reason=None,
),
]
eval_batch: EvaluationBatch[EvalTrajectory, str | None] = EvaluationBatch(
outputs=["X is wrong.", "Y works by Z."],
scores=[0.2, 1.0],
trajectories=trajectories,
)
result = adapter.make_reflective_dataset(
{"instructions": "test"}, eval_batch, ["instructions"]
)
assert "instructions" in result
records = result["instructions"]
assert len(records) == 2
assert records[0]["Inputs"]["question"] == "What is X?"
assert records[0]["Generated Outputs"]["answer"] == "X is wrong."
assert "Expected answer: X is a thing." in records[0]["Feedback"]
assert "Score: 0.20" in records[0]["Feedback"]
assert "Factually incorrect" in records[0]["Feedback"]
assert records[1]["Inputs"]["question"] == "How does Y?"
assert records[1]["Generated Outputs"]["answer"] == "Y works by Z."
assert "Score: 1.00" in records[1]["Feedback"]
assert "N/A" in records[1]["Feedback"]
def test_returns_empty_when_no_trajectories(self, adapter: QAPromptAdapter) -> None:
eval_batch: EvaluationBatch[EvalTrajectory, str | None] = EvaluationBatch(
outputs=[], scores=[], trajectories=None
)
result = adapter.make_reflective_dataset(
{"instructions": "test"}, eval_batch, ["instructions"]
)
assert result == {}
def test_none_answer_becomes_no_answer(self, adapter: QAPromptAdapter) -> None:
trajectories = [
EvalTrajectory(
question="What?",
expected_answer="Answer.",
actual_answer=None,
score=0.0,
judge_reason="Failed",
),
]
eval_batch: EvaluationBatch[EvalTrajectory, str | None] = EvaluationBatch(
outputs=[None],
scores=[0.0],
trajectories=trajectories,
)
result = adapter.make_reflective_dataset(
{"instructions": "test"}, eval_batch, ["instructions"]
)
assert result["instructions"][0]["Generated Outputs"]["answer"] == "(no answer)"
class TestEvaluateAsync:
@pytest.mark.asyncio
async def test_returns_scores_and_outputs(
self,
adapter: QAPromptAdapter,
sample_cases: list[Case[str, str, dict[str, str]]],
) -> None:
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] # ty: ignore[invalid-assignment]
result = await adapter._evaluate_async(
sample_cases, stub_qa, capture_traces=False
)
assert len(result.outputs) == 2
assert len(result.scores) == 2
assert all(o == "X is a thing." for o in result.outputs)
assert all(s == 0.85 for s in result.scores)
assert result.trajectories is None
@pytest.mark.asyncio
async def test_populates_trajectories_when_captured(
self,
adapter: QAPromptAdapter,
sample_cases: list[Case[str, str, dict[str, str]]],
) -> None:
stub_qa = AsyncMock()
stub_qa.answer = AsyncMock(return_value=("An answer.", []))
adapter._judge = AsyncMock(return_value=(0.9, "Almost perfect")) # type: ignore[method-assign] # ty: ignore[invalid-assignment]
result = await adapter._evaluate_async(
sample_cases, stub_qa, capture_traces=True
)
assert result.trajectories is not None
assert len(result.trajectories) == 2
traj = result.trajectories[0]
assert traj.question == "What is X?"
assert traj.expected_answer == "X is a thing."
assert traj.actual_answer == "An answer."
assert traj.score == 0.9
assert traj.judge_reason == "Almost perfect"
@pytest.mark.asyncio
async def test_handles_qa_failure(
self,
adapter: QAPromptAdapter,
sample_cases: list[Case[str, str, dict[str, str]]],
) -> None:
stub_qa = AsyncMock()
stub_qa.answer = AsyncMock(side_effect=RuntimeError("LLM down"))
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)
assert result.trajectories is not None
assert all(t.actual_answer is None for t in result.trajectories)
assert all(
t.judge_reason == "QA agent failed to produce an answer"
for t in result.trajectories
)
class TestReflectionLM:
def test_handles_string_prompt(self) -> None:
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())
result = lm("test prompt")
assert result == "Reflected response"
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())
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] # ty: ignore[invalid-assignment]
messages: list[dict[str, Any]] = [
{"role": "system", "content": "You are helpful."},
{"role": "user", "content": "Hello"},
]
result = lm(messages)
assert result == "Chat response"
assert len(prompts_received) == 1
assert "system: You are helpful." in prompts_received[0]
assert "user: Hello" in prompts_received[0]
class TestEvaluateSync:
def test_delegates_to_evaluate_async(
self,
adapter: QAPromptAdapter,
sample_cases: list[Case[str, str, dict[str, str]]],
) -> None:
expected_batch: EvaluationBatch[EvalTrajectory, str | None] = EvaluationBatch(
outputs=["answer1", "answer2"],
scores=[0.9, 0.8],
trajectories=None,
)
with patch.object(
adapter,
"_evaluate_with_setup",
new_callable=AsyncMock,
return_value=expected_batch,
) as mock_eval:
result = adapter.evaluate(
sample_cases, {"instructions": "my prompt"}, capture_traces=True
)
mock_eval.assert_called_once_with(sample_cases, "my prompt", True)
assert result is expected_batch
class TestProposalAttribute:
def test_propose_new_texts_is_none(self, adapter: QAPromptAdapter) -> None:
assert adapter.propose_new_texts is None
def _make_cases(n: int) -> list[Case[str, str, dict[str, str]]]:
return [
Case(
name=f"q{i}",
inputs=f"Question {i}?",
expected_output=f"Answer {i}.",
metadata={"case_index": str(i)},
)
for i in range(1, n + 1)
]
@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(
key="test",
db_filename="test.lancedb",
document_loader=lambda: None, # type: ignore[arg-type] # ty: ignore[invalid-argument-type]
document_mapper=lambda doc: None,
qa_loader=lambda: None, # type: ignore[arg-type] # ty: ignore[invalid-argument-type]
qa_case_builder=lambda idx, doc: None, # type: ignore[arg-type] # ty: ignore[invalid-argument-type]
)
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)
with (
patch("evaluations.optimization.get_model"),
patch("evaluations.optimization.ReflectionLM"),
patch("gepa.optimize", return_value=gepa_mock_result),
):
result = run_optimization(
spec=spec,
config=AppConfig(),
cases=cases,
num_candidates=10,
db_path=tmp_path / "test.lancedb",
)
assert result["best_score"] == 0.95
assert result["best_prompt"] == "optimized prompt"
assert result["total_calls"] == 10
assert result["num_candidates"] == 3
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"
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=gepa_mock_result),
):
run_optimization(
spec=spec,
config=AppConfig(),
cases=cases,
num_candidates=5,
db_path=tmp_path / "test.lancedb",
output=output_path,
)
assert output_path.read_text() == "saved prompt"
def test_uses_default_prompt_when_spec_has_none(self, tmp_path: Path) -> None:
spec = DatasetSpec(
key="test",
db_filename="test.lancedb",
document_loader=lambda: None, # type: ignore[arg-type] # ty: ignore[invalid-argument-type]
document_mapper=lambda doc: None,
qa_loader=lambda: None, # type: ignore[arg-type] # ty: ignore[invalid-argument-type]
qa_case_builder=lambda idx, doc: None, # type: ignore[arg-type] # ty: ignore[invalid-argument-type]
)
cases = _make_cases(4)
mock_result = MagicMock()
mock_result.best_idx = 0
mock_result.val_aggregate_scores = [0.5]
mock_result.best_candidate = "fallback prompt"
mock_result.total_metric_calls = 1
mock_result.num_candidates = 1
with (
patch("evaluations.optimization.get_model"),
patch("evaluations.optimization.ReflectionLM"),
patch("gepa.optimize", return_value=mock_result) as mock_gepa,
):
result = run_optimization(
spec=spec,
config=AppConfig(),
cases=cases,
num_candidates=1,
db_path=tmp_path / "test.lancedb",
)
# When best_candidate is a string (not dict), it should be used directly
assert result["best_prompt"] == "fallback prompt"
# Verify seed_candidate used QA_SYSTEM_PROMPT (not None)
call_kwargs = mock_gepa.call_args[1]
seed = call_kwargs["seed_candidate"]
assert seed["instructions"] is not None
assert len(seed["instructions"]) > 0
def test_splits_cases_into_train_and_val(self, tmp_path: Path) -> None:
spec = self._make_spec(tmp_path / "test.lancedb")
cases = _make_cases(10)
mock_result = MagicMock()
mock_result.best_idx = 0
mock_result.val_aggregate_scores = [0.7]
mock_result.best_candidate = {"instructions": "prompt"}
mock_result.total_metric_calls = 50
mock_result.num_candidates = 1
with (
patch("evaluations.optimization.get_model"),
patch("evaluations.optimization.ReflectionLM"),
patch("gepa.optimize", return_value=mock_result) as mock_gepa,
):
run_optimization(
spec=spec,
config=AppConfig(),
cases=cases,
num_candidates=5,
db_path=tmp_path / "test.lancedb",
)
call_kwargs = mock_gepa.call_args[1]
assert len(call_kwargs["trainset"]) == 5
assert len(call_kwargs["valset"]) == 5
# Budget = valset_size + num_candidates * (2*minibatch + valset_size)
assert call_kwargs["max_metric_calls"] == 5 + 5 * (2 * 3 + 5)
def test_uses_custom_reflect_model(
self, tmp_path: Path, gepa_mock_result: MagicMock
) -> None:
spec = self._make_spec(tmp_path / "test.lancedb")
cases = _make_cases(4)
reflect_model = ModelConfig(
provider="anthropic", name="claude-sonnet-4-20250514"
)
with (
patch("evaluations.optimization.get_model"),
patch("evaluations.optimization.ReflectionLM") as mock_rlm,
patch("gepa.optimize", return_value=gepa_mock_result),
):
run_optimization(
spec=spec,
config=AppConfig(),
cases=cases,
num_candidates=10,
db_path=tmp_path / "test.lancedb",
reflect_model=reflect_model,
)
mock_rlm.assert_called_once_with(reflect_model, AppConfig())
def test_uses_custom_judge_model(
self, tmp_path: Path, gepa_mock_result: MagicMock
) -> None:
spec = self._make_spec(tmp_path / "test.lancedb")
cases = _make_cases(4)
judge_model = ModelConfig(provider="openai", name="gpt-4o")
with (
patch("evaluations.optimization.get_model") as mock_get_model,
patch("evaluations.optimization.ReflectionLM"),
patch("gepa.optimize", return_value=gepa_mock_result),
):
run_optimization(
spec=spec,
config=AppConfig(),
cases=cases,
num_candidates=10,
db_path=tmp_path / "test.lancedb",
judge_model=judge_model,
)
mock_get_model.assert_called_once_with(judge_model, AppConfig())

2
evaluations/tests/test_skill_runner.py
View file

@ -296,7 +296,7 @@ class TestRunSkillQuestionEndToEnd:
db_path=rag_db,
config=app_config,
question="What is machine learning?",
skill_model=TestModel(),
skill_model=TestModel(call_tools=["search"]),
)
assert isinstance(result, SkillRunResult)

11
uv.lock
View file

@ -1358,15 +1358,6 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/a3/f6/8ef7e4c286deb2709d11ca96a5237caae3ef4876ab3c48095856cfd2df30/genai_prices-0.0.56-py3-none-any.whl", hash = "sha256:dbe86be8f3f556bed1b72209ed36851fec8b01793b3b220f42921a4e7da945f6", size = 68966, upload-time = "2026-03-20T20:33:02.555Z" },
]
[[package]]
name = "gepa"
version = "0.1.1"
source = { registry = "https://pypi.org/simple" }
sdist = { url = "https://files.pythonhosted.org/packages/e4/62/10f5a8f24c075e3b64f952be73ba8e15f0055584bbcdf9ce48d754a36679/gepa-0.1.1.tar.gz", hash = "sha256:643fda01c23de4c9f01306e01305dd69facc29bcb34ad59e4cd07e6621d34aa1", size = 272251, upload-time = "2026-03-16T10:17:53.131Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/dc/b7/8c72dedbb950d88a6f64588fcbc590d2a21e2b9f19b36aa6c5016c54ec75/gepa-0.1.1-py3-none-any.whl", hash = "sha256:71ead7c591eafcc727b83509cdc4182f20264800a6ddf8520d61419daeb47466", size = 244246, upload-time = "2026-03-16T10:17:51.922Z" },
]
[[package]]
name = "ghp-import"
version = "2.1.0"
@ -1538,7 +1529,6 @@ version = "0.47.0"
source = { editable = "evaluations" }
dependencies = [
{ name = "datasets" },
{ name = "gepa" },
{ name = "haiku-rag-slim" },
{ name = "huggingface-hub" },
{ name = "pydantic-ai-slim", extra = ["evals", "logfire"] },
@ -1549,7 +1539,6 @@ dependencies = [
[package.metadata]
requires-dist = [
{ name = "datasets", specifier = ">=4.6.1" },
{ name = "gepa", specifier = ">=0.1.0" },
{ name = "haiku-rag-slim", editable = "haiku_rag_slim" },
{ name = "huggingface-hub", specifier = ">=0.20.0" },
{ name = "pydantic-ai-slim", extras = ["evals", "logfire"], specifier = ">=1.81.0" },