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 at DIFFERENT WING SCOPES (the distinction is load-bearing — Worldtree #393):
  * RAW RANKING — ALL-WING (mimir searches kb+main+fiction, ~9800 rows). Drive `mimir` with the
    descriptive query + variants; report the query mimir ACTUALLY passed (tool_start q=, since it
    reformulates) and where the intended ENTITY lands. Cross-wing dilution can push a one-arm
    vector bridge out of the bge rescue window entirely -> entity ABSENT from top-k.
  * CONSUMER — FICTION (ratatoskr:donut is fiction-scoped, ~1578 rows). Drive Donut N times;
    classify BINDS-ENTITY vs MIS-BINDS-DECOY vs OTHER, and report where the ENTITY and DECOY rank
    in her fiction-scoped results. The ranks tell WHICH failure fired per run: entity PRESENT +
    mis-bind = subject-selection; entity ABSENT + a present decoy = reformulation-induced absence
    (Donut distilled the descriptive phrase to bare tokens that don't carry the vocabulary bridge).
    Empirically Donut mostly does the latter — she distills "the guy with the roid rage" to bare
    "roid rage", so Juicer drops out even at fiction scope and she binds a present decoy (Jack).

Two-seam finding (2026-08-07, v1.0.0b181) — both lose the entity, by DIFFERENT reformulation seams:
(1) mimir PRESERVES the phrase -> cross-wing dilution (kb+main+fiction) drops the entity from its
all-wing top-k; (2) Donut DISTILLS the phrase to bare tokens -> the entity drops even at fiction
scope (ent@None) and she binds a present decoy. The fold's ENT@/DEC@ ranks separate reformulation-
absence from true subject-selection per run. Unifying lever: disambiguating-vocabulary expansion
(the full phrase, or +attribute like "steroid") surfaces the entity at fiction scope — Donut's
expand-runs bind correctly. "dangerous crown" is mostly RESOLVED (one entity). Root gap:
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 _rank_in(res: dict, pattern: str | None) -> int | None:
    """Rank of the first row whose excerpt matches `pattern` (None if absent / no pattern)."""
    if not pattern:
        return None
    rows = res.get("hits", res.get("results", [])) if isinstance(res, dict) else []
    return next((i for i, r in enumerate(rows) if isinstance(r, dict)
                 and re.search(pattern, r.get("excerpt", ""), re.I)), None)


def raw_ranking(base, headers, end_user, case) -> None:
    print("\n  [raw ranking — ALL-WING/mimir] actual tool query (reformulation seam) + entity rank")
    for query in case["variants"]:
        sid = _session(base, headers, "mimir", end_user)  # fresh session per query
        mimir_q, 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)]
        rank = entity_ranks[0] if entity_ranks else "ABSENT (not in top-k)"
        # mimir_q is load-bearing for #393: separates reformulation-at-the-agent-seam
        # (mimir distilled/expanded the phrase) from ranking (the tool ranked it low).
        print(f"    instructed={query!r:28} mimir_q={mimir_q!r:38} entity_rank={rank}")
        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 — FICTION/donut] x{runs} on {case['question']!r}"
          f"  (ent@/dec@ = rank in Donut's fiction-scoped results)")
    for run in range(1, runs + 1):
        sid = _session(base, headers, "ratatoskr:donut", end_user)  # fresh session per run
        q, res, 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
        er, dr = _rank_in(res, case["entity"]), _rank_in(res, case["decoy"])
        print(f"    run{run}: {v:16} ent@{er} dec@{dr}  q={q!r:30} :: {ans.strip()[:56]}")
    # a present entity (ent@ not None) co-occurring with a mis-bind is subject-selection,
    # NOT ranking-absence — the fiction-scope half of the #393 two-mechanism split.
    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()