rune/scripts/gen_attention_capsule_truth.py

#!/usr/bin/env python3
"""Freeze attention_capsule oracle truth for the Rust port battery
(dichotomy wave #2479 tranche, batch: attention_capsule).

Exercises EVERY public function of src/rune/attention_capsule.py across
every branch: AttentionCapsule.to_dict (obj blocks via to_dict, plain-dict
blocks embedded), to_messages (preamble, string blocks, list history,
"user" block, empty content, dict blocks ignored, list-content beats the
user name, later user blocks overwrite), capture_messages (deep copy),
to_prompt (all sections, memories capped at 5, empty capsule), the
estimate_tokens char//4 estimator, and create_capsule (objective
truncation, relevance overlap >= 2 with len > 10 lines, line truncation,
first-5 cap, all 8 requirement patterns incl. substring + case
insensitivity, and all three over-budget trim loops: memories,
constraints, objective halving).

Blocks travel over the wire tagged:
  {"tag": "obj", "name": getattr(.name), "content": getattr(.content),
   "dict": block.to_dict()}   # hasattr(block, "to_dict")
  {"tag": "dict", "dict": block}            # plain dict (no attrs)
An object WITHOUT to_dict cannot be frozen (to_dict embeds the raw
object) and is out of scope; the tag shape still documents it.

Output: /tmp/ac_truth.json — replayed VERBATIM by
src/rust/attention_capsule/src/branch_battery.rs.
"""
from __future__ import annotations

import json
import os
import sys
import warnings
from pathlib import Path

os.environ["RUNE_AC_CORE"] = "python"  # freeze the ORACLE, never the kernel

sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src"))

from rune import attention_capsule as ac  # noqa: E402
from rune.context_broker import ContributionBlock  # noqa: E402

warnings.simplefilter("ignore", DeprecationWarning)  # oracle health only


class Obj:
    """A block object with to_dict (ContributionBlock-shaped) plus any
    name/content attribute shape the case needs."""

    def __init__(self, d=None, name="", content=""):
        self._d = dict(d) if d is not None else {
            "name": name, "priority": 5, "content": content,
            "tokens": 0, "source": "truth-gen", "evictable": True,
        }
        self.name = self._d.get("name", name)
        self.content = self._d.get("content", content)

    def to_dict(self):
        return dict(self._d)


def block_wire(b):
    if hasattr(b, "to_dict"):
        return {"tag": "obj", "name": getattr(b, "name", ""),
                "content": getattr(b, "content", ""), "dict": b.to_dict()}
    if isinstance(b, dict):
        return {"tag": "dict", "dict": b}
    return {"tag": "obj", "name": getattr(b, "name", ""),
            "content": getattr(b, "content", ""), "dict": None}


def capsule_wire(c: ac.AttentionCapsule) -> dict:
    return {
        "active_objective": c.active_objective,
        "current_phase": c.current_phase,
        "relevant_memories": list(c.relevant_memories),
        "active_commitments": list(c.active_commitments),
        "important_constraints": list(c.important_constraints),
        "current_blocker": c.current_blocker,
        "recent_correction": c.recent_correction,
        "critical_assumptions": list(c.critical_assumptions),
        "token_count": c.token_count,
        "blocks": [block_wire(b) for b in c.blocks],
    }


def mk_capsule(**kw) -> ac.AttentionCapsule:
    blocks = kw.pop("blocks", [])
    c = ac.AttentionCapsule(**kw)
    c.blocks = list(blocks)
    return c


def case(name, input, expected):
    return {"name": name, "input": input, "expected": expected}


to_dict_cases = []
to_messages_cases = []
capture_cases = []
to_prompt_cases = []
token_cases = []
create_cases = []

# ── to_dict ──────────────────────────────────────────────────────────────
c = mk_capsule()
to_dict_cases.append(case("to-dict-empty", capsule_wire(c), c.to_dict()))

c = mk_capsule(
    active_objective="ship the release", current_phase="verify",
    relevant_memories=["m1", "m2"], active_commitments=["c1"],
    important_constraints=["k1", "k2"], current_blocker="flaky test",
    recent_correction="wrong branch", critical_assumptions=["a1"],
    token_count=42,
)
to_dict_cases.append(case("to-dict-full-no-blocks", capsule_wire(c), c.to_dict()))

