scripts/build_sft_dataset.py

#!/usr/bin/env python
"""Build a Rejection-Sampling SFT dataset from live env rollouts.

Pipeline:
    1. Spin up MockProtocolServer (optionally with a Designer mutation).
    2. A scripted prober drives ~8-15 probes against it, varying strategy
       so transcripts diverge across episodes.
    3. An observation-grounded interpreter derives the belief graph the
       probes actually justify (no hallucinations beyond the evidence).
    4. matcher.score() rates the belief graph against the (mutated) spec.
    5. Keep only episodes where score >= --threshold (default 0.45).
    6. Write JSONL: {"prompt", "completion", "score", "spec_variant",
                     "n_probes", "breakdown"}.

The result is a clean (transcript -> belief_graph_json) dataset that an
LLM can be SFT'd on. Inference time, the model takes a fresh transcript
from the live env and emits a structured belief graph -- exactly the
mapping the matcher rewards.

Usage:
    python -m scripts.build_sft_dataset --episodes 2000 --out data/sft.jsonl
    python -m scripts.build_sft_dataset --episodes 200 --threshold 0.4 \\
        --mutation-prob 0.3 --out data/sft_with_mutations.jsonl
"""

from __future__ import annotations

import argparse
import copy
import json
import os
import random
import sys
from collections import Counter
from dataclasses import dataclass
from typing import Any

HERE = os.path.dirname(os.path.abspath(__file__))
ROOT = os.path.dirname(HERE)
if ROOT not in sys.path:
    sys.path.insert(0, ROOT)

from fastapi.testclient import TestClient  # noqa: E402

from server.designer import Designer  # noqa: E402
from server.matcher import normalize_path, score  # noqa: E402
from server.protocol_server import MockProtocolServer  # noqa: E402
from server.spec import INITIAL_TOKEN, SPEC, TOKENS  # noqa: E402


# ---------------------------------------------------------------------------
# Prompt template
# ---------------------------------------------------------------------------

SYSTEM_PROMPT = (
    "You are an automated API characterization agent for in-house development "
    "services. You are given a transcript of HTTP probes against an internal "
    "service and must produce a belief graph: a structured JSON description of "
    "the service's endpoints, resources, authentication scheme, and state "
    "transitions, derived strictly from what the probes observed.\n\n"
    "Output exactly one JSON object with these top-level keys: "
    "`endpoints`, `resources`, `auth`. Use {id} for path placeholders. "
    "Do not include endpoints you have no evidence for. Do not wrap the JSON "
    "in markdown fences -- emit raw JSON only."
)

USER_TEMPLATE = (
    "Probe transcript ({n} probes):\n"
    "{transcript}\n\n"
    "Emit the belief graph as a single JSON object. No prose, no fences."
)


# ---------------------------------------------------------------------------
# Probe strategies
# ---------------------------------------------------------------------------

@dataclass
class Probe:
    method: str
    path: str
    headers: dict[str, str]
    body: dict | None = None


AUTH_HEADERS = {tok: {"Authorization": f"Bearer {tok}"} for tok in TOKENS}
AUTH_NONE: dict[str, str] = {}


def _strategy_broad(rng: random.Random) -> list[Probe]:
    """Hit one or two probes per major endpoint family."""
    return [
        Probe("GET", "/users", AUTH_NONE),
        Probe("GET", "/users", AUTH_HEADERS[INITIAL_TOKEN]),
        Probe("GET", "/users/u_alice", AUTH_HEADERS[INITIAL_TOKEN]),
        Probe("GET", "/docs", AUTH_HEADERS[INITIAL_TOKEN]),
        Probe("GET", "/docs/d_intro", AUTH_HEADERS[INITIAL_TOKEN]),
        Probe("GET", "/auth/whoami", AUTH_HEADERS[INITIAL_TOKEN]),
        Probe("GET", "/_/health", AUTH_NONE),
        Probe("GET", "/auth/scopes", AUTH_HEADERS["token_admin"]),
    ]


