@pentatonic-ai/ai-agent-sdk 0.10.4 → 0.10.6

package/packages/memory-engine-v2/scripts/bakeoff_guided_vs_kv.py

@@ -0,0 +1,607 @@
+#!/usr/bin/env python3
+"""bakeoff_guided_vs_kv — compare KV-text vs guided-JSON distiller output.
+
+Replay/validation harness for DISTILL_OUTPUT_MODE. NO live LLM calls by
+default: without --endpoint the script only re-parses existing outputs
+(from distillation_traces or a JSONL file) and reports parse-quality
+metrics. Live mode (both output formats against a vLLM box) requires an
+explicit --endpoint.
+
+Input sources (exactly one):
+  --input PATH      JSONL of sample events. Each line is an event dict:
+                    {"id", "source_kind", "content", "attributes", ...}.
+                    Optional replay fields: "raw_response" (KV-format
+                    model output) and/or "raw_response_guided"
+                    (guided-JSON model output) — when present they are
+                    parsed offline with the matching parser.
+  --from-traces     Pull events + raw KV responses from
+                    distillation_traces (joined to events). user_prompt
+                    is exactly build_event_block() output, so it doubles
+                    as the live-replay input; raw_response gives
+                    KV-format parse stats for free. Connection from
+                    --pg-dsn or PG_DSN env.
+
+Eval-set hygiene (ALWAYS applied — mirrors worker.claim_next_batch's
+claim-time pre-filters so the eval distribution matches what the
+distiller actually processes):
+  - source_kind == 'agent' excluded (never distil the agent's own output)
+  - attributes.source in {seesa, claude-code-plugin, pip-code-ingest}
+    or prefixed openclaw-/triage-/briefing- excluded
+  - bytes-garbage excluded (> 10% U+FFFD replacement chars)
+
+Per-mode metrics:
+  - parse yield: ENT / FCT / REL counts (total + per-event mean)
+  - dropped-record rate (kv: record-shaped lines the parser rejected;
+    guided: schema violations reported by validate_payload + events
+    lost to truncation salvage)
+  - FCT field-count violations (kv: FCT| lines without exactly 6 segments)
+  - pipe-debris instances (a literal `|` inside any parsed field value)
+  - subject-not-in-ENT rate (facts whose subject isn't a declared entity)
+  - junk-entity rate (noise_filter.is_noise_entity_name)
+  - groundedness (entity-name substring spot-check against the source)
+  - latency (live calls; or distillation_traces.llm_chunk_ms offline)
+
+Outputs: <out-dir>/report.md + <out-dir>/records.jsonl
+
+Examples:
+  # offline — replay 500 recent traces, KV parse stats only
+  python bakeoff_guided_vs_kv.py --from-traces --limit 500
+
+  # offline — replay a JSONL that already contains both raw outputs
+  python bakeoff_guided_vs_kv.py --input samples.jsonl
+
+  # LIVE — run both formats against the 7B box (explicit opt-in)
+  python bakeoff_guided_vs_kv.py --from-traces --limit 100 \
+      --endpoint http://172.31.91.6:8005/v1/chat/completions \
+      --model Qwen/Qwen2.5-7B-Instruct
+"""
+
+from __future__ import annotations
+
+import argparse
+import importlib.util
+import json
+import os
+import re
+import statistics
+import sys
+import time
+from pathlib import Path
+from typing import Any
+
+# ----------------------------------------------------------------------
+# Load the worker + schema modules from extractor-async (flat scripts
+# dir, no package) the same way the unit tests do.
+# ----------------------------------------------------------------------
+
+_EXTRACTOR_DIR = Path(__file__).resolve().parent.parent / "extractor-async"
+sys.path.insert(0, str(_EXTRACTOR_DIR))
+
+
+def _load(name: str):
+    spec = importlib.util.spec_from_file_location(name, _EXTRACTOR_DIR / f"{name}.py")
+    assert spec and spec.loader
+    mod = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(mod)
+    return mod
+
+
+worker = _load("worker")
+extraction_schema = _load("extraction_schema")
+noise_filter = _load("noise_filter")
+
+
+# ----------------------------------------------------------------------
+# Eval-set hygiene — mirrors worker.claim_next_batch's pre-filters
+# (source_kind=agent skip, skip-sources, bytes-garbage ratio) plus the
+# known agent-adjacent sources that pollute eval sets.
+# ----------------------------------------------------------------------
+
+EXCLUDE_SOURCES_EXACT = {"seesa", "claude-code-plugin", "pip-code-ingest"}
+EXCLUDE_SOURCE_PREFIXES = ("openclaw-", "triage-", "briefing-")
+GARBAGE_CHAR_RATIO = 0.10  # mirrors DISTILL_GARBAGE_CHAR_RATIO default
+
+
+def is_eval_eligible(event: dict[str, Any]) -> tuple[bool, str]:
+    """(eligible, reason-if-not). Same spirit as claim_next_batch."""
+    if (event.get("source_kind") or "") == "agent":
+        return False, "source_kind=agent"
+    attrs = event.get("attributes") or {}
+    if isinstance(attrs, str):
+        try:
+            attrs = json.loads(attrs)
+        except json.JSONDecodeError:
+            attrs = {}
+    src = str(attrs.get("source") or "")
+    if src in EXCLUDE_SOURCES_EXACT or src.startswith(EXCLUDE_SOURCE_PREFIXES):
+        return False, f"source={src}"
+    content = event.get("content") or ""
+    if content and content.count("�") / len(content) > GARBAGE_CHAR_RATIO:
+        return False, "bytes_garbage"
+    return True, ""
+
+
+# ----------------------------------------------------------------------
+# Per-event metric extraction
+# ----------------------------------------------------------------------
+
+_KV_RECORD_RE = re.compile(r"^(?:[-*]\s+)?(ENT|FCT|REL)\|")
+
+
+def _kv_raw_stats(raw: str) -> dict[str, int]:
+    """Raw-text stats the parser can't see: candidate record lines and
+    FCT lines whose segment count isn't exactly 6."""
+    candidates = 0
+    fct_field_violations = 0
+    for line in raw.splitlines():
+        line = line.strip()
+        m = _KV_RECORD_RE.match(line)
+        if not m:
+            continue
+        candidates += 1
+        if m.group(1) == "FCT" and len(line.split("|")) != 6:
+            fct_field_violations += 1
+    return {"candidates": candidates, "fct_field_violations": fct_field_violations}
+
+
+def _iter_parsed_values(rec: dict[str, Any]):
+    for e in rec.get("entities", []):
+        yield e.get("name") or ""
+        for a in e.get("aliases") or []:
+            yield a
+    for f in rec.get("facts", []):
+        for k in ("subject", "predicate", "object", "statement"):
+            yield f.get(k) or ""
+    for r in rec.get("relationships", []):
+        for k in ("from", "to", "type"):
+            yield r.get(k) or ""
+
+
+def score_event(
+    rec: dict[str, Any], source_text: str
+) -> dict[str, Any]:
+    """Quality metrics for one parsed per-event extraction dict."""
+    ents = rec.get("entities", [])
+    facts = rec.get("facts", [])
+    rels = rec.get("relationships", [])
+    ent_names = {e.get("name", "") for e in ents}
+    src_lower = (source_text or "").lower()
+
+    junk = sum(
+        1 for e in ents
+        if noise_filter.is_noise_entity_name(e.get("type", ""), e.get("name", ""))
+    )
+    # Groundedness spot-check: the entity surface form should literally
+    # appear in the source (case-insensitive). Substring is deliberately
+    # crude — it overcounts grounding for short names, undercounts for
+    # inflected forms — but it's model-free and symmetric across modes.
+    ungrounded = sum(
+        1 for e in ents
+        if e.get("name") and e.get("name", "").lower() not in src_lower
+    )
+    subject_violations = sum(
+        1 for f in facts if f.get("subject") and f.get("subject") not in ent_names
+    )
+    pipe_debris = sum(1 for v in _iter_parsed_values(rec) if "|" in str(v))
+    return {
+        "ent": len(ents),
+        "fct": len(facts),
+        "rel": len(rels),
+        "junk_entities": junk,
+        "ungrounded_entities": ungrounded,
+        "subject_not_in_ent": subject_violations,
+        "pipe_debris": pipe_debris,
+    }
+
+
+# ----------------------------------------------------------------------
+# Aggregation
+# ----------------------------------------------------------------------
+
+
+class ModeStats:
+    def __init__(self, mode: str) -> None:
+        self.mode = mode
+        self.events = 0
+        self.totals: dict[str, int] = {
+            "ent": 0, "fct": 0, "rel": 0,
+            "junk_entities": 0, "ungrounded_entities": 0,
+            "subject_not_in_ent": 0, "pipe_debris": 0,
+            "dropped_records": 0, "fct_field_violations": 0,
+            "schema_violations": 0, "events_lost_to_truncation": 0,
+        }
+        self.latencies_ms: list[float] = []
+
+    def add(self, scores: dict[str, Any]) -> None:
+        self.events += 1
+        for k, v in scores.items():
+            if k in self.totals:
+                self.totals[k] += v
+
+    def rate(self, key: str, denom_key: str = "ent") -> float:
+        denom = self.totals[denom_key]
+        return (self.totals[key] / denom) if denom else 0.0
+
+    def summary(self) -> dict[str, Any]:
+        t = self.totals
+        out: dict[str, Any] = {"mode": self.mode, "events": self.events, **t}
+        out["junk_entity_rate"] = round(self.rate("junk_entities"), 4)
+        out["ungrounded_entity_rate"] = round(self.rate("ungrounded_entities"), 4)
+        out["subject_not_in_ent_rate"] = round(
+            self.rate("subject_not_in_ent", "fct"), 4
+        )
+        if self.latencies_ms:
+            out["latency_ms_p50"] = round(statistics.median(self.latencies_ms), 1)
+            out["latency_ms_mean"] = round(statistics.fmean(self.latencies_ms), 1)
+        return out
+
+
+# ----------------------------------------------------------------------
+# Offline replay
+# ----------------------------------------------------------------------
+
+
+def replay_kv(raw: str, source_text: str, stats: ModeStats) -> dict[str, Any]:
+    parsed = worker._parse_kv_records(raw, expected_n=1)[0]
+    scores = score_event(parsed, source_text)
+    raw_stats = _kv_raw_stats(raw)
+    parsed_count = scores["ent"] + scores["fct"] + scores["rel"]
+    scores["dropped_records"] = max(0, raw_stats["candidates"] - parsed_count)
+    scores["fct_field_violations"] = raw_stats["fct_field_violations"]
+    stats.add(scores)
+    return scores
+
+
+def replay_guided(raw: str, source_text: str, stats: ModeStats) -> dict[str, Any]:
+    parsed = worker._parse_guided_json(raw, expected_n=1)[0]
+    scores = score_event(parsed, source_text)
+    payload = worker._load_guided_payload(raw)
+    if payload is None:
+        scores["events_lost_to_truncation"] = 1
+        scores["schema_violations"] = 0
+    else:
+        scores["events_lost_to_truncation"] = (
+            0 if isinstance(payload.get("events"), list) and payload["events"] else 1
+        )
+        scores["schema_violations"] = len(extraction_schema.validate_payload(payload))
+    stats.add(scores)
+    return scores
+
+
+# ----------------------------------------------------------------------
+# Live mode (explicit --endpoint only)
+# ----------------------------------------------------------------------
+
+_EVENT_HEADER_LINE_RE = re.compile(r"^\[event \d+\]", re.MULTILINE)
+
+
+def _build_live_body(
+    mode: str, user_prompt: str, n: int, model: str, param_style: str
+) -> dict[str, Any]:
+    if mode == "kv":
+        return {
+            "model": model,
+            "messages": [
+                {"role": "system", "content": worker.BATCH_SYSTEM_PROMPT},
+                {"role": "user", "content": user_prompt},
+            ],
+            "temperature": 0.0,
+            "max_tokens": worker.LLM_MAX_TOKENS_PER_EVENT * n,
+        }
+    body: dict[str, Any] = {
+        "model": model,
+        "messages": [
+            {"role": "system", "content": worker.GUIDED_JSON_SYSTEM_PROMPT},
+            {"role": "user", "content": user_prompt},
+        ],
+        "temperature": 0.0,
+        "max_tokens": worker.LLM_MAX_TOKENS_PER_EVENT_JSON * n,
+    }
+    if param_style == "guided_json":
+        body["guided_json"] = extraction_schema.EXTRACTION_SCHEMA
+    else:
+        body["response_format"] = {
+            "type": "json_schema",
+            "json_schema": {
+                "name": "memory_extraction",
+                "strict": True,
+                "schema": extraction_schema.EXTRACTION_SCHEMA,
+            },
+        }
+    return body
+
+
+def run_live(
+    samples: list[dict[str, Any]],
+    endpoint: str,
+    model: str,
+    api_key: str,
+    param_style: str,
+    batch_size: int,
+    timeout: float,
+    kv_stats: ModeStats,
+    guided_stats: ModeStats,
+    records: list[dict[str, Any]],
+) -> None:
+    import httpx  # already a worker dep
+
+    headers = {"Content-Type": "application/json"}
+    if api_key:
+        headers["X-API-Key"] = api_key
+        headers["Authorization"] = f"Bearer {api_key}"
+
+    with httpx.Client(timeout=timeout) as client:
+        for start in range(0, len(samples), batch_size):
+            chunk = samples[start : start + batch_size]
+            n = len(chunk)
+            blocks = []
+            for i, s in enumerate(chunk):
+                if s.get("user_prompt"):
+                    # Trace replay: user_prompt IS build_event_block
+                    # output; renumber the header to the chunk-local
+                    # index so parsing reattaches correctly.
+                    block = _EVENT_HEADER_LINE_RE.sub(
+                        f"[event {i}]", s["user_prompt"], count=1
+                    )
+                else:
+                    block = worker.build_event_block(i, s)
+                blocks.append(block)
+            user_prompt = "\n\n---\n\n".join(blocks)
+
+            for mode, stats in (("kv", kv_stats), ("guided_json", guided_stats)):
+                body = _build_live_body(mode, user_prompt, n, model, param_style)
+                t0 = time.perf_counter()
+                r = client.post(endpoint, json=body, headers=headers)
+                r.raise_for_status()
+                ms = (time.perf_counter() - t0) * 1000
+                stats.latencies_ms.append(ms)
+                data = r.json()
+                text = (
+                    (data.get("choices") or [{}])[0]
+                    .get("message", {})
+                    .get("content", "")
+                ) or ""
+                if mode == "kv":
+                    parsed = worker._parse_kv_records(text, n)
+                    raw_stats = _kv_raw_stats(text)
+                else:
+                    parsed = worker._parse_guided_json(text, n)
+                    payload = worker._load_guided_payload(text)
+                for i, (s, rec) in enumerate(zip(chunk, parsed)):
+                    src = s.get("user_prompt") or s.get("content") or ""
+                    scores = score_event(rec, src)
+                    if mode == "kv":
+                        # Chunk-level raw stats land on the first event of
+                        # the chunk (they aren't attributable per-event).
+                        if i == 0:
+                            parsed_count = sum(
+                                len(p.get("entities", []))
+                                + len(p.get("facts", []))
+                                + len(p.get("relationships", []))
+                                for p in parsed
+                            )
+                            scores["dropped_records"] = max(
+                                0, raw_stats["candidates"] - parsed_count
+                            )
+                            scores["fct_field_violations"] = raw_stats[
+                                "fct_field_violations"
+                            ]
+                    else:
+                        if i == 0:
+                            if payload is None:
+                                scores["events_lost_to_truncation"] = n
+                            else:
+                                got = (
+                                    len(payload.get("events", []))
+                                    if isinstance(payload.get("events"), list)
+                                    else 0
+                                )
+                                scores["events_lost_to_truncation"] = max(0, n - got)
+                                scores["schema_violations"] = len(
+                                    extraction_schema.validate_payload(payload)
+                                )
+                    stats.add(scores)
+                    records.append(
+                        {
+                            "event_id": s.get("id") or s.get("event_id"),
+                            "mode": mode,
+                            "live": True,
+                            "latency_ms_chunk": round(ms, 1),
+                            **scores,
+                        }
+                    )
+
+
+# ----------------------------------------------------------------------
+# Input loading
+# ----------------------------------------------------------------------
+
+TRACES_SQL = """
+SELECT t.event_id, t.user_prompt, t.raw_response, t.llm_chunk_ms,
+       t.system_prompt_hash,
+       e.source_kind, e.content, e.attributes
+FROM distillation_traces t
+JOIN events e ON e.id = t.event_id
+ORDER BY t.created_at DESC
+LIMIT %s
+"""
+
+
+def load_from_traces(pg_dsn: str, limit: int) -> tuple[list[dict[str, Any]], dict[str, int]]:
+    import psycopg
+    import psycopg.rows
+
+    excluded: dict[str, int] = {}
+    samples: list[dict[str, Any]] = []
+    # Over-fetch so hygiene filtering still leaves ~limit rows.
+    with psycopg.connect(pg_dsn) as conn:
+        with conn.cursor(row_factory=psycopg.rows.dict_row) as cur:
+            cur.execute(TRACES_SQL, (limit * 5,))
+            for row in cur:
+                ok, reason = is_eval_eligible(row)
+                if not ok:
+                    excluded[reason] = excluded.get(reason, 0) + 1
+                    continue
+                samples.append(dict(row))
+                if len(samples) >= limit:
+                    break
+    return samples, excluded
+
+
+def load_from_jsonl(path: Path, limit: int) -> tuple[list[dict[str, Any]], dict[str, int]]:
+    excluded: dict[str, int] = {}
+    samples: list[dict[str, Any]] = []
+    with path.open() as fh:
+        for line in fh:
+            line = line.strip()
+            if not line:
+                continue
+            ev = json.loads(line)
+            ok, reason = is_eval_eligible(ev)
+            if not ok:
+                excluded[reason] = excluded.get(reason, 0) + 1
+                continue
+            samples.append(ev)
+            if len(samples) >= limit:
+                break
+    return samples, excluded
+
+
+# ----------------------------------------------------------------------
+# Report
+# ----------------------------------------------------------------------
+
+
+def write_report(
+    out_dir: Path,
+    kv_stats: ModeStats,
+    guided_stats: ModeStats,
+    excluded: dict[str, int],
+    records: list[dict[str, Any]],
+    live: bool,
+) -> None:
+    out_dir.mkdir(parents=True, exist_ok=True)
+    with (out_dir / "records.jsonl").open("w") as fh:
+        for rec in records:
+            fh.write(json.dumps(rec, ensure_ascii=False) + "\n")
+
+    rows = [s.summary() for s in (kv_stats, guided_stats) if s.events]
+    keys = [
+        "events", "ent", "fct", "rel",
+        "dropped_records", "fct_field_violations",
+        "schema_violations", "events_lost_to_truncation",
+        "pipe_debris", "junk_entity_rate", "ungrounded_entity_rate",
+        "subject_not_in_ent_rate", "latency_ms_p50", "latency_ms_mean",
+    ]
+    lines = [
+        "# Guided-JSON vs KV distiller bake-off",
+        "",
+        f"- generated: {time.strftime('%Y-%m-%d %H:%M:%S %Z')}",
+        f"- mode: {'LIVE (called --endpoint)' if live else 'offline replay (no LLM calls)'}",
+        f"- eval-set hygiene exclusions: {json.dumps(excluded) if excluded else 'none'}",
+        "",
+        "| metric | " + " | ".join(r["mode"] for r in rows) + " |",
+        "|---|" + "---|" * len(rows),
+    ]
+    for k in keys:
+        vals = [str(r.get(k, "—")) for r in rows]
+        lines.append(f"| {k} | " + " | ".join(vals) + " |")
+    lines += [
+        "",
+        "Notes:",
+        "- `dropped_records` / `fct_field_violations` / `pipe_debris` are",
+        "  KV-format failure classes; guided decoding makes them",
+        "  structurally impossible (`schema_violations` /",
+        "  `events_lost_to_truncation` are the guided-mode analogues).",
+        "- `ungrounded_entity_rate` is a crude substring spot-check of",
+        "  entity names against the source text — compare across modes,",
+        "  don't read it as an absolute hallucination rate.",
+    ]
+    (out_dir / "report.md").write_text("\n".join(lines) + "\n")
+    print("\n".join(lines))
+    print(f"\nwrote {out_dir / 'report.md'} and {out_dir / 'records.jsonl'}")
+
+
+# ----------------------------------------------------------------------
+# Main
+# ----------------------------------------------------------------------
+
+
+def main() -> int:
+    ap = argparse.ArgumentParser(description=__doc__.split("\n", 1)[0])
+    src = ap.add_mutually_exclusive_group(required=True)
+    src.add_argument("--input", type=Path, help="JSONL of sample events")
+    src.add_argument(
+        "--from-traces", action="store_true",
+        help="pull events + raw KV responses from distillation_traces",
+    )
+    ap.add_argument("--pg-dsn", default=os.environ.get("PG_DSN", ""),
+                    help="Postgres DSN for --from-traces (default: PG_DSN env)")
+    ap.add_argument("--limit", type=int, default=200)
+    ap.add_argument("--endpoint", default="",
+                    help="chat-completions URL. REQUIRED for live mode; "
+                         "omit for offline replay (no LLM calls).")
+    ap.add_argument("--model", default=os.environ.get("LLM_MODEL", "Qwen/Qwen2.5-7B-Instruct"))
+    ap.add_argument("--api-key", default=os.environ.get("LLM_API_KEY", ""))
+    ap.add_argument("--param-style", choices=["response_format", "guided_json"],
+                    default="response_format",
+                    help="how the schema is attached in guided live calls")
+    ap.add_argument("--batch-size", type=int, default=15,
+                    help="events per live LLM call (mirrors EVENTS_PER_LLM_CALL)")
+    ap.add_argument("--timeout", type=float, default=180.0)
+    ap.add_argument("--out-dir", type=Path, default=Path("bakeoff-out"))
+    args = ap.parse_args()
+
+    if args.from_traces:
+        if not args.pg_dsn:
+            ap.error("--from-traces needs --pg-dsn or PG_DSN env")
+        samples, excluded = load_from_traces(args.pg_dsn, args.limit)
+    else:
+        samples, excluded = load_from_jsonl(args.input, args.limit)
+
+    if not samples:
+        print("no eligible samples after hygiene filtering", file=sys.stderr)
+        return 1
+    print(f"{len(samples)} eligible samples "
+          f"({sum(excluded.values())} excluded by hygiene filters)")
+
+    kv_stats = ModeStats("kv")
+    guided_stats = ModeStats("guided_json")
+    records: list[dict[str, Any]] = []
+
+    # Offline replay over whatever raw outputs the input carries.
+    for s in samples:
+        source_text = s.get("user_prompt") or s.get("content") or ""
+        raw_kv = s.get("raw_response")
+        if raw_kv:
+            scores = replay_kv(raw_kv, source_text, kv_stats)
+            if s.get("llm_chunk_ms"):
+                kv_stats.latencies_ms.append(float(s["llm_chunk_ms"]))
+            records.append({"event_id": s.get("event_id") or s.get("id"),
+                            "mode": "kv", "live": False, **scores})
+        raw_guided = s.get("raw_response_guided")
+        if raw_guided:
+            scores = replay_guided(raw_guided, source_text, guided_stats)
+            records.append({"event_id": s.get("event_id") or s.get("id"),
+                            "mode": "guided_json", "live": False, **scores})
+
+    if args.endpoint:
+        run_live(
+            samples, args.endpoint, args.model, args.api_key,
+            args.param_style, args.batch_size, args.timeout,
+            kv_stats, guided_stats, records,
+        )
+    elif not guided_stats.events:
+        print(
+            "note: no guided-JSON outputs in the input and no --endpoint — "
+            "guided column will be empty. Pass --endpoint for a live "
+            "side-by-side run.",
+        )
+
+    write_report(args.out_dir, kv_stats, guided_stats, excluded, records,
+                 live=bool(args.endpoint))
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())