finetune/scripts/audit_actor_sft.py

#!/usr/bin/env python3
"""Strict local audit for Actor SFT assistant JSON completions."""

from __future__ import annotations

import argparse
import json
from collections import Counter
from pathlib import Path
from typing import Any


DEFAULT_PATH = Path("finetune/data/actor_sft_v1.jsonl")
REQUIRED_FIELDS = [
    "intent",
    "line",
    "emotion",
    "gesture",
    "stage_effect",
    "memory_update",
    "tool_request",
]
REQUIRED_FIELD_SET = set(REQUIRED_FIELDS)
FORBIDDEN_TOP_LEVEL_FIELDS = {
    "memory_record",
    "memory_effect",
    "recent_transcript",
    "show_state",
    "held_props",
    "mood",
    "name",
    "latest_prop",
    "latest_audience_action",
    "tool_results",
    "status",
    "result",
    "notes",
    "current_show_phase",
}
ALLOWED_TOOLS = {"inspect_prop", "consult_stage_oracle", "change_lighting"}
SUMMARY_ISSUES = [
    "assistant_invalid_json",
    "assistant_non_object_json",
    "assistant_json_does_not_start_with_object",
    "assistant_json_does_not_end_with_object",
    "assistant_contains_markdown",
    "assistant_continuation_detected",
    "assistant_extra_whitespace_before_or_after_json",
    "missing_required_fields",
    "extra_top_level_fields",
    "duplicate_top_level_keys",
    "forbidden_top_level_fields",
    "line_missing_or_empty",
    "tool_request_keys_not_exact",
    "tool_request_not_object_or_null",
    "tool_request_invalid_tool",
    "tool_request_args_not_object",
    "tool_request_args_not_exact_for_inspect_prop",
    "tool_request_args_not_exact_for_consult_stage_oracle",
    "tool_request_args_not_exact_for_change_lighting",
    "tool_request_reason_missing_or_empty",
    "finale_output_outside_finale_context",
    "finale_context_unavailable",
]

# Match the current app runtime schema in puppet_theater/tools.py.
TOOL_ARGS = {
    "inspect_prop": {"prop"},
    "consult_stage_oracle": {"question"},
    "change_lighting": {"mood"},
}


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("path", nargs="?", type=Path, default=DEFAULT_PATH)
    parser.add_argument("--max-examples", type=int, default=20)
    args = parser.parse_args()

    stats = audit_file(args.path, args.max_examples)
    print_summary(args.path, stats)
    if stats["failure_rows"]:
        raise SystemExit(1)


def audit_file(path: Path, max_examples: int) -> dict[str, Any]:
    stats: dict[str, Any] = {
        "total_rows": 0,
        "failure_rows": 0,
        "issue_counts": Counter(),
        "row_type_counts": Counter(),
        "tool_counts": Counter(),
        "examples": [],
    }
    with path.open("r", encoding="utf-8") as handle:
        for line_number, line in enumerate(handle, start=1):
            if not line.strip():
                continue
            stats["total_rows"] += 1
            row_errors: list[str] = []
            try:
                row = json.loads(line)
            except json.JSONDecodeError as exc:
                row_errors.append(f"row_invalid_json: {exc}")
                record_row_errors(stats, max_examples, line_number, None, None, row_errors)
                continue
            row_id = row.get("id") if isinstance(row, dict) else None
            row_type = row.get("row_type") if isinstance(row, dict) else None
            if isinstance(row_type, str):
                stats["row_type_counts"][row_type] += 1
            assistant_content = extract_assistant_content(row, row_errors)
            show_state = extract_show_state(row)
            if assistant_content is None:
                record_row_errors(stats, max_examples, line_number, row_id, row_type, row_errors)
                continue

            row_errors.extend(validate_exact_json_text(assistant_content))
            parsed, duplicates, parse_error = parse_json_object_with_duplicate_keys(assistant_content)
            if parse_error is not None:
                row_errors.append(f"assistant_invalid_json: {parse_error}")
            elif not isinstance(parsed, dict):
                row_errors.append("assistant_non_object_json")
            else:
                row_errors.extend(validate_assistant_object(parsed, duplicates, show_state))
                tool_request = parsed.get("tool_request")
                if isinstance(tool_request, dict):
                    stats["tool_counts"][tool_request.get("tool", "invalid")] += 1
                elif tool_request is None:
                    stats["tool_counts"]["none"] += 1
                else:
                    stats["tool_counts"]["invalid"] += 1
            record_row_errors(stats, max_examples, line_number, row_id, row_type, row_errors)
    return stats


def extract_assistant_content(row: Any, errors: list[str]) -> str | None:
    if not isinstance(row, dict):
        errors.append("row_non_object")
        return None
    messages = row.get("messages")
    if not isinstance(messages, list) or len(messages) < 3:
        errors.append("messages_missing_assistant")
        return None
    assistant_message = messages[2]
    if not isinstance(assistant_message, dict) or assistant_message.get("role") != "assistant":
        errors.append("assistant_message_invalid")
        return None
    content = assistant_message.get("content")
    if not isinstance(content, str):
        errors.append("assistant_content_not_string")
        return None
    return content


def validate_exact_json_text(content: str) -> list[str]:
    errors: list[str] = []
    stripped = content.strip()
    if content != stripped:
        errors.append("assistant_extra_whitespace_before_or_after_json")
    if not stripped.startswith("{"):
        errors.append("assistant_json_does_not_start_with_object")
    if not stripped.endswith("}"):
        errors.append("assistant_json_does_not_end_with_object")
    if "```" in stripped or stripped.startswith("`"):
        errors.append("assistant_contains_markdown")
    lowered = stripped.lower()
    if "\nassistant" in lowered or "### assistant" in lowered or "<|assistant" in lowered:
        errors.append("assistant_continuation_detected")
    return errors


