Structured output and tool calls

"""constrained.py: free sampling versus grammar-constrained decoding for a JSON reply.

A simulated model classifies support messages and replies with
{"label": <one of three>, "confidence": 0.00-1.00}. Sometimes it adds a preamble,
wraps the JSON in a code fence, or uses single quotes. Constrained decoding masks
every token that cannot extend a valid output. For this tiny schema the set of valid
strings is finite, so its prefix set is the automaton; real engines compile a JSON
Schema or grammar into a finite-state or pushdown automaton instead.
Standard library only; the model is a stand-in.
"""

from __future__ import annotations

import json
import math
import random

LABELS = ("refund", "shipping", "other")
VOCAB = ['{', '}', '"label"', '"label": ', ':', ': ', ',', ', ', '"confidence"',
         '"confidence": ', '"refund"', '"shipping"', '"other"', "'label'", "'refund'",
         "'shipping'", "'other'", "'confidence'", '0', '1', '2', '3', '4', '5', '6', '7',
         '8', '9', '0.', '1.', '.', ' ', '\n', '```', '```json\n', 'Sure', '! ', 'Here',
         ' is', ' the', ' JSON', ':\n', '<eos>']


def valid_strings() -> set[str]:
    out = set()
    for label in LABELS:
        for c in range(101):
            out.add(f'{{"label": "{label}", "confidence": {c // 100}.{c % 100:02d}}}')
    return out


VALID = valid_strings()
PREFIXES = {s[:i] for s in VALID for i in range(len(s) + 1)}


def intended(rng: random.Random, truth: str) -> str:
    """What the model 'wants' to write, before any constraint."""
    label = truth if rng.random() < 0.9 else rng.choice([l for l in LABELS if l != truth])
    body = f'{{"label": "{label}", "confidence": 0.{rng.randrange(60, 99)}}}'
    r = rng.random()
    if r < 0.06:
        return "Sure! Here is the JSON:\n" + body
    if r < 0.10:
        return "```json\n" + body + "\n```"
    if r < 0.13:
        return body.replace('"', "'")
    return body


def logits(text: str, plan: str, rng: random.Random) -> list[float]:
    """High score for tokens that continue the plan (or the JSON inside it), noise elsewhere."""
    core = plan[plan.find("{"):plan.rfind("}") + 1].replace("'", '"')
    out = []
    for tok in VOCAB:
        score = rng.gauss(0.0, 1.0)
        for target, bonus in ((plan, 9.0), (core, 6.0)):
            if target.startswith(text):
                rest = target[len(text):]
                if tok == "<eos>" and rest == "":
                    score += bonus
                elif tok != "<eos>" and rest.startswith(tok):
                    score += bonus + len(tok) * 0.1  # prefer longer matching tokens
        out.append(score)
    return out


def decode(plan: str, rng: random.Random, constrained: bool, max_tokens: int) -> tuple[str, int]:
    text, masked = "", 0
    for _ in range(max_tokens):
        scores = logits(text, plan, rng)
        if constrained:
            for i, tok in enumerate(VOCAB):
                ok = (text in VALID) if tok == "<eos>" else (text + tok) in PREFIXES
                if not ok:
                    scores[i] = -math.inf
                    masked += 1
        tok = VOCAB[max(range(len(VOCAB)), key=lambda i: scores[i])]   # greedy
        if tok == "<eos>":
            break
        text += tok
    return text, masked


def check(text: str, truth: str) -> tuple[bool, bool]:
    try:
        obj = json.loads(text)
        ok = (isinstance(obj, dict) and set(obj) == {"label", "confidence"}
              and obj["label"] in LABELS and 0 <= obj["confidence"] <= 1)
    except (json.JSONDecodeError, TypeError):
        return False, False
    return ok, ok and obj["label"] == truth


