inference.py

#!/usr/bin/env python3
"""Hackathon baseline inference script for TalkingHeadBench.

This script:
- Loads benchmark cases from tests/test_set/.
- Connects to a running TalkingHeadBench OpenEnv server.
- Runs one episode for each task tier: image, clips, weights.
- Uses the OpenAI Python client for all LLM calls.
- Prints per-tier scores and a weighted final score.

Required environment variables:
- API_BASE_URL
- MODEL_NAME
- HF_TOKEN
"""

from __future__ import annotations

import argparse
import json
import os
import re
import sys
from dataclasses import dataclass
from pathlib import Path
from typing import Any

from openai import OpenAI
from pydantic import ValidationError

from client import TalkingHeadBenchEnv
from models import ImageDiagnosticsAction, ParamAnomalyAction, PhonemeRiskAction


PROJECT_ROOT = Path(__file__).resolve().parent
TEST_SET_DIR = PROJECT_ROOT / "tests" / "test_set"

TEST_SET_FILES: dict[str, Path] = {
    "image": TEST_SET_DIR / "subenv1_cases.json",
    "clips": TEST_SET_DIR / "subenv2_cases.json",
    "weights": TEST_SET_DIR / "subenv3_cases.json",
}

SYSTEM_PROMPT = (
    "You are a senior diagnostic engineer for TalkingHeadBench. "
    "Return only strict JSON matching the expected action schema. "
    "Never prescribe absolute numeric parameter values; only directional advice."
)

ACTION_SCHEMA_SUMMARY = """Step 0 (ImageDiagnosticsAction):
- regime_classification: frontal_simple|non_frontal|complex_background|occluded|low_quality
- identified_risk_factors: list[str]
- prompt_issues: list[str]
- recommended_prompt_modifications: list[str]
- image_usability_score: float in [0,1]
- reasoning: str

Step 1 (ParamAnomalyAction):
- config_risk_level: safe|marginal|risky|dangerous
- anomalies: list[{parameter, issue, severity, linked_failure_mode}]
- predicted_failure_modes: list[str]
- directional_fixes: list[{target, direction, rationale, priority}]
- summary: str

Step 2 (PhonemeRiskAction):
- phoneme_risk_ranking: list[{phoneme, risk_score, risk_type, confidence, evidence}]
- predicted_behavior_triggers: list[{trigger_phoneme, triggered_behavior, association_strength, is_intended, concern_level}]
- risky_phoneme_clusters: list[{phonemes, cluster_risk_type, combined_risk_score, interaction_description}]
- model_behavioral_safety: safe|minor_concerns|moderate_risk|high_risk|unsafe
- mitigation_recommendations: list[{target, action, rationale, priority}]
- summary: str
"""

SCHEMA_MODELS = {
    "ImageDiagnosticsAction": ImageDiagnosticsAction,
    "ParamAnomalyAction": ParamAnomalyAction,
    "PhonemeRiskAction": PhonemeRiskAction,
}

STEP_SCHEMAS = [
    "ImageDiagnosticsAction",
    "ParamAnomalyAction",
    "PhonemeRiskAction",
]


@dataclass
class TierRunResult:
    tier: str
    case_id: str | None
    reward: float
    scores: dict[str, Any]
    steps: int
    used_custom_bundle: bool


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Run TalkingHeadBench baseline inference.")
    parser.add_argument(
        "--env-url",
        default=os.getenv("THB_ENV_URL", "http://localhost:8000"),
        help="TalkingHeadBench server base URL (default: THB_ENV_URL or http://localhost:8000)",
    )
    return parser.parse_args()


def require_env(name: str) -> str:
    value = os.getenv(name, "").strip()
    if not value:
        raise SystemExit(f"Missing required environment variable: {name}")
    return value