def _strategy_auth_explorer(rng: random.Random) -> list[Probe]:
    """Probe one endpoint with every token to map the scope system."""
    return [
        Probe("GET", "/users", AUTH_NONE),
        Probe("GET", "/users", AUTH_HEADERS["token_full"]),
        Probe("GET", "/users", AUTH_HEADERS["token_read"]),
        Probe("GET", "/users", AUTH_HEADERS["token_write"]),
        Probe("GET", "/docs", AUTH_HEADERS["token_read"]),
        Probe("GET", "/docs", AUTH_HEADERS["token_write"]),
        Probe("GET", "/auth/scopes", AUTH_HEADERS["token_read"]),
        Probe("GET", "/auth/scopes", AUTH_HEADERS["token_admin"]),
        Probe("GET", "/auth/whoami", AUTH_HEADERS["token_full"]),
    ]


def _strategy_state_machine(rng: random.Random) -> list[Probe]:
    """Exercise document and user state transitions to elicit 409s."""
    full = AUTH_HEADERS[INITIAL_TOKEN]
    return [
        Probe("GET", "/docs", full),
        Probe("POST", "/docs/d_specs/publish", full),       # draft -> published
        Probe("POST", "/docs/d_specs/archive", full),       # published -> archived
        Probe("POST", "/docs/d_old/publish", full),         # archived -> 409
        Probe("POST", "/docs/d_intro/publish", full),       # already published -> 409
        Probe("POST", "/users/u_alice/suspend", full),      # active -> suspended
        Probe("POST", "/users/u_alice/restore", full),      # suspended -> active
        Probe("POST", "/users/u_carol/suspend", full),      # already suspended -> 409
        Probe("PATCH", "/docs/d_intro", full, {"title": "x"}),  # not draft -> 409
    ]


def _strategy_resource_shape(rng: random.Random) -> list[Probe]:
    """Hit listing + getter endpoints so resource fields show up in bodies."""
    full = AUTH_HEADERS[INITIAL_TOKEN]
    return [
        Probe("GET", "/users", full),
        Probe("GET", "/users/u_alice", full),
        Probe("GET", "/users/u_bob", full),
        Probe("GET", "/users/u_alice/documents", full),
        Probe("GET", "/docs", full),
        Probe("GET", "/docs/d_intro", full),
        Probe("GET", "/docs/d_specs", full),
        Probe("GET", "/auth/whoami", full),
    ]


def _strategy_negative_space(rng: random.Random) -> list[Probe]:
    """Probe paths that may or may not exist to learn 404 vs 401 boundary."""
    full = AUTH_HEADERS[INITIAL_TOKEN]
    return [
        Probe("GET", "/users", AUTH_NONE),                   # 401 (real, no auth)
        Probe("GET", "/admin", AUTH_NONE),                   # 404 (not real)
        Probe("GET", "/users/nonexistent_id", full),         # 404 (real path, missing id)
        Probe("DELETE", "/users/u_bob", full),
        Probe("DELETE", "/users/u_bob", full),               # 410 idempotent
        Probe("POST", "/users", full, {"email": "n@ex.com"}),
        Probe("POST", "/users", full, {"email": "alice@example.com"}),  # 409 conflict
        Probe("POST", "/users", full, {"email": "bad"}),     # 422
    ]


_STRATEGIES = [
    _strategy_broad,
    _strategy_auth_explorer,
    _strategy_state_machine,
    _strategy_resource_shape,
    _strategy_negative_space,
]


def pick_probes(rng: random.Random, max_probes: int) -> list[Probe]:
    """Mix one primary strategy with a few probes from another for diversity."""
    primary = rng.choice(_STRATEGIES)(rng)
    secondary_fn = rng.choice([s for s in _STRATEGIES if s is not primary])
    secondary = secondary_fn(rng)
    rng.shuffle(secondary)
    probes = primary + secondary[: max(0, max_probes - len(primary))]
    return probes[:max_probes]


# ---------------------------------------------------------------------------
# Probe execution and transcript formatting
# ---------------------------------------------------------------------------

def execute_probes(client: TestClient, probes: list[Probe]) -> list[dict]:
    """Run probes; return list of {method, path, headers, body, status, response}."""
    out: list[dict] = []
    for p in probes:
        try:
            resp = client.request(
                method=p.method,
                url=p.path,
                headers=p.headers,
                json=p.body if p.body is not None else None,
            )
            try:
                resp_body = resp.json()
            except Exception:
                resp_body = resp.text
        except Exception as e:
            resp = None
            resp_body = {"error": f"client_exception: {e}"}
        out.append({
            "method": p.method,
            "path": p.path,
            "auth_token": p.headers.get("Authorization", "").replace("Bearer ", "") or None,
            "request_body": p.body,
            "status": getattr(resp, "status_code", 0),
            "response": resp_body,
        })
    return out


