pv-agent/tests/test_query_planner.py

"""Tests: Query-Planer (Stufe 1) — Heuristik-Gate, JSON-Parsing, Fallback,
Multi-Query-Fusion und Integration in answer_question."""
import json

import pytest

from agent.config import Config
from agent.generate import answer_question
from agent.query_planner import SubQuery, parse_plan, plan_queries, should_plan
from tests.conftest import FakeOllama


def test_should_plan_gate():
    # Jahreszahl -> temporale Frage -> planen
    assert should_plan("Wie hoch war der Mindestlohn für Friseure im Jahr 2024?")
    # Vergleich
    assert should_plan("Was ist der Unterschied zwischen Abfertigung neu und alt?")
    # lang / mehrfach
    assert should_plan(
        "Wie wird die Überstundenpauschale behandelt und wie wirkt sie sich "
        "auf die Sozialversicherung und die Lohnsteuer aus?"
    )
    # einfach -> Single-Shot
    assert not should_plan("Wie viele Werktage Urlaub stehen Arbeitnehmern zu?")
    assert not should_plan("Was ist Altersteilzeit?")
    assert not should_plan("Wie hoch ist der KV-Mindestlohn im Friseurgewerbe?")


def test_parse_plan_valid_and_fallback():
    subs, qtype = parse_plan(
        'Vorab: {"type": "specific", "queries": [{"text": "mindestlohn friseur", "stand_year": "2024"}, '
        '{"text": "lohnberechnung friseur", "stand_year": null}]}',
        original="Originalfrage?",
    )
    assert [s.text for s in subs] == ["mindestlohn friseur", "lohnberechnung friseur"]
    assert subs[0].stand_year == "2024"
    assert subs[1].stand_year is None
    assert qtype == "specific"

    # Fallbacks: kaputtes JSON, leeres Array, leere Texte
    for raw in ("kein json", '{"queries": []}', '{"queries": [{"text": ""}]}'):
        subs, qtype = parse_plan(raw, original="Originalfrage?")
        assert len(subs) == 1 and subs[0].text == "Originalfrage?"
        assert qtype == "specific"

    # Ungueltiges Jahr -> None erzwingen, Text bleibt; unbekannter scope -> None
    subs, _ = parse_plan(
        '{"queries": [{"text": "x", "stand_year": "98", "scope": "xyz"}]}', original="orig"
    )
    assert subs == [SubQuery(text="x", stand_year=None, scope=None)]


def test_parse_plan_type_survey():
    subs, qtype = parse_plan(
        '{"type": "survey", "queries": [{"text": "wiku personal aktuell 2026 neuerungen", "scope": null}]}',
        original="orig",
    )
    assert qtype == "survey"
    assert subs[0].scope is None


def test_parse_plan_caps_at_three_queries():
    raw = json.dumps(
        {"queries": [{"text": f"q{i}"} for i in range(5)]}
    )
    subs, _ = parse_plan(raw, original="orig")
    assert len(subs) == 3


def test_plan_queries_simple_question_no_llm_call():
    client = FakeOllama(answers=[])  # darf nicht aufgerufen werden
    cfg = Config(planner_enabled=True)
    subs, planned, qtype = plan_queries("Was ist Altersteilzeit?", client, cfg)
    assert planned is False and qtype == "specific"
    assert subs == [SubQuery(text="Was ist Altersteilzeit?")]
    assert client.calls == 0


def test_plan_queries_uses_planner_and_falls_back_on_error():
    cfg = Config(planner_enabled=True)
    client = FakeOllama(
        answers=['{"type": "specific", "queries": [{"text": "atz lohnausgleich"}, {"text": "atz altersteilzeitgeld", "stand_year": null}]}']
    )
    subs, planned, qtype = plan_queries(
        "Wie funktioniert der Lohnausgleich bei Altersteilzeit und was ersetzt das AMS?", client, cfg
    )
    assert planned is True and qtype == "specific" and len(subs) == 2
    assert subs[0].text == "atz lohnausgleich"

    # Fehler -> Originalfrage, geplant False
    failing = FakeOllama(answers=[])
    failing.chat = lambda *a, **k: (_ for _ in ()).throw(RuntimeError("offline"))
    subs, planned, qtype = plan_queries("Wie funktioniert der Lohnausgleich 2026?", failing, cfg)
    assert planned is False and qtype == "specific"
    assert subs[0].text.startswith("Wie funktioniert")


def test_answer_question_planner_integration(mini_index):
    """Planer an: 1. Chat = Planer-JSON, 2. Chat = Antwort. Beide Sub-Queries
    fusionieren in einer Retrieved-Menge (Zitier-Validierung ueber die Union)."""
    cfg = Config(
        kb_dir=mini_index.kb_dir,
        db_path=mini_index.db_path,
        embed_off=True,
        planner_enabled=True,
        planner_max_queries=2,
    )
    client = FakeOllama(answers=[
        json.dumps({"type": "specific", "queries": [
            {"text": "altersteilzeit lohnausgleich", "stand_year": None},
            {"text": "urlaubsanspruch", "stand_year": None},
        ]}),
        "Antwort mit Beleg [lb-min-01] und [lb-min-02].",
    ])
    result = answer_question(
        "Wie funktioniert der Lohnausgleich bei Altersteilzeit und wie viel Urlaub bleibt?",
        cfg, client=client,
    )
    assert client.calls == 2  # Planer + Antwort
    assert result["verified"] is True
    assert result["refused"] is False
    assert set(result["citations"]) == {"lb-min-01", "lb-min-02"}
    assert len(result["planned_queries"]) == 2


def test_answer_question_survey_map_reduce(mini_index):
    """Survey-Frage: 1. Planer (type=survey), 2. Map-Destillat,
    3. Reduce-Antwort. Zitier-Validierung weiterhin gegen die Union."""
    cfg = Config(
        kb_dir=mini_index.kb_dir,
        db_path=mini_index.db_path,
        embed_off=True,
        planner_enabled=True,
        survey_blocks=8,
    )
    client = FakeOllama(answers=[
        json.dumps({"type": "survey", "queries": [
            {"text": "wiku personal aktuell 2026", "stand_year": None},
        ]}),
        # Map-Ausgabe: Destillat je Block mit KB-ID
        "[lb-min-01] Altersteilzeit: Lohnausgleich + ATZ-Geld.\n"
        "[lb-min-02] Urlaub: 5 Wochen je Dienstjahr.",
        # Reduce-Antwort
        "Neuerungen: ATZ-Lohnausgleich [lb-min-01]; Urlaub 5 Wochen [lb-min-02].",
    ])
    result = answer_question(
        "Welche Neuerungen behandeln die Wissensbasis 2026?",
        cfg, client=client,
    )
    assert client.calls == 3  # Planer + Map + Reduce
    assert result["verified"] is True
    assert result["refused"] is False
    assert set(result["citations"]) == {"lb-min-01", "lb-min-02"}


def test_search_multi_fuses_across_queries(mini_index):
    """Multi-Query-Retrieval: Sub-Queries summieren Beitraege; Treffer aus
    beiden Themen erscheinen im Kontext."""
    from agent.retrieve import Retriever

    cfg = Config(kb_dir=mini_index.kb_dir, db_path=mini_index.db_path, embed_off=True)
    r = Retriever(cfg)
    subs = [SubQuery(text="Altersteilzeit Lohnausgleich"),
            SubQuery(text="Urlaubsanspruch")]
    results = r.search_multi(subs, n_entries=8)
    ids = {res.entry_id for res in results}
    assert {"lb-min-01", "lb-min-02"} <= ids  # beide Themen vertreten
    r.close()