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

benchmark RAG and analysis skills via --target

Browse source
This commit is contained in:
Yiorgis Gozadinos 2026-04-28 12:09:19 +03:00
parent 7685a592c2
commit d00befd0c4
No known key found for this signature in database
4 changed files with 640 additions and 15 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

120
evaluations/evaluations/benchmark.py
View file

@ -2,7 +2,7 @@ import asyncio
import shutil
from collections.abc import Mapping
from pathlib import Path
from typing import Any, cast
from typing import Any, Literal, cast
import logfire
import typer
@ -17,6 +17,7 @@ from rich.progress import Progress
from evaluations.config import DatasetSpec
from evaluations.datasets import DATASETS
from evaluations.evaluators import ANSWER_EQUIVALENCE_RUBRIC
from evaluations.skill_runner import SkillFactory, run_skill_question
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
@ -24,6 +25,9 @@ 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
Target = Literal["qa", "rag-skill", "analysis-skill"]
TARGETS: tuple[Target, ...] = ("qa", "rag-skill", "analysis-skill")
load_dotenv(find_dotenv(usecwd=True))
HF_REPO_ID = "ggozad/haiku-rag-eval-dbs"
@ -39,11 +43,14 @@ def build_experiment_metadata(
test_cases: int,
config: AppConfig,
judge_config: ModelConfig | None = None,
target: Target = "qa",
skill_config: ModelConfig | None = None,
) -> dict[str, Any]:
"""Build experiment metadata for Logfire tracking."""
metadata: dict[str, Any] = {
"dataset": dataset_key,
"test_cases": test_cases,
"target": target,
"embedder_provider": config.embeddings.model.provider,
"embedder_model": config.embeddings.model.name,
"embedder_dim": config.embeddings.model.vector_dim,
@ -71,6 +78,16 @@ def build_experiment_metadata(
"judge_enable_thinking": judge_config.enable_thinking,
}
)
if skill_config is not None:
metadata.update(
{
"skill_provider": skill_config.provider,
"skill_model": skill_config.name,
"skill_temperature": skill_config.temperature,
"skill_max_tokens": skill_config.max_tokens,
"skill_enable_thinking": skill_config.enable_thinking,
}
)
return metadata
@ -260,6 +277,18 @@ async def run_retrieval_benchmark(
}
def _skill_factory_for_target(target: Target) -> SkillFactory:
if target == "rag-skill":
from haiku.rag.skills.rag import create_skill
return create_skill
if target == "analysis-skill":
from haiku.rag.skills.analysis import create_skill
return create_skill
raise ValueError(f"target {target!r} is not a skill target")
async def run_qa_benchmark(
spec: DatasetSpec,
config: AppConfig,
@ -267,6 +296,8 @@ async def run_qa_benchmark(
name: str | None = None,
db_path: Path | None = None,
judge_model: ModelConfig | None = None,
target: Target = "qa",
skill_model: ModelConfig | None = None,
) -> ReportCaseFailure[str, str, dict[str, str]] | None:
corpus = spec.qa_loader()
if limit is not None:
@ -298,21 +329,49 @@ async def run_qa_benchmark(
)
db = spec.db_path(db_path)
async with HaikuRAG(db, config=config) as rag:
qa = get_qa_agent(rag, config, system_prompt=spec.resolve_system_prompt(config))
skill_config = skill_model or config.qa.model if target != "qa" else None
eval_name = name if name is not None else f"{spec.key}_qa_evaluation"
experiment_metadata = build_experiment_metadata(
dataset_key=spec.key,
test_cases=len(cases),
config=config,
judge_config=judge_config,
target=target,
skill_config=skill_config,
)
if target == "qa":
async with HaikuRAG(db, config=config) as rag:
qa = get_qa_agent(
rag, config, system_prompt=spec.resolve_system_prompt(config)
)
async def answer_question(question: str) -> str:
answer, _ = await qa.answer(question)
return answer
report = await evaluation_dataset.evaluate(
answer_question,
name=eval_name,
max_concurrency=1,
progress=True,
metadata=experiment_metadata,
)
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:
answer, _ = await qa.answer(question)
return answer
eval_name = name if name is not None else f"{spec.key}_qa_evaluation"
experiment_metadata = build_experiment_metadata(
dataset_key=spec.key,
test_cases=len(cases),
config=config,
judge_config=judge_config,
)
result = await run_skill_question(
skill_factory=skill_factory,
db_path=db,
config=config,
question=question,
skill_model=resolved_skill_model,
)
return result.answer
report = await evaluation_dataset.evaluate(
answer_question,
@ -361,6 +420,8 @@ async def evaluate_dataset(
vacuum_interval: int = 100,
multimodal_only: bool = False,
judge_model: ModelConfig | None = None,
target: Target = "qa",
skill_model: ModelConfig | None = None,
) -> None:
if not skip_db:
console.print(f"Using dataset: {spec.key}", style="bold magenta")
@ -380,7 +441,9 @@ async def evaluate_dataset(
)
if not skip_qa:
console.print("\nRunning QA benchmarks...", style="bold yellow")
console.print(
f"\nRunning QA benchmarks (target={target})...", style="bold yellow"
)
await run_qa_benchmark(
spec,
config,
@ -388,6 +451,8 @@ async def evaluate_dataset(
name=name,
db_path=db_path,
judge_model=judge_model,
target=target,
skill_model=skill_model,
)
@ -462,10 +527,33 @@ def run(
"--judge-model",
help="Judge model as 'provider:name' (e.g. 'ollama:gpt-oss').",
),
target: str = typer.Option(
"qa",
"--target",
help="What to benchmark: qa | 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."
),
),
) -> None:
spec = _resolve_dataset(dataset)
app_config = _load_config(config)
if target not in TARGETS:
raise typer.BadParameter(
f"Unknown target {target!r}. Choose from: {', '.join(TARGETS)}"
)
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(
@ -480,6 +568,8 @@ def run(
vacuum_interval=vacuum_interval,
multimodal_only=multimodal_only,
judge_model=judge_model_config,
target=target_value,
skill_model=skill_model_config,
)
)

92
evaluations/evaluations/skill_runner.py Normal file
View file

@ -0,0 +1,92 @@
from collections.abc import Callable
from dataclasses import dataclass, field
from pathlib import Path
from typing import Protocol, cast
from pydantic_ai.models import Model
from haiku.rag.agents.research.models import Citation
from haiku.rag.config.models import AppConfig
from haiku.rag.store.models.chunk import SearchResult
from haiku.skills import run_skill
from haiku.skills.models import Skill
SkillFactory = Callable[..., Skill]
class _RagLikeState(Protocol):
document_filter: str | None
citation_index: dict[str, Citation]
citations: list[str]
searches: dict[str, list[SearchResult]]
@dataclass
class SkillRunResult:
answer: str
cited_uris: list[str] = field(default_factory=list)
cited_chunk_ids: list[str] = field(default_factory=list)
searched_uris: list[str] = field(default_factory=list)
n_searches: int = 0
async def run_skill_question(
skill_factory: SkillFactory,
db_path: Path,
config: AppConfig,
question: str,
skill_model: str | Model,
document_filter: str | None = None,
request_limit: int | None = None,
) -> SkillRunResult:
"""Run a single question through a skill and return answer + retrieval data.
Builds the skill via ``skill_factory(db_path=..., config=...)`` and
invokes it with a fresh state instance derived from
``skill.state_type``. After the run, citations and searched documents
are extracted from the state for downstream eval scoring.
The skill must produce a state with RAG-skill-shaped fields (citation
index, searches, optional document filter) — i.e. ``RAGState`` or
``AnalysisState`` from ``haiku.rag.skills``.
"""
skill = skill_factory(db_path=db_path, config=config)
if request_limit is not None:
skill.request_limit = request_limit
if skill.state_type is None:
raise ValueError(f"Skill {skill.metadata.name!r} has no state_type")
state = skill.state_type()
typed = cast(_RagLikeState, state)
if document_filter is not None:
typed.document_filter = document_filter
answer, _, _ = await run_skill(skill_model, skill, question, state=state)
cited_chunk_ids: list[str] = list(typed.citations)
seen_cited: set[str] = set()
cited_uris: list[str] = []
for chunk_id in cited_chunk_ids:
citation = typed.citation_index.get(chunk_id)
if citation is None:
continue
if citation.document_uri not in seen_cited:
seen_cited.add(citation.document_uri)
cited_uris.append(citation.document_uri)
seen_searched: set[str] = set()
searched_uris: list[str] = []
for results in typed.searches.values():
for result in results:
uri = result.document_uri
if uri and uri not in seen_searched:
seen_searched.add(uri)
searched_uris.append(uri)
return SkillRunResult(
answer=answer,
cited_uris=cited_uris,
cited_chunk_ids=cited_chunk_ids,
searched_uris=searched_uris,
n_searches=len(typed.searches),
)

130
evaluations/tests/test_benchmark.py
View file

@ -192,3 +192,133 @@ class TestEvaluateDatasetJudgeModel:
mock_qa.assert_called_once()
assert mock_qa.call_args[1]["judge_model"] is custom_judge
class TestExperimentMetadataTargets:
def test_default_target_is_qa(self) -> None:
result = build_experiment_metadata(
dataset_key="test", test_cases=1, config=AppConfig()
)
assert result["target"] == "qa"
assert "skill_provider" not in result
assert "skill_model" not in result
def test_skill_target_includes_skill_config(self) -> None:
skill = ModelConfig(provider="ollama", name="gpt-oss-large", temperature=0.2)
result = build_experiment_metadata(
dataset_key="test",
test_cases=1,
config=AppConfig(),
target="rag-skill",
skill_config=skill,
)
assert result["target"] == "rag-skill"
assert result["skill_provider"] == "ollama"
assert result["skill_model"] == "gpt-oss-large"
assert result["skill_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_skill_model(self) -> None:
skill = 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-skill",
skill_model=skill,
)
mock_qa.assert_called_once()
assert mock_qa.call_args[1]["target"] == "rag-skill"
assert mock_qa.call_args[1]["skill_model"] is skill
@pytest.mark.asyncio
async def test_default_target_is_qa(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"] == "qa"
assert mock_qa.call_args[1]["skill_model"] is None
class TestRunQaBenchmarkSkillTarget:
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_skill_target_uses_run_skill_question(
self, tmp_path: Path
) -> None:
from evaluations.skill_runner import SkillRunResult
skill_run = AsyncMock(return_value=SkillRunResult(answer="from skill"))
with (
patch("evaluations.benchmark.get_model") as mock_get_model,
patch(
"evaluations.benchmark.run_skill_question", new=skill_run
) as mock_run_skill,
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-skill",
)
# When target is rag-skill, HaikuRAG context manager is NOT entered
# (the skill manages its own client via lifespan).
mock_haiku.assert_not_called()
# skill model defaults to qa.model when not provided
skill_call = mock_get_model.call_args_list[-1]
assert skill_call[0][0] == AppConfig().qa.model
assert mock_run_skill is skill_run
@pytest.mark.asyncio
async def test_analysis_skill_target_resolves_factory(self, tmp_path: Path) -> None:
from evaluations.benchmark import _skill_factory_for_target
from haiku.rag.skills.analysis import create_skill as analysis_factory
from haiku.rag.skills.rag import create_skill as rag_factory
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]

313
evaluations/tests/test_skill_runner.py Normal file
View file

@ -0,0 +1,313 @@
import random
from pathlib import Path
from typing import Any
import pytest
from pydantic_ai.models.test import TestModel
from evaluations.skill_runner import SkillRunResult, run_skill_question
from haiku.rag.agents.research.models import Citation
from haiku.rag.client import HaikuRAG
from haiku.rag.config.models import AppConfig
from haiku.rag.embeddings import EmbedderWrapper
from haiku.rag.skills.analysis import (
AnalysisState,
create_skill as create_analysis_skill,
)
from haiku.rag.skills.rag import RAGState, create_skill as create_rag_skill
from haiku.rag.store.models.chunk import SearchResult
VECTOR_DIM = 2560
@pytest.fixture(autouse=True)
def mock_embedder(monkeypatch: pytest.MonkeyPatch):
"""Deterministic embeddings so search is reproducible."""
async def fake_embed_query(self, text):
random.seed(hash(text) % (2**32))
return [random.random() for _ in range(VECTOR_DIM)]
async def fake_embed_documents(self, texts):
result = []
for t in texts:
random.seed(hash(t) % (2**32))
result.append([random.random() for _ in range(VECTOR_DIM)])
return result
monkeypatch.setattr(EmbedderWrapper, "embed_query", fake_embed_query)
monkeypatch.setattr(EmbedderWrapper, "embed_documents", fake_embed_documents)
@pytest.fixture
def app_config():
return AppConfig(environment="skill-runner-test")
@pytest.fixture
async def rag_db(tmp_path: Path):
"""A small two-document database."""
db_path = tmp_path / "test.lancedb"
async with HaikuRAG(db_path, create=True) as rag:
await rag.create_document(
"Artificial intelligence is transforming healthcare and finance.",
title="AI Overview",
uri="test://ai",
)
await rag.create_document(
"Machine learning includes supervised, unsupervised, and reinforcement.",
title="ML Basics",
uri="test://ml",
)
return db_path
class TestRunSkillQuestionMocked:
"""Verify the runner reads state correctly without going through a real skill loop."""
async def test_extracts_cited_and_searched_uris(
self, monkeypatch: pytest.MonkeyPatch, app_config: AppConfig, rag_db: Path
) -> None:
async def fake_run_skill(
model: Any,
skill: Any,
request: str,
state: Any = None,
event_sink: Any = None,
) -> tuple[str, list[Any], list[Any]]:
state.citation_index["c1"] = Citation(
chunk_id="c1",
document_id="d1",
document_uri="test://doc-a",
document_title="A",
content="alpha",
)
state.citation_index["c2"] = Citation(
chunk_id="c2",
document_id="d2",
document_uri="test://doc-b",
document_title="B",
content="beta",
)
state.citations = ["c1", "c2"]
state.searches["q1"] = [
SearchResult(content="x", score=0.9, document_uri="test://doc-a"),
SearchResult(content="y", score=0.8, document_uri="test://doc-c"),
]
state.searches["q2"] = [
SearchResult(content="z", score=0.7, document_uri="test://doc-a"),
]
return "answer", [], []
monkeypatch.setattr("evaluations.skill_runner.run_skill", fake_run_skill)
result = await run_skill_question(
skill_factory=create_rag_skill,
db_path=rag_db,
config=app_config,
question="anything?",
skill_model=TestModel(),
)
assert isinstance(result, SkillRunResult)
assert result.answer == "answer"
assert result.cited_chunk_ids == ["c1", "c2"]
assert result.cited_uris == ["test://doc-a", "test://doc-b"]
assert result.searched_uris == ["test://doc-a", "test://doc-c"]
assert result.n_searches == 2
async def test_skips_chunks_missing_from_index(
self, monkeypatch: pytest.MonkeyPatch, app_config: AppConfig, rag_db: Path
) -> None:
async def fake_run_skill(
model: Any,
skill: Any,
request: str,
state: Any = None,
event_sink: Any = None,
) -> tuple[str, list[Any], list[Any]]:
state.citation_index["c1"] = Citation(
chunk_id="c1",
document_id="d1",
document_uri="test://doc-a",
content="a",
)
state.citations = ["c1", "missing"]
return "ok", [], []
monkeypatch.setattr("evaluations.skill_runner.run_skill", fake_run_skill)
result = await run_skill_question(
skill_factory=create_rag_skill,
db_path=rag_db,
config=app_config,
question="?",
skill_model=TestModel(),
)
assert result.cited_chunk_ids == ["c1", "missing"]
assert result.cited_uris == ["test://doc-a"]
async def test_document_filter_is_set_on_state(
self, monkeypatch: pytest.MonkeyPatch, app_config: AppConfig, rag_db: Path
) -> None:
captured: dict = {}
async def fake_run_skill(
model: Any,
skill: Any,
request: str,
state: Any = None,
event_sink: Any = None,
) -> tuple[str, list[Any], list[Any]]:
captured["filter"] = state.document_filter
captured["state_type"] = type(state)
return "ok", [], []
monkeypatch.setattr("evaluations.skill_runner.run_skill", fake_run_skill)
await run_skill_question(
skill_factory=create_rag_skill,
db_path=rag_db,
config=app_config,
question="?",
skill_model=TestModel(),
document_filter="uri = 'test://ai'",
)
assert captured["filter"] == "uri = 'test://ai'"
assert captured["state_type"] is RAGState
async def test_request_limit_override(
self, monkeypatch: pytest.MonkeyPatch, app_config: AppConfig, rag_db: Path
) -> None:
captured: dict = {}
async def fake_run_skill(
model: Any,
skill: Any,
request: str,
state: Any = None,
event_sink: Any = None,
) -> tuple[str, list[Any], list[Any]]:
captured["request_limit"] = skill.request_limit
return "ok", [], []
monkeypatch.setattr("evaluations.skill_runner.run_skill", fake_run_skill)
await run_skill_question(
skill_factory=create_rag_skill,
db_path=rag_db,
config=app_config,
question="?",
skill_model=TestModel(),
request_limit=42,
)
assert captured["request_limit"] == 42
async def test_request_limit_unset_leaves_skill_default(
self, monkeypatch: pytest.MonkeyPatch, app_config: AppConfig, rag_db: Path
) -> None:
captured: dict = {}
async def fake_run_skill(
model: Any,
skill: Any,
request: str,
state: Any = None,
event_sink: Any = None,
) -> tuple[str, list[Any], list[Any]]:
captured["request_limit"] = skill.request_limit
return "ok", [], []
monkeypatch.setattr("evaluations.skill_runner.run_skill", fake_run_skill)
await run_skill_question(
skill_factory=create_rag_skill,
db_path=rag_db,
config=app_config,
question="?",
skill_model=TestModel(),
)
assert captured["request_limit"] is None
async def test_analysis_skill_uses_analysis_state(
self, monkeypatch: pytest.MonkeyPatch, app_config: AppConfig, rag_db: Path
) -> None:
captured: dict = {}
async def fake_run_skill(
model: Any,
skill: Any,
request: str,
state: Any = None,
event_sink: Any = None,
) -> tuple[str, list[Any], list[Any]]:
captured["state_type"] = type(state)
return "ok", [], []
monkeypatch.setattr("evaluations.skill_runner.run_skill", fake_run_skill)
await run_skill_question(
skill_factory=create_analysis_skill,
db_path=rag_db,
config=app_config,
question="?",
skill_model=TestModel(),
)
assert captured["state_type"] is AnalysisState
async def test_raises_when_skill_has_no_state_type(
self, app_config: AppConfig, rag_db: Path
) -> None:
from haiku.skills.models import Skill, SkillMetadata, SkillSource
def factory(*, db_path, config) -> Skill:
return Skill(
metadata=SkillMetadata(name="bare", description="No state."),
source=SkillSource.ENTRYPOINT,
instructions="Do nothing.",
)
with pytest.raises(ValueError, match="no state_type"):
await run_skill_question(
skill_factory=factory,
db_path=rag_db,
config=app_config,
question="?",
skill_model=TestModel(),
)
class TestRunSkillQuestionEndToEnd:
"""Real skill loop against a real LanceDB. Verifies the wiring beyond mocks."""
async def test_rag_skill_runs_against_real_db(
self,
allow_model_requests: None,
app_config: AppConfig,
rag_db: Path,
) -> None:
result = await run_skill_question(
skill_factory=create_rag_skill,
db_path=rag_db,
config=app_config,
question="What is machine learning?",
skill_model=TestModel(),
)
assert isinstance(result, SkillRunResult)
assert result.answer
assert result.n_searches >= 1
assert all(uri.startswith("test://") for uri in result.searched_uris)
@pytest.fixture
def allow_model_requests():
import pydantic_ai.models
with pydantic_ai.models.override_allow_model_requests(True):
yield