def load_test_set_cases(path: Path) -> list[dict[str, Any]]:
    if not path.exists():
        raise SystemExit(f"Missing test-set file: {path}")

    try:
        payload = json.loads(path.read_text(encoding="utf-8"))
    except json.JSONDecodeError as exc:
        raise SystemExit(f"Invalid JSON in {path}: {exc}") from exc

    cases = payload.get("cases") if isinstance(payload, dict) else None
    if not isinstance(cases, list):
        raise SystemExit(f"Invalid test-set format in {path}: expected top-level 'cases' list")
    if not cases:
        raise SystemExit(f"Test-set file has no cases: {path}")
    return cases


def _coerce_to_dict(value: Any) -> dict[str, Any]:
    if value is None:
        return {}
    if isinstance(value, dict):
        return value
    if hasattr(value, "model_dump"):
        dumped = value.model_dump()
        return dumped if isinstance(dumped, dict) else {"value": dumped}
    if hasattr(value, "dict"):
        dumped = value.dict()
        return dumped if isinstance(dumped, dict) else {"value": dumped}
    if hasattr(value, "__dict__"):
        return {k: v for k, v in vars(value).items() if not k.startswith("_")}
    return {"value": str(value)}


def unpack_step_result(result: Any) -> tuple[dict[str, Any], bool, float | None, dict[str, Any] | None]:
    raw_observation = getattr(result, "observation", result)
    obs = _coerce_to_dict(raw_observation)

    done = bool(getattr(result, "done", obs.get("done", False)))
    reward = getattr(result, "reward", obs.get("reward"))
    scores = getattr(result, "scores", obs.get("scores"))

    if not isinstance(scores, dict):
        scores = _coerce_to_dict(scores) if scores is not None else None

    if obs.get("scores") is None and scores:
        obs["scores"] = scores

    return obs, done, reward, scores


def expected_schema_name(step_index: int, observation: dict[str, Any]) -> str:
    schema = observation.get("expected_action_schema")
    if isinstance(schema, str) and schema:
        return schema
    if 0 <= step_index < len(STEP_SCHEMAS):
        return STEP_SCHEMAS[step_index]
    return "unknown"


def minimal_action(step_index: int) -> dict[str, Any]:
    if step_index == 0:
        return {
            "regime_classification": "frontal_simple",
            "identified_risk_factors": [],
            "prompt_issues": [],
            "recommended_prompt_modifications": [],
            "image_usability_score": 0.5,
            "reasoning": "Fallback action due to parsing failure.",
        }
    if step_index == 1:
        return {
            "config_risk_level": "marginal",
            "anomalies": [],
            "predicted_failure_modes": [],
            "directional_fixes": [],
            "summary": "Fallback action due to parsing failure.",
        }
    if step_index == 2:
        return {
            "phoneme_risk_ranking": [],
            "predicted_behavior_triggers": [],
            "risky_phoneme_clusters": [],
            "model_behavioral_safety": "minor_concerns",
            "mitigation_recommendations": [],
            "summary": "Fallback action due to parsing failure.",
        }
    return {}


def _extract_json_candidate(text: str) -> dict[str, Any]:
    stripped = text.strip()

    if stripped.startswith("```"):
        stripped = re.sub(r"^```(?:json)?\s*", "", stripped)
        stripped = re.sub(r"\s*```$", "", stripped)

    try:
        payload = json.loads(stripped)
        if isinstance(payload, dict):
            return payload
    except json.JSONDecodeError:
        pass

    match = re.search(r"\{.*\}", stripped, flags=re.DOTALL)
    if not match:
        raise ValueError("No JSON object found in model response")

    payload = json.loads(match.group(0))
    if not isinstance(payload, dict):
        raise ValueError("Parsed JSON is not an object")
    return payload


def validate_action_payload(schema_name: str, payload: dict[str, Any]) -> dict[str, Any]:
    model_cls = SCHEMA_MODELS.get(schema_name)
    if model_cls is None:
        return payload
    return model_cls.model_validate(payload).model_dump()


