Blame for inference.py

"""
inference.py
============
Baseline ReAct-style agent for InformationIntegrityEnv.

Two run modes
-------------
benchmark   Run all difficulties over all scenarios (default).
            Prints per-episode [START]/[STEP]/[END] lines, then a final
            summary table and tier-separation check.

single      Run one episode against a specific scenario or a live server.

Environment variables
---------------------
  API_BASE_URL      URL of the LLM provider (OpenAI-compatible)
  MODEL_NAME        Model to use (e.g. gpt-4o, claude-3-5-sonnet, etc.)
  HF_TOKEN          Bearer token / API key for the provider
  ENV_SERVER_URL    URL of the env server (default http://localhost:8000)
  MAX_STEPS         Per-episode step limit (default 20)
  BENCHMARK_SEED    Random seed for reproducible scenario sampling (default 42)

Logging protocol (required by OpenEnv)
---------------------------------------
  [START]  — once at episode start
  [STEP]   — at every step
  [END]    — with final score

Benchmark summary format (required)
-------------------------------------
  difficulty=easy   avg_score=X.XXXX  avg_reward=X.XXXX  n=N
  difficulty=medium avg_score=X.XXXX  avg_reward=X.XXXX  n=N
  difficulty=hard   avg_score=X.XXXX  avg_reward=X.XXXX  n=N

Usage
-----
  # Full benchmark (local, all scenarios)
  python inference.py --benchmark

  # Single episode from dataset
  python inference.py --local --difficulty medium --scenario-id 0

  # Single episode from JSON file
  python inference.py --local --scenario scenario.json

  # Against running env server
  python inference.py --server --difficulty easy --scenario-id 1

  # Write benchmark results to JSON
  python inference.py --benchmark --out results.json
"""

from __future__ import annotations

import argparse
import json
import logging
import os
import random
import sys
import time
from typing import Any, Dict, List, Optional

try:
    from dotenv import load_dotenv
    load_dotenv()
except ImportError:
    pass  # rely on shell environment

import requests as http_requests
from openai import OpenAI

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
)
logger = logging.getLogger("inference")

# ---------------------------------------------------------------------------
# Config from environment
# ---------------------------------------------------------------------------

API_BASE_URL    = os.getenv("API_BASE_URL",    "https://router.huggingface.co/v1")
MODEL_NAME      = os.getenv("MODEL_NAME",      "Qwen/Qwen2.5-72B-Instruct")
HF_TOKEN        = os.getenv("HF_TOKEN")
ENV_SERVER_URL  = os.getenv("ENV_SERVER_URL",  "http://localhost:8000")
MAX_STEPS       = int(os.getenv("MAX_STEPS",   10))
BENCHMARK_SEED  = int(os.getenv("BENCHMARK_SEED", 42))
DIFFICULTY      = os.getenv("DIFFICULTY",      "easy")
SCENARIO_ID     = os.getenv("SCENARIO_ID",     None)   # int index or string scenario id

DIFFICULTIES    = ["easy", "medium", "hard"]


# ---------------------------------------------------------------------------
# LLM client
# ---------------------------------------------------------------------------

def _build_client() -> OpenAI:
    if not HF_TOKEN:
        raise ValueError(
            "No API key found. Set HF_TOKEN environment variable."
        )
    return OpenAI(api_key=HF_TOKEN, base_url=API_BASE_URL)


# ---------------------------------------------------------------------------
# Environment interface — local (in-process)
# ---------------------------------------------------------------------------

class LocalEnvClient:
    """Runs InformationIntegrityEnv in-process against the static dataset."""

    def __init__(
        self,
        difficulty: str,
        scenario_id: Optional[str | int] = None,
        scenario_dict: Optional[dict] = None,
    ) -> None:
        from environment import InformationIntegrityEnv

        if scenario_dict is not None:
            self._env = InformationIntegrityEnv(scenario=scenario_dict)
        else:
            self._env = InformationIntegrityEnv.from_dataset(
                difficulty=difficulty,
                scenario_id=scenario_id,
            )
        self._done = False

    def reset(self, difficulty: Optional[str] = None, scenario_id: Optional[str | int] = None) -> dict:
        obs = self._env.reset(difficulty=difficulty, scenario_id=scenario_id)
        self._done = False
        return obs.model_dump()

    def step(self, action: dict) -> dict:
        from models import Action
        act = Action(**action)
        obs, rew, done, info = self._env.step(act)
        self._done = done
        return {
            "observation": obs.model_dump(),
            "reward":      rew.model_dump(),
            "done":        done,
            "info":        info,
        }

    @property
    def done(self) -> bool:
        return self._done

    def episode_result(self) -> dict:
        return self._env.episode_result().model_dump()