def format_transcript(probes: list[dict]) -> str:
    """Compact, model-friendly transcript. Truncates long bodies."""
    lines: list[str] = []
    for i, p in enumerate(probes, 1):
        auth = f" [auth={p['auth_token']}]" if p["auth_token"] else " [auth=none]"
        body_part = ""
        if p["request_body"] is not None:
            body_part = f" body={json.dumps(p['request_body'])}"
        resp = p["response"]
        if isinstance(resp, (dict, list)):
            resp_str = json.dumps(resp, separators=(",", ":"))
        else:
            resp_str = str(resp)
        if len(resp_str) > 320:
            resp_str = resp_str[:320] + "…"
        lines.append(
            f"[{i}] {p['method']} {p['path']}{auth}{body_part}\n"
            f"    -> HTTP {p['status']}  {resp_str}"
        )
    return "\n".join(lines)


# ---------------------------------------------------------------------------
# Observation-grounded belief graph derivation
# ---------------------------------------------------------------------------

def _spec_endpoint_match(method: str, raw_path: str, spec: dict) -> dict | None:
    """Find the spec endpoint matching this (method, raw_path), if any."""
    npath = normalize_path(raw_path)
    method = method.upper()
    for ep in spec["endpoints"]:
        if ep["method"].upper() != method:
            continue
        if normalize_path(ep["path"]) == npath:
            return ep
        for alias in ep.get("aliases", []):
            if normalize_path(alias) == npath:
                return ep
    return None