def call_chat_completion(
    *,
    client: OpenAI,
    model_name: str,
    messages: list[dict[str, str]],
) -> str:
    response = client.chat.completions.create(
        model=model_name,
        messages=messages,
        temperature=0.2,
        max_tokens=1200,
    )

    choices = response.choices or []
    if not choices:
        raise ValueError("Model response did not include choices")

    content = choices[0].message.content
    if not isinstance(content, str) or not content.strip():
        raise ValueError("Model response did not include text content")

    return content.strip()


def build_user_prompt(
    *,
    tier: str,
    case_id: str | None,
    step_index: int,
    schema_name: str,
    observation: dict[str, Any],
) -> str:
    instruction = observation.get("instruction", "")
    node = observation.get("node", "unknown")
    signals = observation.get("signals", {})

    return (
        f"Task tier: {tier}\n"
        f"Case id: {case_id or 'unknown'}\n"
        f"Environment step index: {step_index}\n"
        f"Current node: {node}\n"
        f"Expected action schema: {schema_name}\n\n"
        f"Action schema summary:\n{ACTION_SCHEMA_SUMMARY}\n"
        f"Instruction:\n{instruction}\n\n"
        f"Signals (JSON):\n{json.dumps(signals, indent=2, sort_keys=True)}\n\n"
        f"Return only one JSON object that satisfies {schema_name}."
    )


def generate_action(
    *,
    client: OpenAI,
    model_name: str,
    tier: str,
    case_id: str | None,
    step_index: int,
    schema_name: str,
    observation: dict[str, Any],
) -> dict[str, Any]:
    user_prompt = build_user_prompt(
        tier=tier,
        case_id=case_id,
        step_index=step_index,
        schema_name=schema_name,
        observation=observation,
    )
    messages = [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": user_prompt},
    ]

    last_error = "unknown"
    last_raw = ""

    for attempt in range(2):
        raw = call_chat_completion(client=client, model_name=model_name, messages=messages)
        last_raw = raw

        try:
            parsed = _extract_json_candidate(raw)
            return validate_action_payload(schema_name, parsed)
        except (ValueError, json.JSONDecodeError, ValidationError) as exc:
            last_error = str(exc)
            if attempt == 0:
                messages.extend(
                    [
                        {"role": "assistant", "content": raw},
                        {
                            "role": "user",
                            "content": (
                                "Your previous output was invalid. "
                                f"Error: {last_error}. "
                                f"Return only corrected JSON for {schema_name}. "
                                "Do not include markdown fences or extra text."
                            ),
                        },
                    ]
                )

    print(
        "Warning: LLM JSON parsing failed after retry. "
        f"Using minimal fallback action for step {step_index}. "
        f"Last error: {last_error}"
    )
    if last_raw:
        print("Raw model output:")
        print(last_raw)
    return minimal_action(step_index)


def _build_custom_bundle_for_tier(tier: str, case: dict[str, Any]) -> dict[str, Any]:
    observation = case.get("observation", {})
    if not isinstance(observation, dict):
        raise ValueError(f"Case {case.get('id', 'unknown')} has non-dict observation")

    case_id = str(case.get("id") or "unknown")

    if tier == "image":
        if "image_obs" in observation and isinstance(observation.get("image_obs"), dict):
            image_obs = observation["image_obs"]
            proposed_config = observation.get("proposed_config", {})
        else:
            image_obs = observation
            proposed_config = observation.get("proposed_config", {})
        if not isinstance(proposed_config, dict):
            proposed_config = {}
        return {
            "case_id": case_id,
            "image_observation": image_obs,
            "param_config": proposed_config,
        }

    if tier == "clips":
        return {
            "case_id": case_id,
            "clip_signal_observations": [observation],
        }

    if tier == "weights":
        return {
            "case_id": case_id,
            "weight_observation": observation,
        }

    raise ValueError(f"Unsupported tier: {tier}")


