pv-agent/tests/test_generate.py

"""Tests: Grounding — Post-Validierung, Verweigerungspflicht, Regenerierung.

Der kritische Teil der Pipeline: keine Antwort mit ungültigen Zitaten
verlässt answer_question.
"""
import pytest

from agent.retrieve import ChunkResult
from agent.generate import (
    REFUSAL_MESSAGE,
    UNCERTAIN_MESSAGE,
    answer_question,
    build_user_content,
    looks_like_refusal,
    strip_think,
    trim_results,
    validate_answer,
)


class TestValidateAnswer:
    def test_valid_citation_passes(self):
        assert validate_answer("ATZ ist Teilzeit [lb-min-01].", ["lb-min-01"]) == []

    def test_unknown_id_is_violation(self):
        v = validate_answer("ATZ ist xyz [lb-atz-99].", ["lb-min-01"])
        assert any("lb-atz-99" in x for x in v)

    def test_no_citation_is_violation(self):
        v = validate_answer("ATZ ist eine Teilzeit.", ["lb-min-01"])
        assert any("keine KB-ID" in x for x in v)

    def test_refusal_without_citation_is_ok(self):
        assert validate_answer(REFUSAL_MESSAGE, ["lb-min-01"]) == []

    def test_bare_id_mention_is_detected(self):
        assert validate_answer("Siehe lb-min-02 für Details.", ["lb-min-02"]) == []


class TestRefusalDetection:
    def test_refusal_phrase(self):
        assert looks_like_refusal("Dazu enthält die Wissensbasis keine Aussage.")

    def test_refusal_phrase_with_typos_folds(self):
        assert looks_like_refusal("Dazu enthält die Wissensbasis keine Aussage!")

    def test_normal_answer_is_no_refusal(self):
        assert not looks_like_refusal("Der Anspruch besteht [lb-min-01].")

    def test_strip_think_removes_block(self):
        open_tag = "<" + "think" + ">"
        close_tag = "</" + "think" + ">"
        text = open_tag + "Reasoning here" + close_tag + "Antwort [lb-min-01]."
        out = strip_think(text)
        assert "Reasoning" not in out
        assert out.strip().startswith("Antwort [lb-min-01].")


class TestUserContent:
    def test_blocks_contain_metadata_header(self, mini_index):
        from agent.retrieve import Retriever

        r = Retriever(mini_index)
        try:
            results = r.search("Altersteilzeit Lohnausgleich")
            content = build_user_content("Was ist ATZ?", results)
            assert "Block 1 — [lb-min-01]" in content
            assert "Stand: 2026-01" in content
            assert "Frage: Was ist ATZ?" in content
        finally:
            r.close()


class TestAnswerQuestion:
    def test_happy_path_verified(self, mini_index, fake_ollama):
        client = fake_ollama(
            answers=["Altersteilzeit ist eine Teilzeit mit Lohnausgleich "
                     "[lb-min-01]. (Stand 2026-01)"]
        )
        result = answer_question(
            "Was ist Altersteilzeit?", mini_index, client=client
        )
        assert result["verified"] is True
        assert result["refused"] is False
        assert result["citations"] == ["lb-min-01"]
        assert result["sources"][0]["id"] == "lb-min-01"
        assert result["sources"][0]["stand"] == "2026-01"
        assert client.calls == 1

    def test_hallucinated_id_regenerates_then_refuses(self, mini_index, fake_ollama):
        client = fake_ollama(answers=[
            "ATZ gilt ab 60. Lebensjahr [lb-atz-99].",
            "ATZ gilt ab 60. Lebensjahr, siehe [lb-atz-99].",
        ])
        result = answer_question(
            "Was ist Altersteilzeit?", mini_index, client=client
        )
        assert result["refused"] is True
        assert result["verified"] is False
        assert result["answer"] == UNCERTAIN_MESSAGE
        assert result["regenerations"] == 1
        assert result["citations"] == []
        assert "draft" in result and "lb-atz-99" in result["draft"]

    def test_regeneration_can_recover(self, mini_index, fake_ollama):
        client = fake_ollama(answers=[
            "ATZ gilt ab 60 [lb-atz-99].",
            "ATZ ist Teilzeit mit Lohnausgleich [lb-min-01].",
        ])
        result = answer_question(
            "Was ist Altersteilzeit?", mini_index, client=client
        )
        assert result["verified"] is True
        assert result["regenerations"] == 1
        assert result["citations"] == ["lb-min-01"]

    def test_empty_retrieval_refuses_deterministically(self, mini_index, fake_ollama):
        client = fake_ollama(answers=["sollte nie aufgerufen werden"])
        result = answer_question(
            "Wie hoch ist der Wechselkurs von Bermuda-Dollar?", mini_index,
            client=client,
        )
        assert result["refused"] is True
        assert result["answer"] == REFUSAL_MESSAGE
        assert result["verified"] is True
        assert client.calls == 0  # kein LLM-Call bei leerem Retrieval

    def test_model_refusal_is_kept(self, mini_index, fake_ollama):
        client = fake_ollama(answers=[
            f"Zu dieser Frage: {REFUSAL_MESSAGE}"
        ])
        result = answer_question(
            "Was ist Altersteilzeit?", mini_index, client=client
        )
        assert result["refused"] is True
        assert result["verified"] is True  # Regel-4-konforme Verweigerung
        assert client.calls == 1

    def test_top_k_limits_context(self, mini_index, fake_ollama):
        client = fake_ollama(answers=["Teilzeit [lb-min-01]."])
        result = answer_question(
            "Altersteilzeit Urlaub Lohnverrechnung", mini_index,
            client=client, top_k=1,
        )
        main = [s for s in result["sources"]]
        assert result["n_context"] >= 1
        # Haupt-Blöcke auf top_k begrenzt; cross_ref-Erweiterungen dürfen dazu
        assert len([s for s in main]) <= result["n_context"]

def test_trim_results_drops_tail_under_budget():
    """Prompt-Budget: Tail-Blöcke (niedrig gerankt) ganz weglassen, nicht truncieren."""
    def blk(n, chars):
        return ChunkResult(
            chunk_id=n, entry_id=f"lb-min-{n:02d}", section="Zusammenfassung",
            text="x" * chars, title=f"Titel {n}", stand="2026-01", work="W",
            chapter="C", topic="t", tags=[], legal_bases=[], cross_refs=[],
        )
    blocks = [blk(i, 30_000) for i in range(1, 6)]  # 5 x ~30 KB = 150 KB
    trimmed = trim_results(blocks, 90_000)
    assert 6 >= len(trimmed) >= 6 or len(trimmed) == 5  # 5 Bloecke, Mindesthoehe 6 greift nicht
    # Groesserer Fall: 14 Bloecke, Budget schneidet hinten weg
    blocks = [blk(i, 8_000) for i in range(1, 15)]  # ~112 KB
    trimmed = trim_results(blocks, 90_000)
    assert 6 <= len(trimmed) < 14
    # Erste Bloecke bleiben (best gerankt)
    assert trimmed[0].entry_id == "lb-min-01"
    assert sum(len(b.text) for b in trimmed) + len(trimmed) * 64 <= 90_000 + 8_000
    # Unbegrenzt: Original unveraendert
    assert trim_results(blocks, None) is blocks