dp-zp-agent/tools/crzp/answering/context.py

"""Follow-up questions for /crzp: a signed, stateless conversation token.

A /crzp answer returns conversation_context.token holding the understood query
(who, role, title, topic, years, type, school, intent) and the permalink ids of
the shown works. The host sends it back with the next message of the same chat
("a v roku 2023?", "a z akých škôl boli tí študenti?", "o čom je tá druhá?").
The backend keeps no chat state: the token is HMAC-signed with a key derived
from SEARCH_API_KEY (domain-separated from the ZP Wiki tokens) and expires in
24 hours. Nothing in it is shown to the user; the answer header always states
the conditions actually applied.
"""
from __future__ import annotations

import base64
import hashlib
import hmac
import json
import os
import re
import time

from tools.crzp.answering.understand import ROLE_CUES, CrzpQuery, _intent
from tools.crzp.source.parse import fold

SCHEMA = "crzp-1"
LIFETIME = 24 * 60 * 60
MAX_TOKEN = 16000
MAX_SIDS = 30
CARRIED = {"intent", "people", "title", "topics", "surface", "phrases", "school_name", "years", "work_types", "school", "unit",
           "alternatives", "target_role", "from_model", "sids"}

CONTINUES = re.compile(r"^(?:a|aj|a co|co tak|a ako|a kto|a kolko|a ktore|a v|a z|a pre|a len|iba|len|teraz)\b")
LIST_CUE = re.compile(r"\b(?:ktor\w*|vypis\w*|zoznam\w*|ukaz\w*|ake prace|ake su|aku pracu|nazvy|list|which)\b")
ORDINALS = [(re.compile(r"\b(?:prv\w*|first)\b"), 0), (re.compile(r"\b(?:druh\w*|second)\b"), 1),
            (re.compile(r"\b(?:tret\w*|third)\b"), 2), (re.compile(r"\b(?:stvrt\w*|fourth)\b"), 3),
            (re.compile(r"\b(?:piat\w*|fifth)\b"), 4), (re.compile(r"\b(?:posledn\w*|last)\b"), -1)]
SINGLE = re.compile(r"\b(?:ta|tu|tej|tou|tato|tuto|tejto) (?:praca|pracu|prace|praci|pracou|bakalark\w*|diplomov\w*)\b"
                    r"|\bo com (?:je|bola|bolo)\b|\bjej (?:abstrakt|klucov|oponent|skolitel|vedu)")


class ContextError(ValueError):
    """invalid_context or expired_context (both are public, safe reasons)."""


def _key() -> bytes | None:
    secret = os.getenv("SEARCH_API_KEY", "")
    return hashlib.sha256(("zp-crzp-context-v1\0" + secret).encode()).digest() if secret else None


def _b64(value: bytes) -> str:
    return base64.urlsafe_b64encode(value).decode().rstrip("=")


def _unb64(value: str) -> bytes:
    return base64.b64decode(value + "=" * (-len(value) % 4), altchars=b"-_", validate=True)


def issue(question: str, query: CrzpQuery, shown: list[dict], options: list[dict] | None = None) -> dict | None:
    """conversation_context for the result, or None (no key configured, nothing to continue).
    options: the namesakes offered in a clarification ({"name", "school"})."""
    key = _key()
    if key is None or not (query.people or query.title or query.topics or query.school):
        return None
    rows = [r for r in shown if r.get("sid")][:MAX_SIDS]
    payload = {"schema": SCHEMA, "expires": int(time.time()) + LIFETIME, "question": question,
               "query": query.model_dump(mode="json", include=CARRIED - {"sids"}),
               "sids": [r["sid"] for r in rows], "options": (options or [])[:8]}
    body = _b64(json.dumps(payload, ensure_ascii=False, separators=(",", ":")).encode())
    token = body + "." + _b64(hmac.digest(key, body.encode(), "sha256"))
    if len(token) > MAX_TOKEN:
        return None
    return {"token": token, "work_options": [{"id": r["sid"], "label": r["title"]} for r in rows[:10]],
            "meeting_options": []}


