ratatoskr/docs/diagnostics/descriptive_query_binding.py

#!/usr/bin/env python3
"""Descriptive-query subject-binding probe — the canonical fixture for Worldtree #393.

#393: an ATTRIBUTE/descriptive question ("the guy with the roid rage") reformulates to a
token query that matches MULTIPLE distinct canon subjects on a shared word ("rage"), so the
intended entity can be absent from top-k while topically-adjacent decoys rank above it. The
name-check (`names_subject`) can't help — the caller has no name to pass until the attribute
is resolved to an entity, which is the open problem. Downstream, the consumer sometimes binds
to the wrong co-retrieved subject and cross-contaminates details (a confident, fluent mis-bind
assembled from real-but-mismatched rows, not a hallucination).

Two regimes, one variable set:
  * RAW RANKING (persona-independent): drive the `mimir` librarian with the descriptive query
    and its variants; show the top-k rows flagged for the intended ENTITY vs the DECOY. The
    load-bearing fact is "intended entity absent from top-k."
  * CONSUMER: drive `ratatoskr:donut` N times on the descriptive question; classify each answer
    BINDS-ENTITY vs MIS-BINDS-DECOY vs OTHER and report the rate.

Two-regime finding (2026-08-07, personal instance): roid-rage mis-binding SURVIVES the #389 arc
(b172 -> v1.0.0b181); "dangerous crown" is mostly RESOLVED by the bge rerank + source annotations
(it resolves to essentially one entity, "the Enchanted Crown of the Sepsis Whore"). The gap is
attribute->entity resolution, upstream of names_subject by construction.

Self-contained: the only third-party dependency is httpx (`uv run --with httpx`). Config from env:
WORLDTREE_API_URL, WORLDTREE_API_KEY, RATATOSKR_END_USER_ID. No secrets stored here.

Usage:
  uv run --with httpx python docs/diagnostics/descriptive_query_binding.py
  uv run --with httpx python docs/diagnostics/descriptive_query_binding.py --runs 10
  uv run --with httpx python docs/diagnostics/descriptive_query_binding.py --case roid-rage
"""
from __future__ import annotations

import argparse
import json
import os
import re
from collections import Counter

import httpx

# (label, descriptive question, raw-ranking query variants, intended-entity regex, decoy regex|None)
CASES = [
    {
        "label": "roid-rage",
        "question": "Tell me about the guy with the roid rage.",
        "variants": ["the guy with the roid rage", "roid rage"],
        "entity": r"juicer",
        "decoy": r"\bJack\b",
    },
    {
        "label": "dangerous-crown",
        "question": "Tell me about that dangerous crown.",
        "variants": ["that dangerous crown", "dangerous crown"],
        "entity": r"sepsis|crown of the sepsis whore",
        "decoy": None,
    },
]


def _cfg() -> tuple[str, dict, str]:
    base = os.environ.get("WORLDTREE_API_URL", "http://10.250.50.152:8081")
    key = os.environ.get("WORLDTREE_API_KEY")
    if not key:
        raise SystemExit("WORLDTREE_API_KEY unset — source env.sh first.")
    end_user = os.environ.get("RATATOSKR_END_USER_ID", "ratatoskr-tui")
    return base, {"Authorization": f"Bearer {key}"}, end_user


def _session(base: str, headers: dict, agent_id: str, end_user: str) -> str:
    r = httpx.post(f"{base}/sessions", json={"agent_id": agent_id, "end_user_id": end_user},
                   headers=headers, timeout=30)
    r.raise_for_status()
    return r.json()["session_id"]


def _drive(base: str, headers: dict, sid: str, content: str) -> tuple[str | None, dict, str]:
    """POST a turn; return (tool_query, first tool_result dict, accumulated answer text)."""
    q, result, parts = None, {}, []
    with httpx.stream("POST", f"{base}/sessions/{sid}/messages", json={"content": content},
                      headers=headers, timeout=180) as r:
        for line in r.iter_lines():
            if not line.startswith("data: "):
                continue
            ev = json.loads(line[6:])
            t = ev.get("type")
            if t == "tool_start" and q is None:
                q = (ev.get("arguments") or {}).get("query")
            elif t == "tool_result" and not result:
                result = ev.get("result") if isinstance(ev.get("result"), dict) else {}
            elif t == "text":
                v = ev.get("text") or ev.get("content") or ev.get("delta")
                if isinstance(v, str):
                    parts.append(v)
            elif t == "done":
                break
    return q, result, "".join(parts)


def _flag(excerpt: str, entity: str, decoy: str | None) -> str:
    if re.search(entity, excerpt, re.I):
        return "ENTITY"
    if decoy and re.search(decoy, excerpt):
        return "DECOY "
    return "      "


def raw_ranking(base, headers, end_user, case) -> None:
    print("\n  [raw ranking] intended entity absent from top-k?")
    for query in case["variants"]:
        sid = _session(base, headers, "mimir", end_user)  # fresh session per query
        _, res, _ = _drive(base, headers, sid, f"Use search_library to find: {query}")
        rows = res.get("results", []) if isinstance(res, dict) else []
        entity_ranks = [i for i, r in enumerate(rows) if isinstance(r, dict)
                        and re.search(case["entity"], r.get("excerpt", ""), re.I)]
        top = entity_ranks[0] if entity_ranks else None
        print(f"    query={query!r:36} entity first appears at rank "
              f"{top if top is not None else 'ABSENT (not in top-k)'}")
        for i, r in enumerate(rows[:6]):
            if isinstance(r, dict):
                ex = (r.get("excerpt") or "").replace("\n", " ")
                print(f"      #{i} [{_flag(ex, case['entity'], case['decoy'])}] "
                      f"{r.get('score')}  {ex[:74]}")


def consumer(base, headers, end_user, case, runs) -> Counter:
    verdicts: Counter = Counter()
    print(f"\n  [consumer] Donut x{runs} on {case['question']!r}")
    for run in range(1, runs + 1):
        sid = _session(base, headers, "ratatoskr:donut", end_user)  # fresh session per run
        q, _, ans = _drive(base, headers, sid, case["question"])
        binds = bool(re.search(case["entity"], ans, re.I))
        mis = bool(case["decoy"]) and bool(re.search(case["decoy"], ans)) and not binds
        v = "BINDS-ENTITY" if binds else ("MIS-BINDS-DECOY" if mis else "OTHER")
        verdicts[v] += 1
        print(f"    run{run}: {v:16} q={q!r:34} :: {ans.strip()[:80]}")
    print(f"    >>> {case['label']}: {dict(verdicts)}")
    return verdicts


def main() -> None:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--runs", type=int, default=6, help="consumer repeats per case")
    ap.add_argument("--case", help="run only this case label (e.g. roid-rage)")
    ns = ap.parse_args()
    base, headers, end_user = _cfg()
    cases = [c for c in CASES if ns.case in (None, c["label"])]
    if not cases:
        raise SystemExit(f"no case matching {ns.case!r} (have: {[c['label'] for c in CASES]})")
    for case in cases:
        print(f"\n{'='*72}\n# {case['label']}")
        raw_ranking(base, headers, end_user, case)
        consumer(base, headers, end_user, case, ns.runs)


if __name__ == "__main__":
    main()