c = mk_capsule(
    active_objective="o",
    blocks=[
        ContributionBlock("system", 10, "You are terse.", 4, "gen", False),
        {"name": "user", "content": "plain dict block"},
        ContributionBlock("user", 9, "final ask", 3, "gen", False),
        {"arbitrary": True, "nested": [1, 2]},
    ],
)
to_dict_cases.append(case("to-dict-mixed-blocks", capsule_wire(c), c.to_dict()))

c = mk_capsule(
    relevant_memories=["héllo wörld", "日本語メモ"], current_blocker="café ☕",
)
to_dict_cases.append(case("to-dict-unicode", capsule_wire(c), c.to_dict()))

# ── to_messages ──────────────────────────────────────────────────────────
c = mk_capsule()
to_messages_cases.append(case(
    "msg-empty-capsule-no-preamble",
    {"capsule": capsule_wire(c), "system_preamble": ""},
    c.to_messages()))

to_messages_cases.append(case(
    "msg-preamble-only",
    {"capsule": capsule_wire(c), "system_preamble": "Be terse."},
    c.to_messages("Be terse.")))

c = mk_capsule(blocks=[ContributionBlock("sys", 5, "block one", 2, "gen", True)])
to_messages_cases.append(case(
    "msg-string-block-no-preamble",
    {"capsule": capsule_wire(c), "system_preamble": ""},
    c.to_messages()))

c = mk_capsule(blocks=[
    ContributionBlock("a", 5, "first", 2, "gen", True),
    ContributionBlock("b", 5, "", 0, "gen", True),  # empty content skipped
    ContributionBlock("c", 5, "second", 2, "gen", True),
])
to_messages_cases.append(case(
    "msg-string-blocks-order-empty-skipped",
    {"capsule": capsule_wire(c), "system_preamble": "preamble"},
    c.to_messages("preamble")))

history = [{"role": "user", "content": "hi"},
           {"role": "assistant", "content": "hello"}]
c = mk_capsule(blocks=[
    ContributionBlock("history", 5, [dict(m) for m in history], 8, "gen", True),
    ContributionBlock("user", 9, "final question", 3, "gen", False),
])
to_messages_cases.append(case(
    "msg-history-then-user",
    {"capsule": capsule_wire(c), "system_preamble": ""},
    c.to_messages()))

c = mk_capsule(blocks=[
    ContributionBlock("sys", 5, "rules", 2, "gen", True),
    ContributionBlock("history", 5, [dict(m) for m in history], 8, "gen", True),
    {"name": "user", "content": "dict user should be IGNORED"},
    ContributionBlock("user", 9, "", 0, "gen", False),  # empty user dropped
    {"name": "user", "content": "second user wins"},
])
to_messages_cases.append(case(
    "msg-full-shape-dict-blocks-ignored",
    {"capsule": capsule_wire(c), "system_preamble": "preamble"},
    c.to_messages("preamble")))

c = mk_capsule(blocks=[
    ContributionBlock("user", 9, "first user", 3, "gen", False),
    ContributionBlock("user", 9, "second user", 3, "gen", False),
])
to_messages_cases.append(case(
    "msg-two-user-blocks-later-wins",
    {"capsule": capsule_wire(c), "system_preamble": ""},
    c.to_messages()))

c = mk_capsule(blocks=[
    ContributionBlock("user", 9, [dict(m) for m in history], 8, "gen", True),
])
to_messages_cases.append(case(
    "msg-user-named-list-content-is-history",
    {"capsule": capsule_wire(c), "system_preamble": ""},
    c.to_messages()))

c = mk_capsule(blocks=[
    Obj(name="user", content="unnamed-to_dict-obj"),  # to_dict unused here
])
to_messages_cases.append(case(
    "msg-obj-block-getattr-defaults",
    {"capsule": capsule_wire(c), "system_preamble": ""},
    c.to_messages()))