def run_tier_episode(
    *,
    env_url: str,
    tier: str,
    case: dict[str, Any],
    client: OpenAI,
    model_name: str,
) -> TierRunResult:
    bundle = _build_custom_bundle_for_tier(tier, case)
    planned_case_id = str(case.get("id") or "unknown")

    with TalkingHeadBenchEnv(base_url=env_url).sync() as env:
        result = env.reset(mode=tier, custom_bundle=bundle)
        step_index = 0

        while True:
            observation, done, reward, scores = unpack_step_result(result)
            case_id = observation.get("case_id") or planned_case_id
            schema_name = expected_schema_name(step_index, observation)

            if done:
                final_reward = float(reward if reward is not None else 0.0)
                return TierRunResult(
                    tier=tier,
                    case_id=str(case_id) if case_id is not None else None,
                    reward=final_reward,
                    scores=scores or {},
                    steps=step_index,
                    used_custom_bundle=True,
                )

            action = generate_action(
                client=client,
                model_name=model_name,
                tier=tier,
                case_id=str(case_id) if case_id is not None else None,
                step_index=step_index,
                schema_name=schema_name,
                observation=observation,
            )

            result = env.step(action)
            step_index += 1

            if step_index > 8:
                raise RuntimeError(f"Episode exceeded expected step count for tier={tier}")


def extract_tier_score(result: TierRunResult) -> float:
    key_by_tier = {
        "image": "subenv1_score",
        "clips": "subenv2_score",
        "weights": "subenv3_score",
    }
    score_key = key_by_tier.get(result.tier)
    if score_key and isinstance(result.scores, dict) and score_key in result.scores:
        value = result.scores.get(score_key)
        if isinstance(value, (int, float)):
            return float(value)
    return float(result.reward)


def main() -> None:
    args = parse_args()

    api_base_url = require_env("API_BASE_URL")
    model_name = require_env("MODEL_NAME")
    hf_token = require_env("HF_TOKEN")

    tier_cases: dict[str, list[dict[str, Any]]] = {
        tier: load_test_set_cases(path) for tier, path in TEST_SET_FILES.items()
    }

    client = OpenAI(api_key=hf_token, base_url=api_base_url)

    print("TalkingHeadBench baseline inference")
    print(f"Environment URL: {args.env_url}")
    print(f"LLM base URL: {api_base_url}")
    print(f"Model: {model_name}")

    run_order = ["image", "clips", "weights"]
    tier_results: dict[str, TierRunResult] = {}

    for tier in run_order:
        print(f"\nRunning tier: {tier}")
        case = tier_cases[tier][0]
        result = run_tier_episode(
            env_url=args.env_url,
            tier=tier,
            case=case,
            client=client,
            model_name=model_name,
        )
        tier_results[tier] = result

        print(
            f"tier={tier} case_id={result.case_id} reward={result.reward:.4f} "
            f"steps={result.steps} custom_bundle={result.used_custom_bundle}"
        )

    s1 = extract_tier_score(tier_results["image"])
    s2 = extract_tier_score(tier_results["clips"])
    s3 = extract_tier_score(tier_results["weights"])
    weighted_final = 0.25 * s1 + 0.35 * s2 + 0.40 * s3

    report = {
        "per_subenv_scores": {
            "subenv1_score": round(s1, 6),
            "subenv2_score": round(s2, 6),
            "subenv3_score": round(s3, 6),
        },
        "weighted_final_score": round(weighted_final, 6),
        "tiers": {
            tier: {
                "case_id": tier_results[tier].case_id,
                "reward": round(tier_results[tier].reward, 6),
                "steps": tier_results[tier].steps,
                "scores": tier_results[tier].scores,
                "loaded_cases": len(tier_cases[tier]),
            }
            for tier in run_order
        },
    }

    print("\nStructured score report")
    print(json.dumps(report, indent=2, sort_keys=True))


if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        sys.exit(130)