# ---------------------------------------------------------------------------
# Environment interface — remote server
# ---------------------------------------------------------------------------

class ServerEnvClient:
    """Communicates with a running FastAPI env server."""

    def __init__(self, base_url: str = ENV_SERVER_URL) -> None:
        self._base = base_url.rstrip("/")
        self._session_id: Optional[str] = None

    def reset(
        self,
        difficulty: Optional[str] = None,
        scenario_id: Optional[str | int] = None,
        scenario_dict: Optional[dict] = None,
    ) -> dict:
        payload: dict = {"max_steps": MAX_STEPS}
        if scenario_dict:
            payload["scenario"] = scenario_dict
        elif difficulty:
            payload["difficulty"] = difficulty
            if scenario_id is not None:
                payload["scenario_id"] = scenario_id
        r = http_requests.post(f"{self._base}/reset", json=payload, timeout=30)
        r.raise_for_status()
        data = r.json()
        self._session_id = data["session_id"]
        return data["observation"]

    def step(self, action: dict) -> dict:
        payload = {"session_id": self._session_id, "action": action}
        r = http_requests.post(f"{self._base}/step", json=payload, timeout=30)
        r.raise_for_status()
        return r.json()

    def episode_result(self) -> dict:
        # Server does not expose a dedicated episode_result endpoint;
        # grader_score is included in the final step's info dict.
        return {}


# ---------------------------------------------------------------------------
# System prompt
# ---------------------------------------------------------------------------

SYSTEM_PROMPT = SYSTEM_PROMPT = """You are a high-precision content moderation agent.

Your goal is to make the MOST ACCURATE moderation decision based on evidence.

You MUST follow a structured reasoning process:

1. First verify claims.
2. Then evaluate sources carefully:
   - Prefer HIGH reliability sources
   - Consider stance (support/refute/neutral)
3. Avoid redundant actions.
4. Only decide when evidence is sufficient.

Decision guidelines:
- REMOVE → if content is clearly false and harmful
- FLAG   → if uncertain, conflicting evidence, or partially misleading
- ALLOW  → if claims are well-supported by reliable sources

Action rules:
- You MUST verify at least one claim
- You MUST check at least two sources
- NEVER repeat the same action
- Be efficient (avoid unnecessary steps)

Output STRICT JSON ONLY:
{"type": "...", ...}
"""


def _observation_prompt(obs_dict: dict) -> str:
    content = obs_dict.get("content", "")
    claims  = obs_dict.get("claims", [])
    sources = obs_dict.get("available_sources", [])
    history = obs_dict.get("interaction_history", [])
    flags   = obs_dict.get("status_flags", {})
    step    = obs_dict.get("step_count", 0)

    # Track what has already been done from full history
    checked_source_ids  = {
        h["action_payload"].get("source_id")
        for h in history
        if h.get("action_type") == "check_source" and h.get("action_payload", {}).get("source_id")
    }
    verified_claim_ids  = {
        h["action_payload"].get("claim_id")
        for h in history
        if h.get("action_type") == "verify_claim" and h.get("action_payload", {}).get("claim_id")
    }

    claims_text = "\n".join(
        f"  [{c['id']}] {c['text']} — status: {c.get('verification_status','unverified')}"
        + (" [ALREADY VERIFIED — do not repeat]" if c['id'] in verified_claim_ids else "")
        for c in claims
    )
    sources_text = "\n".join(
        f"  [{s['id']}] {s['title']} | stance={s['stance']} | reliability={s['reliability']:.2f}"
        + (" [ALREADY CHECKED — do not repeat]" if s['id'] in checked_source_ids else "")
        for s in sources
    )

    last_result = ""
    if history:
        last = history[-1]
        last_result = f"\nLast step result: {json.dumps(last.get('result', {}), ensure_ascii=False)[:300]}"

    unchecked_sources  = [s['id'] for s in sources  if s['id']  not in checked_source_ids]
    unverified_claims  = [c['id'] for c in claims   if c['id']  not in verified_claim_ids]
    sufficient         = flags.get('sufficient_evidence', False)

    next_step_hint = ""
    if not sufficient:
        if unverified_claims:
            next_step_hint = f"\nHINT: You still need to verify claims: {unverified_claims}. Use verify_claim next."
        elif unchecked_sources:
            next_step_hint = f"\nHINT: You still need to check sources: {unchecked_sources}. Use check_source next."
    else:
        next_step_hint = "\nHINT: You have sufficient evidence. Use the 'decide' action now (allow/flag/remove)."

    return f"""STEP {step}

CONTENT UNDER REVIEW:
{content}

CLAIMS:
{claims_text}

AVAILABLE SOURCES:
{sources_text}

STATUS FLAGS:
  decision_made={flags.get('decision_made', False)}
  sufficient_evidence={sufficient}
  redundant_actions={flags.get('redundant_actions', 0)}
{last_result}{next_step_hint}

Choose your next action (JSON only):"""