# ── capture_messages ─────────────────────────────────────────────────────
for name, msgs in [
    ("capture-basic", [{"role": "user", "content": "hi"}]),
    ("capture-empty", []),
    ("capture-nested", [{"role": "user", "content": "x",
                         "extra": {"deep": [1, 2, {"k": None}]}}]),
    ("capture-two", [{"role": "system", "content": "s"},
                     {"role": "user", "content": "u"}]),
]:
    c = mk_capsule()
    got = c.capture_messages(msgs)
    capture_cases.append(case(name, [json.loads(json.dumps(m)) for m in msgs], got))
    assert got == c.sent_messages

# ── to_prompt ────────────────────────────────────────────────────────────
c = mk_capsule()
to_prompt_cases.append(case("prompt-empty", capsule_wire(c), c.to_prompt()))

c = mk_capsule(active_objective="fix the flaky test")
to_prompt_cases.append(case("prompt-objective-only", capsule_wire(c), c.to_prompt()))

c = mk_capsule(
    active_objective="ship", current_phase="review",
    relevant_memories=[f"mem-{i}" for i in range(7)],  # capped at 5
    active_commitments=["commit-a", "commit-b"],
    important_constraints=["no new deps", "keep API"],
    current_blocker="waiting on CI", recent_correction="rebased",
    critical_assumptions=["gitea up", "kernel fresh"],
)
to_prompt_cases.append(case("prompt-all-sections-memories-capped",
                            capsule_wire(c), c.to_prompt()))

c = mk_capsule(relevant_memories=["single"])
to_prompt_cases.append(case("prompt-single-memory", capsule_wire(c), c.to_prompt()))

c = mk_capsule(active_objective="目标：发布版本", current_blocker="阻塞：测试 ☕")
to_prompt_cases.append(case("prompt-unicode", capsule_wire(c), c.to_prompt()))

# ── estimate_tokens ──────────────────────────────────────────────────────
for name, text in [
    ("tokens-empty", ""),
    ("tokens-three", "abc"),
    ("tokens-four", "abcd"),
    ("tokens-five", "abcde"),
    ("tokens-phrase", "hello world, this is a test"),
    ("tokens-unicode", "héllo wörld ★★★★ 日本語"),
]:
    token_cases.append(case(name, text, ac.estimate_tokens(text)))

# ── create_capsule ───────────────────────────────────────────────────────
ctx = ("Previous attempt failed due to import error. Benchmark uses pytest. "
       "Similar fix in commit abc123.")
c = ac.create_capsule("Fix the failing test", ctx)
create_cases.append(case(
    "create-basic-relevance",
    {"task": "Fix the failing test", "context": ctx, "max_tokens": 600},
    capsule_wire(c)))
assert c.active_objective == "Fix the failing test"
assert c.token_count <= 600

c = ac.create_capsule("No context task")
create_cases.append(case(
    "create-no-context",
    {"task": "No context task", "context": "", "max_tokens": 600},
    capsule_wire(c)))

lines = [
    "short",                        # len <= 10 → skipped
    "",                             # empty → skipped
    "one word only update",         # overlap: task words {update,...}
    "update the failing test now",  # 3+ overlap → kept
    "totally unrelated line here",  # 0-1 overlap → skipped
    "update failing test again please",  # kept
    "x" * 11,                       # len 11 but no overlap → skipped
    "update failing test and update more update",  # kept
    "update failing test fourth line",             # kept
    "update failing test fifth line",              # kept
    "update failing test sixth line",              # kept → cap at 5
    "update failing test seventh line",            # over cap → dropped
]
c = ac.create_capsule("update the failing test", "\n".join(lines))
create_cases.append(case(
    "create-relevance-filters-and-cap",
    {"task": "update the failing test", "context": "\n".join(lines),
     "max_tokens": 600},
    capsule_wire(c)))
assert len(c.relevant_memories) <= 5

long_line = "update the failing test " + "y" * 120
c = ac.create_capsule("update the failing test", long_line)
create_cases.append(case(
    "create-line-truncated-to-100",
    {"task": "update the failing test", "context": long_line,
     "max_tokens": 600},
    capsule_wire(c)))