def restore(token: str) -> dict:
    key = _key()
    try:
        if key is None or not isinstance(token, str) or not 1 <= len(token) <= MAX_TOKEN:
            raise ValueError
        body, signature = token.split(".")
        if not hmac.compare_digest(_unb64(signature), hmac.digest(key, body.encode(), "sha256")):
            raise ValueError
        value = json.loads(_unb64(body))
        if value.get("schema") != SCHEMA or not isinstance(value.get("query"), dict):
            raise ValueError
        if value["expires"] < time.time():
            raise ContextError("expired_context")
        value["query"] = CrzpQuery.model_validate(value["query"])
        return value
    except ContextError:
        raise
    except (ValueError, KeyError, TypeError, UnicodeError):
        raise ContextError("invalid_context") from None


def follow_up(question: str, new: CrzpQuery, previous: dict) -> tuple[CrzpQuery, str | None]:
    """(query to answer, clarification or None). A message with its own subject and no
    reference to the previous answer is answered on its own; otherwise the previous
    conditions are kept and the new ones (person, topic, year, type, school, intent) replace them."""
    low = fold(question)
    own_subject = bool(new.people or new.title or new.topics)
    if own_subject and not CONTINUES.search(low):
        return new, None  # "Koľko prác viedol Staš?" stands on its own even with a token
    old: CrzpQuery = previous["query"]
    merged = old.model_copy(deep=True)
    if new.people:
        merged.people = new.people
        role_named = any(re.search(cue, question.lower()) for cue in ROLE_CUES.values())
        for person, before in zip(merged.people, old.people):
            if not role_named:
                person.role = before.role  # "a čo Staš?" keeps the asked role (whatever the model guessed)
        merged.from_model = new.from_model
    if new.title:
        merged.title = new.title
    if new.topics:
        merged.topics, merged.surface, merged.alternatives = new.topics, new.surface, new.alternatives
        merged.phrases = new.phrases
    merged.years = new.years or old.years
    merged.work_types = new.work_types or old.work_types
    if new.school or new.unit:
        merged.school, merged.unit = new.school, new.unit
    explicit = _intent(question.lower())
    if explicit != "list" or LIST_CUE.search(low) or new.intent != "list" and not own_subject:
        merged.intent = new.intent if explicit == "list" else explicit
        merged.target_role = new.target_role
    if merged.people and merged.target_role == merged.people[0].role:
        merged.target_role = None
    if merged.intent == "relation" and merged.target_role is None:
        merged.intent = "list"
    merged.sids, merged.searched, merged.relaxed_author, merged.unverified_only = [], [], False, []
    sids = previous.get("sids") or []
    number = re.fullmatch(r"\s*(?:cislo\s+)?(\d)\s*[.)]?\s*", low)
    ordinal = int(number.group(1)) - 1 if number else next((index for cue, index in ORDINALS if cue.search(low)), None)
    options = previous.get("options") or []
    if options and merged.people and not new.people:  # the answer to "Ktorú osobu myslíš?"
        chosen = None
        if ordinal is not None and 0 <= ordinal < len(options):
            chosen = options[ordinal]
        elif new.school:
            from tools.crzp.answering.understand import SCHOOLS
            matching = [o for o in options if fold(SCHOOLS[new.school][0]) == fold(o["school"])]
            chosen = matching[0] if len(matching) == 1 else None
        if chosen:
            merged.people[0].name = merged.people[0].typed = chosen["name"]
            merged.school_name = chosen["school"]
            merged.from_model = True  # the exact CRZP name: no other name forms
        return merged, None
    if ordinal is not None and not own_subject:
        if not sids:
            return merged, "Predchádzajúca odpoveď nemala žiadnu prácu; uveď názov práce alebo osobu."
        if ordinal >= len(sids):
            return merged, f"Predchádzajúca odpoveď mala {len(sids)} {'práce' if len(sids) < 5 else 'prác'}; ktorú myslíš?"
        merged.sids = [sids[ordinal]]
    elif SINGLE.search(low) and not own_subject:
        if len(sids) != 1:
            return merged, ("Predchádzajúca odpoveď mala viac prác; ktorú myslíš? Napíš napr. „tá druhá“ "
                            "alebo názov práce.")
        merged.sids = list(sids)
    return merged, None