def main() -> None:
    n = 2000
    print(f"{n} simulated replies; schema: label in {LABELS}, confidence 0.00-1.00")
    print(f"{'mode':<28}{'parses + valid':>15}{'right label':>13}{'masked/reply':>14}")
    for constrained, max_tokens in ((False, 40), (True, 40), (True, 6)):
        rng = random.Random(11)
        valid = right = masked = 0
        sample = ""
        for _ in range(n):
            truth = rng.choice(LABELS)
            plan = intended(rng, truth)
            text, m = decode(plan, rng, constrained, max_tokens)
            v, r = check(text, truth)
            valid += v
            right += r
            masked += m
            if not v and not sample:
                sample = text
        mode = ("constrained" if constrained else "free sampling") + f", max {max_tokens} tok"
        print(f"{mode:<28}{valid / n:>15.1%}{right / n:>13.1%}{masked / n:>14.0f}")
        if sample:
            print(f"    first invalid output: {sample!r}")


if __name__ == "__main__":
    main()

2000 simulated replies; schema: label in ('refund', 'shipping', 'other'), confidence 0.00-1.00
mode                         parses + valid  right label  masked/reply
free sampling, max 40 tok             86.7%        78.1%             0
    first invalid output: 'Sure! Here is the JSON:\n{"label": "shipping", "confidence": 0.91}'
constrained, max 40 tok              100.0%        90.1%           479
constrained, max 6 tok                 0.0%         0.0%           246
    first invalid output: '{"label": "shipping", "confidence": '

"""tool_harness.py: the model proposes, the harness decides.

A harness receives tool calls from a simulated model and, before anything runs,
checks: the tool exists and is allowed for this caller, the arguments match the JSON
Schema subset declared for the tool, business rules hold, irreversible effects have
approval, and a retried call with the same idempotency key does not run twice.
Standard library only; the orders and the model's calls are invented.
"""

from __future__ import annotations

import hashlib
import re
from dataclasses import dataclass, field
from typing import Any, Callable

ORDERS = {"ord_8841": {"status": "delivered", "total": 42.00, "owner": "u_17"}}


@dataclass
class Tool:
    name: str
    effect: str                      # "read", "reversible", "irreversible"
    schema: dict[str, Any]           # a JSON Schema subset: type, properties, required
    run: Callable[[dict[str, Any]], dict[str, Any]]
    rules: Callable[[dict[str, Any], str], str | None] = lambda a, u: None


def type_ok(value: Any, spec: dict[str, Any]) -> str | None:
    t = spec["type"]
    if t == "string" and not isinstance(value, str):
        return "must be a string"
    if t == "number" and (isinstance(value, bool) or not isinstance(value, (int, float))):
        return "must be a number"
    if "pattern" in spec and not re.search(spec["pattern"], value):
        return f"must match {spec['pattern']!r}"
    if "enum" in spec and value not in spec["enum"]:
        return f"must be one of {spec['enum']}"
    if "exclusiveMinimum" in spec and not value > spec["exclusiveMinimum"]:
        return f"must be > {spec['exclusiveMinimum']}"
    return None


def validate(args: dict[str, Any], schema: dict[str, Any]) -> str | None:
    props = schema["properties"]
    for key in schema["required"]:
        if key not in args:
            return f"missing required argument {key!r}"
    for key, value in args.items():
        if key not in props:
            return f"unknown argument {key!r}"           # additionalProperties: false
        err = type_ok(value, props[key])
        if err:
            return f"{key} {err}"
    return None


def refund_rules(args: dict[str, Any], user: str) -> str | None:
    order = ORDERS.get(args["order_id"])
    if order is None or order["owner"] != user:
        return "no such order for this user"             # same answer: do not leak existence
    if order["status"] != "delivered":
        return "only delivered orders can be refunded"
    if args["amount"] > order["total"]:
        return f"amount exceeds order total {order['total']:.2f}"
    return None


ORDER_ID = {"type": "string", "pattern": "^ord_[0-9]+$"}
TOOLS = {
    "lookup_order": Tool("lookup_order", "read",
                         {"properties": {"order_id": ORDER_ID}, "required": ["order_id"]},
                         lambda a: {k: v for k, v in ORDERS[a["order_id"]].items() if k != "owner"},
                         lambda a, u: None if ORDERS.get(a["order_id"], {}).get("owner") == u
                         else "no such order for this user"),
    "issue_refund": Tool("issue_refund", "irreversible",
                         {"properties": {"order_id": ORDER_ID,
                                         "amount": {"type": "number", "exclusiveMinimum": 0},
                                         "reason": {"type": "string",
                                                    "enum": ["damaged", "late", "other"]}},
                          "required": ["order_id", "amount", "reason"]},
                         lambda a: {"refund_id": "r_" + hashlib.sha256(a["order_id"].encode()).hexdigest()[:6]},
                         refund_rules),
}


