tests/integration/test_systems.py

from __future__ import annotations

from collections import deque

from agentic_search.llm.openai_compatible import LLMResponse
from agentic_search.systems import (
    AgenticAgent,
    FixedPipeline,
    SearchAsCodeAccuracyAgent,
    SearchAsCodeAgent,
    SearchAsCodeLatencyAgent,
)


class ScriptedClient:
    def __init__(self, responses: list[str]) -> None:
        self.responses = deque(responses)

    def chat(self, messages, *, budget, seed, max_tokens=None):
        budget.reserve_llm()
        budget.record_tokens(10, 5)
        return LLMResponse(
            content=self.responses.popleft(),
            model="scripted",
            input_tokens=10,
            output_tokens=5,
            latency_ms=1.0,
            endpoint="mock://model",
            cached=False,
        )

    @staticmethod
    def estimated_cost_usd(input_tokens: int, output_tokens: int) -> float:
        return 0.0


def test_fixed_pipeline_uses_shared_primitives(runtime) -> None:
    result = FixedPipeline(sparse_k=3, dense_k=3, merge_k=3, rerank_k=2).run(
        "Who wrote notes about the Analytical Engine?", runtime
    )

    assert result.status == "completed"
    assert result.final_ref is not None
    assert result.usage.retrieval_calls == 2
    assert result.usage.total_reranked_docs > 0


def test_sac_refines_then_stops(runtime) -> None:
    client = ScriptedClient(
        [
            "```python\ns = search_sparse('Analytical Engine', k=3)\nresult = rerank('Analytical Engine', s, k=2)\n```",
            "DONE",
        ]
    )
    result = SearchAsCodeAgent(client, max_turns=3).run("Who wrote notes about the Analytical Engine?", runtime)

    assert result.status == "completed"
    assert result.usage.llm_calls == 2
    assert result.final_ref is not None


def test_sac_keeps_last_valid_context_after_an_invalid_refinement(runtime) -> None:
    client = ScriptedClient(
        [
            "```python\nresult = search_sparse('Analytical Engine', k=3)\n```",
            "This is invalid because it is explanation text.",
            "DONE",
        ]
    )
    result = SearchAsCodeAgent(client, max_turns=3).run("Who wrote notes about the Analytical Engine?", runtime)

    assert result.status == "completed"
    assert result.final_ref is not None
    assert result.error is not None


def test_sac_latency_uses_one_compact_controller_turn(runtime) -> None:
    client = ScriptedClient(
        ["```python\ns = search_sparse('Analytical Engine', k=3)\nresult = rerank('Analytical Engine', s, k=2)\n```"]
    )
    result = SearchAsCodeLatencyAgent(client, max_tokens=384).run(
        "Who wrote notes about the Analytical Engine?", runtime
    )

    assert result.status == "completed"
    assert result.usage.llm_calls == 1
    assert result.final_ref is not None


def test_sac_latency_uses_deterministic_fallback_for_invalid_compact_code(runtime) -> None:
    client = ScriptedClient(
        [
            "```python\nresult = merge(dedup(merge([search_sparse('Ada', k=3), search_dense('Ada', k=3)])), k=3)\n```"
        ]
    )
    result = SearchAsCodeLatencyAgent(client, max_tokens=384).run("Who wrote notes?", runtime)

    assert result.status == "completed"
    assert result.usage.llm_calls == 1
    assert result.final_ref is not None
    assert result.error is not None


def test_sac_accuracy_verifies_then_applies_targeted_patch(runtime) -> None:
    client = ScriptedClient(
        [
            "```python\ns = search_sparse('Analytical Engine', k=3)\nresult = rerank('Analytical Engine', s, k=2)\n```",
            '{"sufficient":false,"missing_queries":["Ada Lovelace"]}',
            "```python\nresult = search_dense('Ada Lovelace', k=3)\n```",
        ]
    )
    result = SearchAsCodeAccuracyAgent(
        client,
        initial_max_tokens=1024,
        verifier_max_tokens=384,
        max_patches=1,
        context_chars=8000,
    ).run("Who wrote notes about the Analytical Engine?", runtime)

    assert result.status == "completed"
    assert result.usage.llm_calls == 3
    assert result.final_ref is not None


def test_sac_accuracy_recovers_patch_that_references_prior_result(runtime) -> None:
    client = ScriptedClient(
        [
            "```python\nresult = search_sparse('Analytical Engine', k=3)\n```",
            '{"sufficient":false,"missing_queries":["Ada Lovelace"]}',
            "```python\nnew = search_dense('Ada Lovelace', k=3)\nresult = merge([result, new], k=3)\n```",
        ]
    )
    result = SearchAsCodeAccuracyAgent(
        client,
        initial_max_tokens=1024,
        verifier_max_tokens=384,
        max_patches=1,
        context_chars=8000,
    ).run("Who wrote notes about the Analytical Engine?", runtime)

    assert result.status == "completed"
    assert result.final_ref is not None
    assert result.error is not None
    assert "deterministic verifier-query fallback" in result.error


def test_agentic_selects_tools_and_finishes(runtime) -> None:
    client = ScriptedClient(
        [
            '{"tool":"search_sparse","args":{"query":"Analytical Engine","k":3}}',
            '{"tool":"rerank","args":{"query":"Analytical Engine","ref":"r1","k":2}}',
            '{"tool":"finish","args":{"ref":"r2"}}',
        ]
    )
    result = AgenticAgent(client, max_steps=4).run("Who wrote notes about the Analytical Engine?", runtime)

    assert result.status == "completed"
    assert result.usage.llm_calls == 3
    assert result.final_ref is not None


def test_agentic_can_finish_multiple_named_results(runtime) -> None:
    client = ScriptedClient(
        [
            '{"tool":"search_sparse","args":{"query":"Ada Lovelace","k":3}}',
            '{"tool":"search_sparse","args":{"query":"Charles Babbage","k":3}}',
            '{"tool":"finish","args":{"ref":["r1","r2"]}}',
        ]
    )
    result = AgenticAgent(client, max_steps=4).run("Who wrote notes about the Analytical Engine?", runtime)

    assert result.status == "completed"
    assert result.final_ref is not None


def test_agentic_recovers_unknown_reference_on_next_step(runtime) -> None:
    client = ScriptedClient(
        [
            '{"tool":"search_sparse","args":{"query":"Ada Lovelace","k":3}}',
            '{"tool":"rerank","args":{"query":"Ada Lovelace","ref":"r99","k":2}}',
            '{"tool":"rerank","args":{"query":"Ada Lovelace","ref":"r1","k":2}}',
            '{"tool":"finish","args":{"ref":"r2"}}',
        ]
    )
    result = AgenticAgent(client, max_steps=4).run("Who wrote notes?", runtime)

    assert result.status == "completed"
    assert result.final_ref is not None
    assert result.error is not None
    assert "recovered invalid Agentic action" in result.error