assert all(len(m) <= 100 for m in c.relevant_memories)

c = ac.create_capsule("I need to fix and build, please create and update the change; want to?")
create_cases.append(case(
    "create-all-requirement-patterns",
    {"task": "I need to fix and build, please create and update the change; want to?",
     "context": "", "max_tokens": 600},
    capsule_wire(c)))

c = ac.create_capsule("PREFIX the FIXTURE change")
create_cases.append(case(
    "create-pattern-substring-case-insensitive",
    {"task": "PREFIX the FIXTURE change", "context": "", "max_tokens": 600},
    capsule_wire(c)))

c = ac.create_capsule("the weather is nice")
create_cases.append(case(
    "create-no-requirement-patterns",
    {"task": "the weather is nice", "context": "", "max_tokens": 600},
    capsule_wire(c)))

c = ac.create_capsule("B" * 260, "", max_tokens=600)
create_cases.append(case(
    "create-task-truncated-200",
    {"task": "B" * 260, "context": "", "max_tokens": 600},
    capsule_wire(c)))
assert c.active_objective == "B" * 200

# trim loop 1: memories popped until under budget
fat_ctx = "\n".join(f"update the failing test with memory number {i} padded" for i in range(30))
c = ac.create_capsule("update the failing test", fat_ctx, max_tokens=100)
create_cases.append(case(
    "create-trim-memories",
    {"task": "update the failing test", "context": fat_ctx,
     "max_tokens": 100},
    capsule_wire(c)))
assert c.token_count <= 100

# trim loop 2: constraints popped when memories exhausted
c = ac.create_capsule("A" * 400 + " fix build create update change please need to want to",
                      "", max_tokens=100)
create_cases.append(case(
    "create-trim-constraints",
    {"task": "A" * 400 + " fix build create update change please need to want to",
     "context": "", "max_tokens": 100},
    capsule_wire(c)))
assert c.token_count <= 100

# trim loop 3: objective halved while len > 50
c = ac.create_capsule("A" * 1000, "", max_tokens=25)
create_cases.append(case(
    "create-trim-objective-halving",
    {"task": "A" * 1000, "context": "", "max_tokens": 25},
    capsule_wire(c)))
assert c.token_count <= 25

# boundary: exactly at budget → no trims
c = ac.create_capsule("abcd", "", max_tokens=1)  # 4 chars → 1 token
create_cases.append(case(
    "create-exact-budget-no-trim",
    {"task": "abcd", "context": "", "max_tokens": 1},
    capsule_wire(c)))

# unicode: char-based lengths, not bytes
c = ac.create_capsule("日本語のテスト" * 40, "", max_tokens=30)
create_cases.append(case(
    "create-unicode-char-semantics",
    {"task": "日本語のテスト" * 40, "context": "", "max_tokens": 30},
    capsule_wire(c)))

# ── oracle self-test health ──────────────────────────────────────────────
ac.test_capsule()  # raises on oracle regression

truth = {
    "frozen": "attention_capsule to_dict/to_messages/capture_messages/to_prompt/"
              "estimate_tokens/create_capsule",
    "oracle": "src/rune/attention_capsule.py",
    "note": "Blocks travel tagged: obj {tag,name,content,dict} (hasattr "
            "to_dict) vs dict {tag,dict} (plain dict — to_messages ignores "
            "it via getattr, to_dict embeds it). DeprecationWarnings are a "
            "facade concern; the kernel is warning-free.",
    "to_dict": to_dict_cases,
    "to_messages": to_messages_cases,
    "capture_messages": capture_cases,
    "to_prompt": to_prompt_cases,
    "estimate_tokens": token_cases,
    "create_capsule": create_cases,
}

out = Path("/tmp/ac_truth.json")
out.write_text(json.dumps(truth, indent=1, ensure_ascii=False) + "\n")
print(f"froze {len(to_dict_cases)} to_dict + {len(to_messages_cases)} to_messages "
      f"+ {len(capture_cases)} capture_messages + {len(to_prompt_cases)} to_prompt "
      f"+ {len(token_cases)} estimate_tokens + {len(create_cases)} create_capsule "
      f"-> {out}")