def parse_json_object_with_duplicate_keys(content: str) -> tuple[Any, list[str], str | None]:
    duplicate_keys: list[str] = []

    def hook(pairs: list[tuple[str, Any]]) -> dict[str, Any]:
        seen: set[str] = set()
        for key, _value in pairs:
            if key in seen and key not in duplicate_keys:
                duplicate_keys.append(key)
            seen.add(key)
        return dict(pairs)

    try:
        return json.loads(content, object_pairs_hook=hook), duplicate_keys, None
    except json.JSONDecodeError as exc:
        return None, duplicate_keys, str(exc)


def validate_assistant_object(value: dict[str, Any], duplicate_keys: list[str], show_state: dict[str, Any] | None) -> list[str]:
    errors: list[str] = []
    keys = set(value)
    missing = sorted(REQUIRED_FIELD_SET - keys)
    extra = sorted(keys - REQUIRED_FIELD_SET)
    forbidden = sorted(FORBIDDEN_TOP_LEVEL_FIELDS & keys)
    if missing:
        errors.append(f"missing_required_fields={missing}")
    if extra:
        errors.append(f"extra_top_level_fields={extra}")
    if duplicate_keys:
        errors.append(f"duplicate_top_level_keys={sorted(duplicate_keys)}")
    if forbidden:
        errors.append(f"forbidden_top_level_fields={forbidden}")
    if "line" in value and (not isinstance(value["line"], str) or not value["line"].strip()):
        errors.append("line_missing_or_empty")
    errors.extend(validate_tool_request(value.get("tool_request")))
    if value.get("intent") == "deliver_finale" or value.get("stage_effect") == "final_bow_lights":
        errors.extend(validate_finale_context(show_state))
    return errors


def validate_tool_request(value: Any) -> list[str]:
    if value is None:
        return []
    if not isinstance(value, dict):
        return ["tool_request_not_object_or_null"]
    errors: list[str] = []
    keys = set(value)
    if keys != {"tool", "args", "reason"}:
        errors.append(f"tool_request_keys_not_exact={sorted(keys)}")
    tool = value.get("tool")
    if tool not in ALLOWED_TOOLS:
        errors.append(f"tool_request_invalid_tool={tool!r}")
        return errors
    args = value.get("args")
    if not isinstance(args, dict):
        errors.append("tool_request_args_not_object")
        return errors
    expected_args = TOOL_ARGS[tool]
    if set(args) != expected_args:
        errors.append(f"tool_request_args_not_exact_for_{tool}: expected={sorted(expected_args)} got={sorted(args)}")
    for key, arg_value in args.items():
        if not isinstance(arg_value, str) or not arg_value.strip():
            errors.append(f"tool_request_arg_invalid={key}")
    reason = value.get("reason")
    if not isinstance(reason, str) or not reason.strip():
        errors.append("tool_request_reason_missing_or_empty")
    return errors


def validate_finale_context(show_state: dict[str, Any] | None) -> list[str]:
    if show_state is None:
        return ["finale_context_unavailable"]
    if show_state.get("story_phase") == "finale" or show_state.get("finale_requested") is True:
        return []
    return ["finale_output_outside_finale_context"]


def extract_show_state(row: Any) -> dict[str, Any] | None:
    if not isinstance(row, dict):
        return None
    messages = row.get("messages")
    if not isinstance(messages, list) or len(messages) < 2 or not isinstance(messages[1], dict):
        return None
    content = messages[1].get("content")
    if not isinstance(content, str):
        return None
    marker = "show_state JSON:"
    next_marker = "\nactor JSON:"
    if marker not in content:
        return None
    start = content.index(marker) + len(marker)
    end = content.find(next_marker, start)
    raw_json = content[start:end if end != -1 else None].strip()
    try:
        value = json.loads(raw_json)
    except json.JSONDecodeError:
        return None
    return value if isinstance(value, dict) else None


def record_row_errors(
    stats: dict[str, Any],
    max_examples: int,
    line_number: int,
    row_id: Any,
    row_type: Any,
    errors: list[str],
) -> None:
    if not errors:
        return
    stats["failure_rows"] += 1
    for error in errors:
        stats["issue_counts"][issue_name(error)] += 1
    if len(stats["examples"]) < max_examples:
        stats["examples"].append(
            {
                "line_number": line_number,
                "id": row_id,
                "row_type": row_type,
                "errors": errors,
            }
        )


def issue_name(error: str) -> str:
    return error.split("=", 1)[0].split(":", 1)[0]


def print_summary(path: Path, stats: dict[str, Any]) -> None:
    print(f"file: {path}")
    print(f"total rows: {stats['total_rows']}")
    print(f"strict failure rows: {stats['failure_rows']}")
    print_distribution("row_type distribution", stats["row_type_counts"])
    print_distribution("tool_request distribution", stats["tool_counts"])
    print_issue_distribution(stats["issue_counts"])
    if stats["examples"]:
        print("examples:")
        for example in stats["examples"]:
            print(
                f"- line {example['line_number']} id={example['id']} "
                f"row_type={example['row_type']}: {'; '.join(example['errors'])}"
            )


def print_distribution(title: str, values: Counter) -> None:
    print(f"{title}:")
    if not values:
        print("  none")
        return
    for key, count in sorted(values.items()):
        print(f"  {key}: {count}")


def print_issue_distribution(values: Counter) -> None:
    print("issue distribution:")
    for key in SUMMARY_ISSUES:
        print(f"  {key}: {values.get(key, 0)}")
    extra_keys = sorted(key for key in values if key not in SUMMARY_ISSUES)
    for key in extra_keys:
        print(f"  {key}: {values[key]}")


if __name__ == "__main__":
    main()