# ---------------------------------------------------------------------------
# LLM action call
# ---------------------------------------------------------------------------

def _llm_action(llm: OpenAI, messages: list, step: int, retries: int = 3) -> dict:
    """Call LLM and parse JSON action. Retries on parse failure."""
    last_exc = None
    for attempt in range(retries):
        try:
            resp = llm.chat.completions.create(
                model=MODEL_NAME,
                messages=messages,
                temperature=0.0,
                max_tokens=256,
                timeout=30,
            )
            raw = resp.choices[0].message.content.strip()
            if raw.startswith("```"):
                raw = raw.split("\n", 1)[-1].rsplit("```", 1)[0].strip()
            return json.loads(raw)
        except json.JSONDecodeError as exc:
            logger.warning("Step %d attempt %d: JSON parse error — %s", step, attempt + 1, exc)
            last_exc = exc
        except Exception as exc:
            logger.error("LLM call failed at step %d: %s", step, exc)
            last_exc = exc
            if attempt < retries - 1:
                time.sleep(2 ** attempt)

    # Raise so the episode fails loudly rather than silently looping
    raise RuntimeError(f"LLM call failed after {retries} attempts at step {step}: {last_exc}")


# ---------------------------------------------------------------------------
# Explanation builder
# ---------------------------------------------------------------------------

def _build_explanation(obs_dict: dict, decision: Optional[str]) -> str:
    try:
        from explainability import ExplainabilityEngine, build_reasoning_from_episode
        from models import Observation
        obs = Observation(**obs_dict)
        reasoning = build_reasoning_from_episode(obs=obs, decision=decision or "none")
        return ExplainabilityEngine().generate_explanation(reasoning)
    except Exception as exc:
        return f"(explanation unavailable: {exc})"


# ---------------------------------------------------------------------------
# OpenEnv-compliant logging helpers
# ---------------------------------------------------------------------------

def log_start(task: str, env: str, model: str) -> None:
    print(f"[START] task={task} env={env} model={model}", flush=True)


def log_step(step: int, action: str, reward: float, done: bool, error: str = "null") -> None:
    print(
        f"[STEP] step={step} action={action} reward={reward:.2f} "
        f"done={str(done).lower()} error={error}",
        flush=True,
    )


def log_end(success: bool, steps: int, rewards: List[float], score: float = 0.0) -> None:
    rewards_str = ",".join(f"{r:.2f}" for r in rewards)
    print(
        f"[END] success={str(success).lower()} steps={steps} score={score:.4f} rewards={rewards_str}",
        flush=True,
    )


# ---------------------------------------------------------------------------
# OpenEnv single-episode runner (primary entry point)
# ---------------------------------------------------------------------------

def run_single_episode(llm: OpenAI) -> dict:
    """
    Run one episode controlled entirely by DIFFICULTY and SCENARIO_ID env vars.
    Prints strict OpenEnv-format logs: [START], [STEP]×n, [END].
    """
    from tasks.registry import TASKS
    from environment import InformationIntegrityEnv

    # Resolve scenario_id: prefer int index, fall back to string id
    sid: Optional[str | int] = SCENARIO_ID
    if sid is not None:
        try:
            sid = int(sid)
        except (ValueError, TypeError):
            pass  # keep as string scenario id

    task         = TASKS[DIFFICULTY]
    scenario_dict = task.get_scenario(sid if sid is not None else 0)
    task_name = scenario_dict.get('id', f'{DIFFICULTY}_unknown')

    env = InformationIntegrityEnv(scenario=scenario_dict)
    obs = env.reset(difficulty=DIFFICULTY, scenario_id=sid)

    log_start(task=task_name, env="information-integrity-env", model=MODEL_NAME)

    messages: List[Dict[str, str]] = [{"role": "system", "content": SYSTEM_PROMPT}]
    rewards: List[float] = []
    steps_taken = 0
    score       = 0.0

    for step in range(1, MAX_STEPS + 1):
        obs_dict = obs.model_dump()
        messages.append({"role": "user", "content": _observation_prompt(obs_dict)})

        try:
            action_dict = _llm_action(llm, messages, step)
        except RuntimeError as exc:
            logger.error("Aborting episode: %s", exc)
            break
        messages.append({"role": "assistant", "content": json.dumps(action_dict)})

        error_str   = "null"
        step_reward = 0.0
        done        = False

        try:
            from models import Action
            obs, reward_obj, done, info = env.step(Action(**action_dict))
            step_reward = reward_obj.score
        except Exception as exc:
            error_str = str(exc)[:80].replace("\n", " ")

        rewards.append(step_reward)
        steps_taken = step

        log_step(
            step=step,
            action=action_dict.get("type", "unknown"),
            reward=step_reward,
            done=done,
            error=error_str,
        )

        if done:
            try:
                ep_result = env.episode_result()
                score = float(ep_result.grader_score or 0.0)
            except Exception:
                score = 0.0
            score = min(max(score, 0.005), 0.995)
            break

    # Compute grader score even if episode didn't complete (LLM failed / max steps)
    if score == 0.0:
        try:
            ep_result = env.episode_result()
            score = float(ep_result.grader_score or 0.0)
        except Exception:
            score = 0.0
    score = min(max(score, 0.005), 0.995)  # strictly in (0, 1) as required by validator

    success = score >= 0.5
    log_end(success=success, steps=steps_taken, rewards=rewards, score=score)

    return {"score": score, "steps": steps_taken, "success": success}