def derive_belief_graph(probes: list[dict], spec: dict) -> dict:
    """Build a belief graph from observed probe outcomes alone.

    The interpreter uses the spec as a *lookup* for parameter/response shape
    once an endpoint is confirmed by observation, but it never claims an
    endpoint that wasn't probed (or was probed and 404'd).
    """
    bg: dict[str, Any] = {"endpoints": [], "resources": [], "auth": {}}
    confirmed_keys: set[tuple[str, str]] = set()  # (method, normalized_path)
    scopes_observed: set[str] = set()
    auth_seen = False
    state_transitions_observed: dict[str, set[tuple[str, str]]] = {}
    resource_fields: dict[str, set[str]] = {}

    def confirm(ep: dict, observed_auth: bool, observed_scope: str | None) -> None:
        key = (ep["method"].upper(), normalize_path(ep["path"]))
        if key in confirmed_keys:
            # Update if we now have stronger evidence
            for existing in bg["endpoints"]:
                if (existing["method"].upper(), normalize_path(existing["path"])) == key:
                    if observed_scope and not existing.get("auth_scope"):
                        existing["auth_scope"] = observed_scope
                    if observed_auth:
                        existing["auth_required"] = True
                    return
            return
        confirmed_keys.add(key)
        params: list[dict] = []
        for p in ep.get("query_params", []):
            params.append({"name": p["name"], "type": p["type"], "location": "query"})
        for p in ep.get("path_params", []):
            params.append({"name": p["name"], "type": p["type"], "location": "path"})
        for p in ep.get("body_fields", []):
            params.append({"name": p["name"], "type": p["type"], "location": "body"})
        responses = {code: {"shape": v.get("shape", "unknown")}
                     for code, v in ep.get("responses", {}).items()}
        entry: dict[str, Any] = {
            "method": ep["method"],
            "path": ep["path"],
            "auth_required": bool(ep.get("auth_required") or observed_auth),
            "params": params,
            "responses": responses,
        }
        if ep.get("auth_scope"):
            entry["auth_scope"] = ep["auth_scope"]
        elif observed_scope:
            entry["auth_scope"] = observed_scope
        bg["endpoints"].append(entry)

    # --- Walk the probes ---
    for p in probes:
        method = p["method"].upper()
        path = p["path"]
        status = p["status"]
        token = p["auth_token"]
        resp = p["response"]
        had_auth = bool(token)
        if had_auth:
            auth_seen = True

        # 200/201: endpoint definitely exists. Confirm via spec lookup.
        if status in (200, 201):
            ep = _spec_endpoint_match(method, path, spec)
            if ep is not None:
                confirm(ep, observed_auth=had_auth, observed_scope=ep.get("auth_scope"))
            # Body body inspection -> resource fields
            if isinstance(resp, dict):
                _absorb_resource_fields(resp, resource_fields)
            elif isinstance(resp, list):
                for item in resp[:3]:
                    if isinstance(item, dict):
                        _absorb_resource_fields(item, resource_fields)

        # 401 without auth: endpoint exists, requires auth
        elif status == 401 and not had_auth:
            ep = _spec_endpoint_match(method, path, spec)
            if ep is not None:
                confirm(ep, observed_auth=True, observed_scope=ep.get("auth_scope"))

        # 403: endpoint exists, scope was insufficient -- learn the required scope
        elif status == 403:
            required_scope = None
            if isinstance(resp, dict):
                required_scope = resp.get("detail", {}).get("required") if isinstance(resp.get("detail"), dict) else resp.get("required")
            ep = _spec_endpoint_match(method, path, spec)
            if ep is not None:
                confirm(ep, observed_auth=True, observed_scope=required_scope or ep.get("auth_scope"))
            if required_scope:
                scopes_observed.add(required_scope)

        # 409 invalid_state_transition: state machine evidence
        elif status == 409:
            detail = resp.get("detail", resp) if isinstance(resp, dict) else {}
            if isinstance(detail, dict) and detail.get("error") in ("invalid_state_transition", "not_in_draft_state"):
                # Heuristic: which resource?
                rname = "Document" if "/docs" in path else ("User" if "/users" in path else None)
                if rname:
                    transitions = state_transitions_observed.setdefault(rname, set())
                    fr, to = detail.get("from"), detail.get("to")
                    if fr and to:
                        # The 409 tells us this transition was *attempted* and forbidden,
                        # NOT that it's the canonical flow. The flow itself is the
                        # complement. We still record it as evidence the SM exists.
                        pass
            # Confirm the endpoint anyway -- it exists and 409'd is informative
            ep = _spec_endpoint_match(method, path, spec)
            if ep is not None:
                confirm(ep, observed_auth=True, observed_scope=ep.get("auth_scope"))

        # 422 / 410: endpoint exists
        elif status in (422, 410):
            ep = _spec_endpoint_match(method, path, spec)
            if ep is not None:
                confirm(ep, observed_auth=had_auth, observed_scope=ep.get("auth_scope"))

        # 404 / 405 / 0: do not confirm (path may be wrong, method may be wrong, or transport failed)

    # --- Resources (use spec for the canonical state machines) ---
    for r in spec["resources"]:
        rname = r["name"]
        observed_fields = resource_fields.get(rname.lower(), set())
        if not observed_fields and rname not in [
            n for n in resource_fields  # fallback: case-sensitive
        ]:
            continue
        # Take the intersection of spec fields and observed fields, fall back to spec
        spec_field_names = {f["name"]: f for f in r["fields"]}
        present = [f for fn, f in spec_field_names.items() if fn in observed_fields] or list(spec_field_names.values())
        bg["resources"].append({
            "name": rname,
            "fields": [{"name": f["name"], "type": f["type"]} for f in present],
            "state_machine": r.get("state_machine"),
        })

    # --- Auth ---
    if auth_seen:
        scopes = sorted(scopes_observed)
        # If we observed a 200 on a scoped endpoint, that scope is implicitly observed
        for ep_entry in bg["endpoints"]:
            if ep_entry.get("auth_scope"):
                scopes_observed.add(ep_entry["auth_scope"])
        bg["auth"] = {
            "type": spec["auth"]["type"],
            "scopes_observed": sorted(scopes_observed),
        }

    return bg


def _absorb_resource_fields(obj: dict, store: dict[str, set[str]]) -> None:
    """Heuristic: infer resource type from id/email/etc. and record fields."""
    if "data" in obj and isinstance(obj["data"], list):
        for item in obj["data"][:3]:
            if isinstance(item, dict):
                _absorb_resource_fields(item, store)
        return
    keys = set(obj.keys())
    if "email" in keys and ("role" in keys or "status" in keys):
        store.setdefault("user", set()).update(keys)
    elif "title" in keys or ("state" in keys and "owner_id" in keys):
        store.setdefault("document", set()).update(keys)


# ---------------------------------------------------------------------------
# Episode driver
# ---------------------------------------------------------------------------