@dataclass
class Harness:
    user: str
    allowed: set[str]
    approvals: set[str] = field(default_factory=set)      # idempotency keys a human approved
    done: dict[str, dict[str, Any]] = field(default_factory=dict)
    executed: int = 0

    def handle(self, call: dict[str, Any]) -> str:
        name, args, key = call.get("name"), call.get("arguments", {}), call.get("idempotency_key")
        tool = TOOLS.get(name)
        if tool is None or name not in self.allowed:
            return f"DENY   {name}: not an allowed tool"
        err = validate(args, tool.schema) or tool.rules(args, self.user)
        if err:
            return f"REJECT {name}: {err}  (returned to the model so it can correct the call)"
        if tool.effect != "read":
            if not key:
                return f"REJECT {name}: writes need an idempotency key"
            if key in self.done:
                return f"DEDUP  {name}: key {key} already ran, returning {self.done[key]}"
            if tool.effect == "irreversible" and key not in self.approvals:
                return f"HOLD   {name}: waiting for human approval of {args}"
        result = tool.run(args)
        self.executed += 1
        if key:
            self.done[key] = result
        return f"RUN    {name}: {result}"


def main() -> None:
    h = Harness(user="u_17", allowed={"lookup_order", "issue_refund"})
    calls = [
        {"name": "lookup_order", "arguments": {"order_id": "8841"}},
        {"name": "lookup_order", "arguments": {"order_id": "ord_8841"}},
        {"name": "issue_refund", "arguments": {"order_id": "ord_8841", "amount": "42",
                                               "reason": "damaged"}},
        {"name": "issue_refund", "arguments": {"order_id": "ord_8841", "amount": 420.0,
                                               "reason": "damaged"}, "idempotency_key": "k1"},
        {"name": "issue_refund", "arguments": {"order_id": "ord_8841", "amount": 42.0,
                                               "reason": "damaged"}, "idempotency_key": "k2"},
        "approve k2",
        {"name": "issue_refund", "arguments": {"order_id": "ord_8841", "amount": 42.0,
                                               "reason": "damaged"}, "idempotency_key": "k2"},
        {"name": "issue_refund", "arguments": {"order_id": "ord_8841", "amount": 42.0,
                                               "reason": "damaged"}, "idempotency_key": "k2"},
        {"name": "delete_account", "arguments": {"user": "u_17"}},
        {"name": "lookup_order", "arguments": {"order_id": "ord_9999", "verbose": True}},
    ]
    for i, call in enumerate(calls, 1):
        if isinstance(call, str):
            h.approvals.add(call.split()[1])
            print(f"{i:>2}. human approves {call.split()[1]}")
            continue
        print(f"{i:>2}. {h.handle(call)}")
    print(f"\ntool executions: {h.executed} (one read, one refund) out of {len(calls) - 1} proposals")


if __name__ == "__main__":
    main()

 1. REJECT lookup_order: order_id must match '^ord_[0-9]+$'  (returned to the model so it can correct the call)
 2. RUN    lookup_order: {'status': 'delivered', 'total': 42.0}
 3. REJECT issue_refund: amount must be a number  (returned to the model so it can correct the call)
 4. REJECT issue_refund: amount exceeds order total 42.00  (returned to the model so it can correct the call)
 5. HOLD   issue_refund: waiting for human approval of {'order_id': 'ord_8841', 'amount': 42.0, 'reason': 'damaged'}
 6. human approves k2
 7. RUN    issue_refund: {'refund_id': 'r_2ccfd1'}
 8. DEDUP  issue_refund: key k2 already ran, returning {'refund_id': 'r_2ccfd1'}
 9. DENY   delete_account: not an allowed tool
10. REJECT lookup_order: unknown argument 'verbose'  (returned to the model so it can correct the call)

tool executions: 2 (one read, one refund) out of 9 proposals