From adc83e5f816c78cb936cb0da83581ba67a7e919c Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Mon, 9 Mar 2026 12:06:45 +0200 Subject: [PATCH] Tests for gepa --- evaluations/tests/__init__.py | 0 evaluations/tests/test_optimization.py | 266 +++++++++++++++++++++++++ 2 files changed, 266 insertions(+) create mode 100644 evaluations/tests/__init__.py create mode 100644 evaluations/tests/test_optimization.py diff --git a/evaluations/tests/__init__.py b/evaluations/tests/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/evaluations/tests/test_optimization.py b/evaluations/tests/test_optimization.py new file mode 100644 index 00000000..dcfb35bf --- /dev/null +++ b/evaluations/tests/test_optimization.py @@ -0,0 +1,266 @@ +from pathlib import Path +from typing import Any +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from pydantic_evals import Case + +from gepa.core.adapter import EvaluationBatch + +from evaluations.optimization import ( + EvalTrajectory, + QAPromptAdapter, + ReflectionLM, +) +from haiku.rag.config.models import AppConfig + + +@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: + mock_qa = AsyncMock() + mock_qa.answer = AsyncMock(return_value=("X is a thing.", [])) + + 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 + ) + + 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: + mock_qa = AsyncMock() + mock_qa.answer = AsyncMock(return_value=("An answer.", [])) + + 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 + ) + + 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: + mock_qa = AsyncMock() + mock_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 + ) + + 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: + 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 + + lm = ReflectionLM(model_config=AppConfig().qa.model, config=AppConfig()) + lm._agent = mock_agent + + 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 + + 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) + + 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 + + +class TestProposalAttribute: + def test_propose_new_texts_is_none(self, adapter: QAPromptAdapter) -> None: + assert adapter.propose_new_texts is None