# ---------------------------------------------------------------------------
# Single episode runner (legacy — used by benchmark loop and CLI)
# ---------------------------------------------------------------------------

def run_episode(
    env: LocalEnvClient | ServerEnvClient,
    llm: OpenAI,
    difficulty: Optional[str] = None,
    scenario_id: Optional[str | int] = None,
) -> dict:
    """Run one full episode using compliant log_start/log_step/log_end."""
    messages: List[Dict[str, str]] = [{"role": "system", "content": SYSTEM_PROMPT}]

    obs_dict = env.reset(difficulty=difficulty, scenario_id=scenario_id)

    task_name = f"{difficulty}_{scenario_id}" if difficulty and scenario_id is not None \
        else obs_dict.get("task_difficulty", "unknown")
    log_start(task=task_name, env="information-integrity-env", model=MODEL_NAME)

    rewards: List[float] = []
    step        = 0
    grader_score: Optional[float] = None

    for step in range(1, MAX_STEPS + 1):
        messages.append({"role": "user", "content": _observation_prompt(obs_dict)})

        try:
            action_dict = _llm_action(llm, messages, step)
        except RuntimeError as exc:
            break
        messages.append({"role": "assistant", "content": json.dumps(action_dict)})

        error_str   = "null"
        step_reward = 0.0
        done        = False

        try:
            result      = env.step(action_dict)
            obs_dict    = result["observation"]
            step_reward = result["reward"].get("score", 0.0)
            done        = result["done"]
        except Exception as exc:
            error_str = str(exc)[:80].replace("\n", " ")

        rewards.append(step_reward)
        log_step(step=step, action=action_dict.get("type", "unknown"),
                 reward=step_reward, done=done, error=error_str)

        if done:
            break

    if isinstance(env, LocalEnvClient):
        try:
            grader_score = env.episode_result().get("grader_score")
        except Exception:
            pass

    score   = max(float(grader_score or 0.005), 0.005)  # strictly > 0 as required by validator
    success = score >= 0.5
    log_end(success=success, steps=step, rewards=rewards, score=score)

    return {
        "difficulty":    difficulty,
        "scenario_id":   scenario_id,
        "steps":         step,
        "grader_score":  score,      # always a float >= 0.005
        "total_reward": sum(rewards),
    }



# ---------------------------------------------------------------------------
# Benchmark loop
# ---------------------------------------------------------------------------

