266 lines
9.1 KiB
Python
266 lines
9.1 KiB
Python
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
|