tests/conftest.py

from __future__ import annotations

from pathlib import Path

import pytest

from agentic_search.budget import BudgetManager
from agentic_search.config import load_config
from agentic_search.data.schemas import Document
from agentic_search.retrieval import DenseRetriever, DocumentStore, Reranker, SparseRetriever
from agentic_search.sdk.runtime import SearchRuntime
from agentic_search.tracing.schemas import RunMetadata
from agentic_search.tracing.writer import TraceWriter


@pytest.fixture
def documents() -> list[Document]:
    return [
        Document(
            doc_id="ada",
            title="Ada Lovelace",
            text="Ada Lovelace wrote notes about Charles Babbage's Analytical Engine.",
            source="fixture",
            dataset="fixture",
        ),
        Document(
            doc_id="babbage",
            title="Charles Babbage",
            text="Charles Babbage designed the Analytical Engine in London.",
            source="fixture",
            dataset="fixture",
        ),
        Document(
            doc_id="other",
            title="Other",
            text="This unrelated passage discusses a different machine.",
            source="fixture",
            dataset="fixture",
        ),
    ]


@pytest.fixture
def runtime(tmp_path: Path, documents: list[Document]) -> SearchRuntime:
    config = load_config()
    metadata = RunMetadata(
        run_id="test-run",
        example_id="test-example",
        system="test",
        config_hash="config",
        corpus_hash="corpus",
        model_name="test-model",
        model_endpoints=[],
        seed=42,
    )
    return SearchRuntime(
        store=DocumentStore(documents),
        sparse=SparseRetriever(documents),
        dense=DenseRetriever(documents, backend="hashing", dimensions=64),
        reranker=Reranker(backend="lexical", model_name="unused", device="cpu", batch_size=2, max_length=128),
        budget=BudgetManager(config.budget),
        trace=TraceWriter(tmp_path, metadata),
    )