def run_benchmark(llm: OpenAI, out_path: Optional[str] = None) -> Dict[str, Any]:
    """
    Loop over all difficulties × all scenarios in each dataset.
    Prints per-episode logs, then the required summary table.

    Returns
    -------
    dict with per-difficulty aggregates and all raw episode results.
    """
    from tasks.registry import TASKS

    all_results: Dict[str, List[dict]] = {d: [] for d in DIFFICULTIES}
    rng = random.Random(BENCHMARK_SEED)

    for difficulty in DIFFICULTIES:
        task       = TASKS[difficulty]
        scenarios  = task.get_all_scenarios()
        n          = len(scenarios)
        logger.info("=== Benchmarking difficulty=%s  (%d scenarios) ===", difficulty, n)

        for idx, sc_dict in enumerate(scenarios):
            # Build a fresh local env for each scenario
            env = LocalEnvClient(
                difficulty=difficulty,
                scenario_id=idx,
            )
            result = run_episode(
                env=env,
                llm=llm,
                difficulty=difficulty,
                scenario_id=idx,
            )
            result["scenario_index"] = idx
            result["scenario_id"]    = sc_dict.get("id", str(idx))
            all_results[difficulty].append(result)

    # ---- Summary ----
    summary: Dict[str, dict] = {}
    print()  # blank line before summary
    for difficulty in DIFFICULTIES:
        episodes = all_results[difficulty]
        n = len(episodes)

        scores  = [e["grader_score"] for e in episodes]  # always float >= 0.005
        rewards = [e["total_reward"]  for e in episodes]

        avg_score  = sum(scores) / len(scores) if scores else 0.005
        avg_reward = sum(rewards) / len(rewards) if rewards else 0.0

        summary[difficulty] = {
            "avg_score":  round(avg_score,  4),
            "avg_reward": round(avg_reward, 4),
            "n":          n,
        }
        print(
            f"difficulty={difficulty:<6} "
            f"avg_score={avg_score:.4f}  "
            f"avg_reward={avg_reward:.4f}  "
            f"n={n}"
        )

    # ---- Difficulty tier separation check ----
    easy_s   = summary["easy"]["avg_score"]
    medium_s = summary["medium"]["avg_score"]
    hard_s   = summary["hard"]["avg_score"]

    if easy_s > medium_s > hard_s:
        print("Difficulty tiers well separated: easy > medium > hard ✓")
    else:
        print(
            "WARNING: Difficulty tiers not well separated "
            f"(easy={easy_s:.4f}, medium={medium_s:.4f}, hard={hard_s:.4f})"
        )

    benchmark_output = {
        "summary":  summary,
        "episodes": all_results,
    }

    if out_path:
        with open(out_path, "w", encoding="utf-8") as fh:
            json.dump(benchmark_output, fh, indent=2, ensure_ascii=False)
        logger.info("Benchmark results written to %s", out_path)

    return benchmark_output


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------

def _parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="InformationIntegrity benchmark agent",
        formatter_class=argparse.RawTextHelpFormatter,
    )

    mode = parser.add_mutually_exclusive_group()
    mode.add_argument(
        "--benchmark", action="store_true",
        help="Run full benchmark over all difficulties and all scenarios (default mode)",
    )
    mode.add_argument(
        "--local", action="store_true",
        help="Run a single episode in-process (requires --difficulty or --scenario)",
    )
    mode.add_argument(
        "--server", action="store_true",
        help="Run a single episode against a running env server",
    )

    parser.add_argument("--difficulty", default="medium",
                        choices=DIFFICULTIES,
                        help="Difficulty tier (used with --local or --server)")
    parser.add_argument("--scenario-id", default=None,
                        help="Scenario index (int) or id (str); None = random")
    parser.add_argument("--scenario", default=None,
                        help="Path to a pre-collected scenario JSON file (overrides --difficulty)")
    parser.add_argument("--out", default=None,
                        help="Write episode/benchmark result JSON to this path")

    return parser.parse_args()


def main() -> None:
    args   = _parse_args()
    llm    = _build_client()

    # ---------- Default: OpenEnv single-episode mode (env-var driven) ----------
    if not args.benchmark and not args.local and not args.server:
        run_single_episode(llm)
        return

    # ---------- Benchmark ----------
    if args.benchmark:
        run_benchmark(llm, out_path=args.out)
        return

    # ---------- Single episode ----------
    scenario_id: Optional[str | int] = args.scenario_id
    if scenario_id is not None:
        try:
            scenario_id = int(scenario_id)
        except (ValueError, TypeError):
            pass  # keep as string id

    # Load scenario from file if specified
    scenario_dict: Optional[dict] = None
    if args.scenario:
        with open(args.scenario, "r", encoding="utf-8") as fh:
            scenario_dict = json.load(fh)

    if args.local:
        env = LocalEnvClient(
            difficulty=args.difficulty,
            scenario_id=scenario_id,
            scenario_dict=scenario_dict,
        )
    else:
        env = ServerEnvClient(base_url=ENV_SERVER_URL)

    result = run_episode(
        env=env,
        llm=llm,
        difficulty=args.difficulty,
        scenario_id=scenario_id,
    )

    if args.out:
        with open(args.out, "w", encoding="utf-8") as fh:
            json.dump(result, fh, indent=2, ensure_ascii=False)
        logger.info("Episode result written to %s", args.out)


if __name__ == "__main__":
    main()