pv-agent/agent/query_planner.py

"""Query-Planer für komplexe Fragen (Stufe 1, M6).

Ein kleiner LLM-Call zerlegt komplexe Fragen in 1-3 unabhängige Sub-Queries
(+ optional Geltungsjahr). Ein Heuristik-Gate entscheidet, ob der Planer-Call
überhaupt läuft — einfache Fragen bleiben deterministischer Single-Shot
(keine Zusatz-Latenz). Parse-/Call-Fehler fallen auf die Original-Frage
zurück; die Grounding-Regeln werden dadurch nie berührt.
"""
from __future__ import annotations

import json
import re
from dataclasses import dataclass

YEAR_RE = re.compile(r"\b(?:19|20)\d{2}\b")
JSON_RE = re.compile(r"\{.*\}", re.DOTALL)

AGGREGATION_RE = re.compile(
    r"neuerungen|übersicht|zusammenfassung|alle\s|übersicht", re.I
)
COMPARISON_RE = re.compile(r"unterschied|vergleic|\bbzw\.|\bsowie\b", re.I)

PLANNER_PROMPT = """Du planst Suchanfragen für eine Wissensbasis zur österreichischen
Personalverrechnung (kuratierte Briefings, Kollektivverträge je Branche und
Geltungsjahr, Gesetzesparagraphen).

Frage: "{question}"

Zerlege die Frage in 1-3 unabhängige Suchanfragen, die zusammen die Frage
beantworten. Regeln:
- Besteht die Frage aus einem Suchthema, gib genau eine Suchanfrage zurück
  (dann die Frage unverändert, leicht stichwortartig gekürzt).
- Suchbegriffe statt Sätze (ohne Frageformulierung, ohne "Kollektivvertrag"-
  Wortballast).
- Fragt die Frage nach einem bestimmten Jahr (Geltung/Stand), setze
  "stand_year" auf dieses vierstellige Jahr, sonst null.
- Setze "scope" je Suchanfrage: "gesetz", wenn nach der gesetzlichen/
  allgemeinen Grundlage gefragt ist (nur Gesetze und Fachbriefings, ohne
  Branchen-Kollektivverträge); "kv", wenn ausdrücklich nach kollektivvertrag-
  lichen Branchenregelungen gefragt ist; null für alles anderes.
- Setze "type" auf "survey", wenn die Frage eine Übersicht über viele
  Dokumente verlangt (Neuerungen, Entwicklungen, alle …); sonst "specific".
- Antworte ausschließlich mit JSON, ohne Erklärung:
{{"type": "specific", "queries": [{{"text": "...", "stand_year": null, "scope": null}}]}}"""


@dataclass
class SubQuery:
    text: str
    stand_year: str | None = None  # YYYY, wenn nach einem Geltungsjahr gefragt
    scope: str | None = None        # "gesetz" (ohne Branchen-KV) | "kv" | None


def should_plan(question: str) -> bool:
    """Heuristik-Gate: nur komplexe Fragen bekommen einen Planer-Call.

    Signale: Jahreszahl, Vergleichs-/Aggregationsmarker, langer Text,
    Mehrfach-Konjunktion. Einfache Fragen bleiben Single-Shot (Latenz)."""
    q = question.strip()
    if YEAR_RE.search(q):
        return True
    words = q.split()
    if len(words) >= 14:
        return True
    if COMPARISON_RE.search(q):
        return True
    if AGGREGATION_RE.search(q):
        return True
    if " und " in q.casefold() and len(words) >= 10:
        return True
    return False


def parse_plan(raw: str, original: str) -> tuple[list[SubQuery], str]:
    """Robustes JSON-Parsing; jeder Fehler → [Originalfrage]. Liefert
    (Sub-Queries, Fragetyp 'survey' | 'specific')."""
    try:
        match = JSON_RE.search(raw)
        if not match:
            raise ValueError("kein JSON-Objekt")
        data = json.loads(match.group(0))
        items = data.get("queries")
        if not isinstance(items, list) or not items:
            raise ValueError("leeres Plan-Array")
        subs: list[SubQuery] = []
        for item in items[:3]:
            text = str(item.get("text", "")).strip()
            if not text:
                raise ValueError("leere Sub-Query")
            year = item.get("stand_year")
            year = str(year) if year and re.fullmatch(r"20\d{2}", str(year)) else None
            scope = item.get("scope")
            scope = scope if scope in ("gesetz", "kv") else None
            subs.append(SubQuery(text=text, stand_year=year, scope=scope))
        qtype = data.get("type") if data.get("type") in ("survey", "specific") else "specific"
        return subs, qtype
    except (ValueError, TypeError, KeyError, json.JSONDecodeError):
        return [SubQuery(text=original)], "specific"
        return [SubQuery(text=original)]


def plan_queries(
    question: str,
    client,
    cfg,
) -> tuple[list[SubQuery], bool, str]:
    """Liefert (Sub-Queries, geplant?, Fragetyp) — Call-/Parse-Fehler →
    Original als Einzel-Query, Typ 'specific'. Der Planer-Call ist klein
    (Frage ohne Kontext, kurzes num_predict); Temperature 0."""
    if not should_plan(question):
        return [SubQuery(text=question)], False, "specific"
    prompt = PLANNER_PROMPT.format(question=question.strip())
    try:
        raw = client.chat(
            cfg.planner_model or cfg.answer_model,
            [
                {"role": "system", "content": "Du antwortest ausschließlich mit JSON."},
                {"role": "user", "content": prompt},
            ],
            temperature=0.0,
            num_ctx=cfg.num_ctx,
            num_predict=cfg.planner_num_predict,
            think=False,
        )
    except Exception:
        return [SubQuery(text=question)], False, "specific"
    subs, qtype = parse_plan(raw, question)
    return subs, True, qtype