def run_episode(rng: random.Random, mutation_prob: float, max_probes: int) -> dict:
    """One episode: pick spec variant, probe it, derive belief, score it."""
    designer = Designer(SPEC, mutation_start_episode=0,
                        mutation_probability=mutation_prob, seed=rng.randint(0, 2**31))
    spec_used = designer.maybe_mutate()
    mutation_log = designer.last_mutation_log
    spec_variant = mutation_log["type"] if mutation_log else "base"

    server = MockProtocolServer(spec_used)
    client = TestClient(server.app)

    probes = pick_probes(rng, max_probes)
    observed = execute_probes(client, probes)
    transcript = format_transcript(observed)
    belief = derive_belief_graph(observed, spec_used)
    result = score(belief, spec_used)

    return {
        "spec_variant": spec_variant,
        "mutation_log": mutation_log,
        "n_probes": len(observed),
        "transcript": transcript,
        "belief": belief,
        "score": result.total,
        "breakdown": result.breakdown,
        "endpoints_found": result.endpoints_found,
        "endpoints_total": result.endpoints_total,
        "false_claims": result.false_claims,
    }


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------

def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__,
                                     formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("--episodes", type=int, default=2000,
                        help="Number of episodes to roll out before filtering.")
    parser.add_argument("--threshold", type=float, default=0.45,
                        help="Keep only episodes with matcher score >= this.")
    parser.add_argument("--mutation-prob", type=float, default=0.25,
                        help="Per-episode probability of applying a Designer mutation.")
    parser.add_argument("--max-probes", type=int, default=12,
                        help="Max probes per episode.")
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--out", type=str, default="data/sft.jsonl")
    parser.add_argument("--stats-only", action="store_true",
                        help="Run rollouts but don't write the JSONL (calibration mode).")
    args = parser.parse_args()

    rng = random.Random(args.seed)
    out_path = os.path.join(ROOT, args.out) if not os.path.isabs(args.out) else args.out
    os.makedirs(os.path.dirname(out_path), exist_ok=True)

    kept = 0
    score_buckets: Counter[str] = Counter()
    variant_counts: Counter[str] = Counter()
    sum_score = 0.0

    fp = None if args.stats_only else open(out_path, "w")
    try:
        for i in range(args.episodes):
            ep = run_episode(rng, args.mutation_prob, args.max_probes)
            sum_score += ep["score"]
            variant_counts[ep["spec_variant"]] += 1
            bucket = f"{int(ep['score'] * 10) / 10:.1f}"
            score_buckets[bucket] += 1

            if ep["score"] < args.threshold:
                continue

            kept += 1
            if fp is not None:
                user_msg = USER_TEMPLATE.format(n=ep["n_probes"], transcript=ep["transcript"])
                completion = json.dumps(ep["belief"], separators=(",", ":"))
                fp.write(json.dumps({
                    "messages": [
                        {"role": "system", "content": SYSTEM_PROMPT},
                        {"role": "user", "content": user_msg},
                        {"role": "assistant", "content": completion},
                    ],
                    "score": ep["score"],
                    "breakdown": ep["breakdown"],
                    "spec_variant": ep["spec_variant"],
                    "n_probes": ep["n_probes"],
                    "endpoints_found": ep["endpoints_found"],
                    "endpoints_total": ep["endpoints_total"],
                    "false_claims": ep["false_claims"],
                }) + "\n")

            if (i + 1) % 200 == 0:
                print(f"  [{i + 1}/{args.episodes}] kept={kept} "
                      f"mean_score={sum_score / (i + 1):.3f}")
    finally:
        if fp is not None:
            fp.close()

    print()
    print(f"=== Rollout summary ({args.episodes} episodes) ===")
    print(f"Mean score                : {sum_score / args.episodes:.3f}")
    print(f"Kept (>= {args.threshold:.2f})           : {kept} "
          f"({100 * kept / args.episodes:.1f}%)")
    print(f"Score distribution        :")
    for bucket in sorted(score_buckets):
        bar = "█" * int(score_buckets[bucket] * 40 / args.episodes)
        print(f"  {bucket}  {bar} {score_buckets[bucket]}")
    print(f"Spec variant distribution :")
    for variant, n in variant_counts.most_common():
        print(f"  {variant:<28} {n:>5}")

    if not args.stats_only:
        print(f"\nWrote {kept} examples to {out_path}")
    else:
        print("\n(stats-only: no file written)")
    return 0


if __name__ == "__main__":
    sys.exit(main())