From fd327996b894a562c7bce4b4bfd8745b989d55cf Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Wed, 18 Mar 2026 17:14:11 +0200 Subject: [PATCH] Configurable judge and reflect models for evaluations --- CHANGELOG.md | 9 +++ docs/benchmarks.md | 3 +- docs/tuning.md | 2 + evaluations/evaluations/benchmark.py | 46 ++++++++--- evaluations/evaluations/evaluators/judge.py | 12 ++- evaluations/evaluations/optimization.py | 11 ++- evaluations/tests/test_benchmark.py | 85 ++++++++++++++++++++- evaluations/tests/test_optimization.py | 50 +++++++++++- haiku_rag_slim/haiku/rag/utils.py | 12 +++ tests/test_utils.py | 18 +++++ 10 files changed, 228 insertions(+), 20 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 35bd4a2f..0863dd9c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,15 @@ # Changelog ## [Unreleased] +### Added + +- **Configurable judge and reflect models**: `evaluations run` and `evaluations optimize` accept `--judge-model provider:name`; `optimize` also accepts `--reflect-model provider:name`. Both fall back to `config.qa.model` when not specified. +- **`parse_model_option`**: Utility in `haiku.rag.utils` for parsing `provider:name` strings into `ModelConfig` + +### Changed + +- **LLMJudge**: Custom evaluator now accepts `ModelConfig` instead of a model name string + ## [0.34.1] - 2026-03-16 ### Added diff --git a/docs/benchmarks.md b/docs/benchmarks.md index 59bc248e..9add1242 100644 --- a/docs/benchmarks.md +++ b/docs/benchmarks.md @@ -60,6 +60,7 @@ evaluations run repliqa --config /path/to/haiku.rag.yaml --db /path/to/custom.la - `--skip-qa` - Skip QA benchmark - `--limit N` - Limit number of test cases - `--name NAME` - Override the evaluation name +- `--judge-model PROVIDER:NAME` - Override the LLM judge model (default: `config.qa.model`) If no config file is specified, the script searches standard locations: `./haiku.rag.yaml`, user config directory, then falls back to defaults. @@ -84,7 +85,7 @@ If no config file is specified, the script searches standard locations: `./haiku ### QA Accuracy -For question-answering evaluation, `pydantic-evals` coordinates an LLM judge (Ollama `qwen3`) to determine whether answers are correct. Accuracy is the fraction of correctly answered questions. +For question-answering evaluation, `pydantic-evals` coordinates an LLM judge to determine whether answers are correct. By default the judge uses the same model as QA (`config.qa.model`); override with `--judge-model provider:name`. Accuracy is the fraction of correctly answered questions. ## RepliQA diff --git a/docs/tuning.md b/docs/tuning.md index 9a0de721..4dedf3d0 100644 --- a/docs/tuning.md +++ b/docs/tuning.md @@ -89,6 +89,8 @@ evaluations optimize wix --output optimized_prompt.txt | `--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: diff --git a/evaluations/evaluations/benchmark.py b/evaluations/evaluations/benchmark.py index ea1cfb91..1d664b18 100644 --- a/evaluations/evaluations/benchmark.py +++ b/evaluations/evaluations/benchmark.py @@ -22,16 +22,12 @@ 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 +from haiku.rag.utils import get_model, parse_model_option load_dotenv(find_dotenv(usecwd=True)) HF_REPO_ID = "ggozad/haiku-rag-eval-dbs" -JUDGE_MODEL_CONFIG = ModelConfig( - provider="ollama", name="gpt-oss", enable_thinking=False, temperature=0.0 -) - logfire.configure(send_to_logfire="if-token-present", service_name="evals") logfire.instrument_pydantic_ai() configure_cli_logging() @@ -272,6 +268,7 @@ async def run_qa_benchmark( limit: int | None = None, name: str | None = None, db_path: Path | None = None, + judge_model: ModelConfig | None = None, ) -> ReportCaseFailure[str, str, dict[str, str]] | None: corpus = spec.qa_loader() if limit is not None: @@ -282,7 +279,8 @@ async def run_qa_benchmark( for index, doc in enumerate(corpus, start=1) ] - judge_model = get_model(JUDGE_MODEL_CONFIG, config) + judge_config = judge_model or config.qa.model + judge = get_model(judge_config, config) evaluation_dataset = EvalDataset[str, str, dict[str, str]]( name=spec.key, @@ -292,7 +290,7 @@ async def run_qa_benchmark( rubric=ANSWER_EQUIVALENCE_RUBRIC, include_input=True, include_expected_output=True, - model=judge_model, + model=judge, assertion={ "evaluation_name": "answer_equivalent", "include_reason": True, @@ -315,7 +313,7 @@ async def run_qa_benchmark( dataset_key=spec.key, test_cases=len(cases), config=config, - judge_config=JUDGE_MODEL_CONFIG, + judge_config=judge_config, ) report = await evaluation_dataset.evaluate( @@ -364,6 +362,7 @@ async def evaluate_dataset( db_path: Path | None, vacuum_interval: int = 100, multimodal_only: bool = False, + judge_model: ModelConfig | None = None, ) -> None: if not skip_db: console.print(f"Using dataset: {spec.key}", style="bold magenta") @@ -384,7 +383,14 @@ async def evaluate_dataset( if not skip_qa: console.print("\nRunning QA benchmarks...", style="bold yellow") - await run_qa_benchmark(spec, config, limit=limit, name=name, db_path=db_path) + await run_qa_benchmark( + spec, + config, + limit=limit, + name=name, + db_path=db_path, + judge_model=judge_model, + ) app = typer.Typer(help="Run retrieval and QA benchmarks for configured datasets.") @@ -453,9 +459,15 @@ def run( "--multimodal-only", help="Only evaluate queries requiring image understanding.", ), + judge_model: str | None = typer.Option( + None, + "--judge-model", + help="Judge model as 'provider:name' (e.g. 'ollama:gpt-oss').", + ), ) -> None: spec = _resolve_dataset(dataset) app_config = _load_config(config) + judge_model_config = parse_model_option(judge_model) if judge_model else None asyncio.run( evaluate_dataset( @@ -469,6 +481,7 @@ def run( db_path=db, vacuum_interval=vacuum_interval, multimodal_only=multimodal_only, + judge_model=judge_model_config, ) ) @@ -489,6 +502,16 @@ def optimize( 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' (e.g. 'ollama:gpt-oss').", + ), + 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 @@ -505,6 +528,9 @@ def optimize( 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, @@ -512,6 +538,8 @@ def optimize( num_candidates=num_candidates, db_path=db, output=output, + judge_model=judge_model_config, + reflect_model=reflect_model_config, ) diff --git a/evaluations/evaluations/evaluators/judge.py b/evaluations/evaluations/evaluators/judge.py index 14b6d83e..6a1d538a 100644 --- a/evaluations/evaluations/evaluators/judge.py +++ b/evaluations/evaluations/evaluators/judge.py @@ -35,10 +35,14 @@ class LLMJudgeResponseSchema(BaseModel): class LLMJudge: """LLM-as-judge for evaluating answer equivalence using Pydantic AI.""" - def __init__(self, model: str = "gpt-oss", config: AppConfig | None = None): - model_config = ModelConfig( - provider="ollama", name=model, enable_thinking=True, temperature=0.0 - ) + def __init__( + self, + model_config: ModelConfig | None = None, + config: AppConfig | None = None, + ): + if model_config is None: + effective_config = config or AppConfig() + model_config = effective_config.qa.model model_obj = get_model(model_config, config) # Create Pydantic AI agent diff --git a/evaluations/evaluations/optimization.py b/evaluations/evaluations/optimization.py index 39c066ad..6a768406 100644 --- a/evaluations/evaluations/optimization.py +++ b/evaluations/evaluations/optimization.py @@ -11,7 +11,6 @@ from pydantic_evals.evaluators.llm_as_a_judge import judge_input_output_expected from gepa.core.adapter import EvaluationBatch -from evaluations.benchmark import JUDGE_MODEL_CONFIG 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 @@ -216,22 +215,26 @@ def run_optimization( 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_model = get_model(JUDGE_MODEL_CONFIG, config) + 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_model, + judge_model=judge, ) - reflection_lm = ReflectionLM(config.qa.model, config) + reflect_config = reflect_model or config.qa.model + reflection_lm = ReflectionLM(reflect_config, config) seed_prompt = spec.resolve_system_prompt(config) or QA_SYSTEM_PROMPT seed_candidate = {"instructions": seed_prompt} diff --git a/evaluations/tests/test_benchmark.py b/evaluations/tests/test_benchmark.py index 1e62090b..b3a2aa5b 100644 --- a/evaluations/tests/test_benchmark.py +++ b/evaluations/tests/test_benchmark.py @@ -1,5 +1,5 @@ from pathlib import Path -from unittest.mock import patch +from unittest.mock import AsyncMock, patch import pytest import typer @@ -8,7 +8,10 @@ from evaluations.benchmark import ( _load_config, _resolve_dataset, build_experiment_metadata, + evaluate_dataset, + run_qa_benchmark, ) +from evaluations.config import DatasetSpec from haiku.rag.config.models import AppConfig, ModelConfig @@ -110,3 +113,83 @@ class TestLoadConfig: with patch("evaluations.benchmark.find_config_file", return_value=None): config = _load_config(None) assert config == AppConfig() + + +class TestRunQaBenchmarkJudgeModel: + def _make_spec(self) -> DatasetSpec: + return DatasetSpec( + key="test", + db_filename="test.lancedb", + document_loader=lambda: None, + document_mapper=lambda doc: None, + qa_loader=lambda: [], + qa_case_builder=lambda idx, doc: None, + ) + + @pytest.mark.asyncio + async def test_uses_custom_judge_model(self, tmp_path: Path) -> None: + custom_judge = ModelConfig(provider="openai", name="gpt-4o") + + with ( + patch("evaluations.benchmark.get_model") as mock_get_model, + patch("evaluations.benchmark.HaikuRAG"), + patch("evaluations.benchmark.get_qa_agent"), + ): + mock_get_model.return_value = "fake-model" + await run_qa_benchmark( + self._make_spec(), + AppConfig(), + db_path=tmp_path / "test.lancedb", + judge_model=custom_judge, + ) + + mock_get_model.assert_called_once_with(custom_judge, AppConfig()) + + @pytest.mark.asyncio + async def test_defaults_to_judge_model_config(self, tmp_path: Path) -> None: + with ( + patch("evaluations.benchmark.get_model") as mock_get_model, + patch("evaluations.benchmark.HaikuRAG"), + patch("evaluations.benchmark.get_qa_agent"), + ): + mock_get_model.return_value = "fake-model" + await run_qa_benchmark( + self._make_spec(), + AppConfig(), + db_path=tmp_path / "test.lancedb", + ) + + mock_get_model.assert_called_once_with(AppConfig().qa.model, AppConfig()) + + +class TestEvaluateDatasetJudgeModel: + @pytest.mark.asyncio + async def test_threads_judge_model_to_qa_benchmark(self) -> None: + custom_judge = ModelConfig( + provider="anthropic", name="claude-sonnet-4-20250514" + ) + + with patch( + "evaluations.benchmark.run_qa_benchmark", new_callable=AsyncMock + ) as mock_qa: + await evaluate_dataset( + spec=DatasetSpec( + key="test", + db_filename="test.lancedb", + document_loader=lambda: None, + document_mapper=lambda doc: None, + qa_loader=lambda: [], + qa_case_builder=lambda idx, doc: None, + ), + config=AppConfig(), + skip_db=True, + skip_retrieval=True, + skip_qa=False, + limit=None, + name=None, + db_path=None, + judge_model=custom_judge, + ) + + mock_qa.assert_called_once() + assert mock_qa.call_args[1]["judge_model"] is custom_judge diff --git a/evaluations/tests/test_optimization.py b/evaluations/tests/test_optimization.py index cf04ffc1..55e2b81e 100644 --- a/evaluations/tests/test_optimization.py +++ b/evaluations/tests/test_optimization.py @@ -15,7 +15,7 @@ from evaluations.optimization import ( ReflectionLM, run_optimization, ) -from haiku.rag.config.models import AppConfig +from haiku.rag.config.models import AppConfig, ModelConfig @pytest.fixture @@ -413,3 +413,51 @@ class TestRunOptimization: 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()) diff --git a/haiku_rag_slim/haiku/rag/utils.py b/haiku_rag_slim/haiku/rag/utils.py index ca6bfa7a..aac551d0 100644 --- a/haiku_rag_slim/haiku/rag/utils.py +++ b/haiku_rag_slim/haiku/rag/utils.py @@ -15,6 +15,18 @@ if TYPE_CHECKING: from haiku.rag.config.models import AppConfig, ModelConfig +def parse_model_option(value: str) -> "ModelConfig": + """Parse a 'provider:name' string into a ModelConfig.""" + from haiku.rag.config.models import ModelConfig + + parts = value.split(":", 1) + if len(parts) != 2 or not parts[0] or not parts[1]: + raise ValueError( + f"Invalid model format '{value}'. Expected 'provider:name' (e.g. 'ollama:gpt-oss')." + ) + return ModelConfig(provider=parts[0], name=parts[1]) + + def cosine_similarity(vec1: list[float], vec2: list[float]) -> float: """Compute cosine similarity between two vectors.""" dot_product = sum(a * b for a, b in zip(vec1, vec2)) diff --git a/tests/test_utils.py b/tests/test_utils.py index 4c2951bf..cba14326 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -640,3 +640,21 @@ async def test_is_up_to_date(monkeypatch): is_current, running, latest = await is_up_to_date() assert is_current is True assert running >= latest + + +# --- parse_model_option tests --- + + +def test_parse_model_option(): + from haiku.rag.utils import parse_model_option + + result = parse_model_option("anthropic:claude-sonnet-4-20250514") + assert result.provider == "anthropic" + assert result.name == "claude-sonnet-4-20250514" + + # Colons in name are preserved + assert parse_model_option("openai:gpt-4o:latest").name == "gpt-4o:latest" + + for bad in ["just-a-name", ":model", "provider:"]: + with pytest.raises(ValueError, match="Invalid model format"): + parse_model_option(bad)