488 lines
18 KiB
Python
488 lines
18 KiB
Python
from pathlib import Path
|
|
from unittest.mock import AsyncMock, patch
|
|
|
|
import pytest
|
|
import typer
|
|
|
|
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
|
|
|
|
|
|
class TestBuildExperimentMetadata:
|
|
def test_basic_metadata(self) -> None:
|
|
config = AppConfig()
|
|
result = build_experiment_metadata(
|
|
dataset_key="test",
|
|
test_cases=42,
|
|
config=config,
|
|
)
|
|
|
|
assert result["dataset"] == "test"
|
|
assert result["test_cases"] == 42
|
|
assert result["embedder_provider"] == config.embeddings.model.provider
|
|
assert result["embedder_model"] == config.embeddings.model.name
|
|
assert result["embedder_dim"] == config.embeddings.model.vector_dim
|
|
assert result["chunk_size"] == config.processing.chunk_size
|
|
assert result["search_limit"] == config.search.limit
|
|
assert result["qa_provider"] == config.qa.model.provider
|
|
assert result["qa_model"] == config.qa.model.name
|
|
assert "judge_provider" not in result
|
|
|
|
def test_with_judge_config(self) -> None:
|
|
config = AppConfig()
|
|
judge = ModelConfig(
|
|
provider="ollama", name="gpt-oss", enable_thinking=False, temperature=0.0
|
|
)
|
|
result = build_experiment_metadata(
|
|
dataset_key="test",
|
|
test_cases=10,
|
|
config=config,
|
|
judge_config=judge,
|
|
)
|
|
|
|
assert result["judge_provider"] == "ollama"
|
|
assert result["judge_model"] == "gpt-oss"
|
|
assert result["judge_temperature"] == 0.0
|
|
assert result["judge_enable_thinking"] is False
|
|
|
|
def test_no_reranker(self) -> None:
|
|
config = AppConfig()
|
|
result = build_experiment_metadata(
|
|
dataset_key="test", test_cases=1, config=config
|
|
)
|
|
assert result["rerank_provider"] is None
|
|
assert result["rerank_model"] is None
|
|
|
|
def test_with_reranker(self) -> None:
|
|
config = AppConfig()
|
|
config.reranking.model = ModelConfig(
|
|
provider="cross-encoder", name="mixedbread-ai/mxbai-rerank-base-v2"
|
|
)
|
|
result = build_experiment_metadata(
|
|
dataset_key="test", test_cases=1, config=config
|
|
)
|
|
assert result["rerank_provider"] == "cross-encoder"
|
|
assert result["rerank_model"] == "mixedbread-ai/mxbai-rerank-base-v2"
|
|
|
|
|
|
class TestResolveDataset:
|
|
def test_valid_dataset(self) -> None:
|
|
spec = _resolve_dataset("wix")
|
|
assert spec.key == "wix"
|
|
|
|
def test_case_insensitive(self) -> None:
|
|
spec = _resolve_dataset("WIX")
|
|
assert spec.key == "wix"
|
|
|
|
def test_unknown_dataset_raises(self) -> None:
|
|
with pytest.raises(typer.BadParameter, match="Unknown dataset 'nonexistent'"):
|
|
_resolve_dataset("nonexistent")
|
|
|
|
def test_error_lists_valid_datasets(self) -> None:
|
|
with pytest.raises(typer.BadParameter, match="wix"):
|
|
_resolve_dataset("nonexistent")
|
|
|
|
|
|
class TestLoadConfig:
|
|
def test_explicit_path(self, tmp_path: Path) -> None:
|
|
config_file = tmp_path / "test.yaml"
|
|
config_file.write_text("search:\n limit: 42\n")
|
|
config = _load_config(config_file)
|
|
assert config.search.limit == 42
|
|
|
|
def test_explicit_path_not_found(self, tmp_path: Path) -> None:
|
|
with pytest.raises(typer.BadParameter, match="Config file not found"):
|
|
_load_config(tmp_path / "nonexistent.yaml")
|
|
|
|
def test_none_falls_back_to_find_config(self, tmp_path: Path) -> None:
|
|
config_file = tmp_path / "haiku.rag.yaml"
|
|
config_file.write_text("search:\n limit: 99\n")
|
|
with patch("evaluations.benchmark.find_config_file", return_value=config_file):
|
|
config = _load_config(None)
|
|
assert config.search.limit == 99
|
|
|
|
def test_none_no_config_uses_defaults(self) -> None:
|
|
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, # type: ignore[arg-type] # ty: ignore[invalid-argument-type]
|
|
document_mapper=lambda doc: None,
|
|
qa_loader=lambda: [], # 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]
|
|
)
|
|
|
|
@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.run_capability_question", new_callable=AsyncMock
|
|
),
|
|
):
|
|
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_any_call(custom_judge, AppConfig())
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_defaults_to_pinned_judge_model(self, tmp_path: Path) -> None:
|
|
from evaluations.benchmark import DEFAULT_JUDGE_MODEL
|
|
|
|
with (
|
|
patch("evaluations.benchmark.get_model") as mock_get_model,
|
|
patch(
|
|
"evaluations.benchmark.run_capability_question", new_callable=AsyncMock
|
|
),
|
|
):
|
|
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_any_call(DEFAULT_JUDGE_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, # type: ignore[arg-type] # ty: ignore[invalid-argument-type]
|
|
document_mapper=lambda doc: None,
|
|
qa_loader=lambda: [], # 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]
|
|
),
|
|
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
|
|
|
|
|
|
class TestExperimentMetadataTargets:
|
|
def test_default_target_is_rag_capability(self) -> None:
|
|
result = build_experiment_metadata(
|
|
dataset_key="test", test_cases=1, config=AppConfig()
|
|
)
|
|
assert result["target"] == "rag-capability"
|
|
assert "capability_provider" not in result
|
|
assert "capability_model" not in result
|
|
|
|
def test_capability_target_includes_capability_config(self) -> None:
|
|
capability = ModelConfig(
|
|
provider="ollama", name="gpt-oss-large", temperature=0.2
|
|
)
|
|
result = build_experiment_metadata(
|
|
dataset_key="test",
|
|
test_cases=1,
|
|
config=AppConfig(),
|
|
target="rag-capability",
|
|
capability_config=capability,
|
|
)
|
|
assert result["target"] == "rag-capability"
|
|
assert result["capability_provider"] == "ollama"
|
|
assert result["capability_model"] == "gpt-oss-large"
|
|
assert result["capability_temperature"] == 0.2
|
|
|
|
|
|
class TestEvaluateDatasetTarget:
|
|
def _spec(self) -> 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: [], # 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]
|
|
)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_threads_target_and_capability_model(self) -> None:
|
|
capability = ModelConfig(provider="ollama", name="gpt-oss")
|
|
with patch(
|
|
"evaluations.benchmark.run_qa_benchmark", new_callable=AsyncMock
|
|
) as mock_qa:
|
|
await evaluate_dataset(
|
|
spec=self._spec(),
|
|
config=AppConfig(),
|
|
skip_db=True,
|
|
skip_retrieval=True,
|
|
skip_qa=False,
|
|
limit=None,
|
|
name=None,
|
|
db_path=None,
|
|
target="rag-capability",
|
|
capability_model=capability,
|
|
)
|
|
|
|
mock_qa.assert_called_once()
|
|
assert mock_qa.call_args[1]["target"] == "rag-capability"
|
|
assert mock_qa.call_args[1]["capability_model"] is capability
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_default_target_is_rag_capability(self) -> None:
|
|
with patch(
|
|
"evaluations.benchmark.run_qa_benchmark", new_callable=AsyncMock
|
|
) as mock_qa:
|
|
await evaluate_dataset(
|
|
spec=self._spec(),
|
|
config=AppConfig(),
|
|
skip_db=True,
|
|
skip_retrieval=True,
|
|
skip_qa=False,
|
|
limit=None,
|
|
name=None,
|
|
db_path=None,
|
|
)
|
|
assert mock_qa.call_args[1]["target"] == "rag-capability"
|
|
assert mock_qa.call_args[1]["capability_model"] is None
|
|
|
|
|
|
class TestRunQaBenchmarkCapabilityTarget:
|
|
def _spec(self, tmp_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: [], # 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]
|
|
)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_rag_capability_target_uses_run_capability_question(
|
|
self, tmp_path: Path
|
|
) -> None:
|
|
from evaluations.capability_runner import CapabilityRunResult
|
|
|
|
capability_run = AsyncMock(
|
|
return_value=CapabilityRunResult(answer="from capability")
|
|
)
|
|
with (
|
|
patch("evaluations.benchmark.get_model") as mock_get_model,
|
|
patch(
|
|
"evaluations.benchmark.run_capability_question", new=capability_run
|
|
) as mock_run_capability,
|
|
patch("evaluations.benchmark.HaikuRAG") as mock_haiku,
|
|
):
|
|
mock_get_model.return_value = "fake-model"
|
|
await run_qa_benchmark(
|
|
self._spec(tmp_path),
|
|
AppConfig(),
|
|
db_path=tmp_path / "test.lancedb",
|
|
target="rag-capability",
|
|
)
|
|
|
|
# When target is rag-capability, HaikuRAG context manager is NOT entered
|
|
# (the capability manages its own client via lifespan).
|
|
mock_haiku.assert_not_called()
|
|
# capability model defaults to qa.model when not provided
|
|
capability_call = mock_get_model.call_args_list[-1]
|
|
assert capability_call[0][0] == AppConfig().qa.model
|
|
assert mock_run_capability is capability_run
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_analysis_capability_target_resolves_factory(
|
|
self, tmp_path: Path
|
|
) -> None:
|
|
from evaluations.benchmark import _capability_factory_for_target
|
|
from haiku.rag.capabilities.analysis import (
|
|
create_capability as analysis_factory,
|
|
)
|
|
from haiku.rag.capabilities.rag import create_capability as rag_factory
|
|
|
|
assert _capability_factory_for_target("rag-capability") is rag_factory
|
|
assert _capability_factory_for_target("analysis-capability") is analysis_factory
|
|
with pytest.raises(ValueError, match="not a capability target"):
|
|
_capability_factory_for_target("unknown") # type: ignore[arg-type] # ty: ignore[invalid-argument-type]
|
|
|
|
|
|
class TestCitationEvaluatorWiring:
|
|
def test_returns_map_twin_for_map_evaluator(self) -> None:
|
|
from evaluations.benchmark import _citation_evaluator_for
|
|
from evaluations.evaluators import CitationMAPEvaluator, MAPEvaluator
|
|
|
|
result = _citation_evaluator_for(MAPEvaluator())
|
|
assert isinstance(result, CitationMAPEvaluator)
|
|
|
|
def test_returns_none_for_no_evaluator(self) -> None:
|
|
from evaluations.benchmark import _citation_evaluator_for
|
|
|
|
assert _citation_evaluator_for(None) is None
|
|
|
|
|
|
class TestAttachRelevantUris:
|
|
def test_joins_by_question(self) -> None:
|
|
from pydantic_evals import Case
|
|
|
|
from evaluations.benchmark import _attach_relevant_uris
|
|
from evaluations.config import RetrievalSample
|
|
from evaluations.evaluators import MAPEvaluator
|
|
|
|
cases: list[Case[str, str, dict]] = [
|
|
Case(name="c1", inputs="What is X?", expected_output="X is a thing"),
|
|
Case(
|
|
name="c2",
|
|
inputs="What is Y?",
|
|
expected_output="Y is another",
|
|
metadata={"existing": "value"},
|
|
),
|
|
Case(
|
|
name="c3",
|
|
inputs="What is Z?",
|
|
expected_output="not in retrieval set",
|
|
),
|
|
]
|
|
|
|
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: [], # 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]
|
|
retrieval_loader=lambda: [ # type: ignore[arg-type] # ty: ignore[invalid-argument-type]
|
|
{"q": "What is X?", "uris": ("uri-x",)},
|
|
{"q": "What is Y?", "uris": ("uri-y1", "uri-y2")},
|
|
],
|
|
retrieval_mapper=lambda d: RetrievalSample(
|
|
question=d["q"], expected_uris=d["uris"]
|
|
),
|
|
retrieval_evaluator=MAPEvaluator(),
|
|
)
|
|
|
|
_attach_relevant_uris(cases, spec, limit=None)
|
|
|
|
assert cases[0].metadata == {"relevant_uris": ["uri-x"]}
|
|
assert cases[1].metadata == {
|
|
"existing": "value",
|
|
"relevant_uris": ["uri-y1", "uri-y2"],
|
|
}
|
|
# case c3 has no matching retrieval sample — metadata untouched
|
|
assert cases[2].metadata is None
|
|
|
|
def test_no_op_without_retrieval_loader(self) -> None:
|
|
from pydantic_evals import Case
|
|
|
|
from evaluations.benchmark import _attach_relevant_uris
|
|
|
|
cases: list[Case[str, str, dict]] = [
|
|
Case(name="c1", inputs="q", expected_output="a"),
|
|
]
|
|
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: [], # 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]
|
|
)
|
|
_attach_relevant_uris(cases, spec, limit=None)
|
|
assert cases[0].metadata is None
|
|
|
|
|
|
class TestFilterQaCorpus:
|
|
def test_keeps_only_matching_ids(self) -> None:
|
|
from datasets import Dataset
|
|
|
|
from evaluations.benchmark import _filter_qa_corpus
|
|
|
|
corpus = Dataset.from_list(
|
|
[{"id": "a", "q": 1}, {"id": "b", "q": 2}, {"id": "c", "q": 3}]
|
|
)
|
|
out = _filter_qa_corpus(corpus, {"a", "c"})
|
|
assert [r["id"] for r in out] == ["a", "c"]
|
|
|
|
def test_none_returns_corpus_unchanged(self) -> None:
|
|
from datasets import Dataset
|
|
|
|
from evaluations.benchmark import _filter_qa_corpus
|
|
|
|
corpus = Dataset.from_list([{"id": "a"}])
|
|
assert _filter_qa_corpus(corpus, None) is corpus
|
|
|
|
|
|
class TestLoadCaseIds:
|
|
def test_reads_strips_and_drops_blanks(self, tmp_path: Path) -> None:
|
|
from evaluations.benchmark import _load_case_ids
|
|
|
|
f = tmp_path / "ids.txt"
|
|
f.write_text("finqa_dev_16\n finqa_dev_66 \n\n\nfinqa_dev_113\n")
|
|
assert _load_case_ids(f) == {"finqa_dev_16", "finqa_dev_66", "finqa_dev_113"}
|
|
|
|
def test_none_path_returns_none(self) -> None:
|
|
from evaluations.benchmark import _load_case_ids
|
|
|
|
assert _load_case_ids(None) is None
|
|
|
|
|
|
class TestEvaluateDatasetCaseIds:
|
|
def _spec(self) -> 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: [], # 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]
|
|
)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_threads_case_ids_to_qa_benchmark(self) -> None:
|
|
from evaluations.benchmark import evaluate_dataset
|
|
|
|
with patch(
|
|
"evaluations.benchmark.run_qa_benchmark", new_callable=AsyncMock
|
|
) as mock_qa:
|
|
await evaluate_dataset(
|
|
spec=self._spec(),
|
|
config=AppConfig(),
|
|
skip_db=True,
|
|
skip_retrieval=True,
|
|
skip_qa=False,
|
|
limit=None,
|
|
name=None,
|
|
db_path=None,
|
|
case_ids={"finqa_dev_16", "finqa_dev_66"},
|
|
)
|
|
mock_qa.assert_called_once()
|
|
assert mock_qa.call_args[1]["case_ids"] == {"finqa_dev_16", "finqa_dev_66"}
|