scripts/inference.py

"""
SYNAPSE-X scripts/inference.py
================================
Multi-task inference script used for benchmarking and the report.py helper.
Runs all four tasks and writes inference_results.json.

This is the SCRIPTS version (benchmarking helper). The official hackathon
submission entrypoint is the root-level inference.py.
"""

import json
import os
import re
import sys
import time
from pathlib import Path
from typing import List, Optional

PROJECT_ROOT = Path(__file__).resolve().parents[1]
if str(PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(PROJECT_ROOT))

try:
    from openai import OpenAI
    _OPENAI_AVAILABLE = True
except ImportError:
    OpenAI = None
    _OPENAI_AVAILABLE = False

from agents.baseline import select_action as select_baseline_action
from env.environment import SynapseXEnvironment
from env.grader import TASK_REGISTRY, TASK_SEEDS, grade
from env.models import Action, ActionPayload, Observation


API_BASE_URL     = os.environ.get("API_BASE_URL", "https://router.huggingface.co/v1")
MODEL_NAME       = os.environ.get("MODEL_NAME",   "meta-llama/Meta-Llama-3-8B-Instruct")
HF_TOKEN         = os.environ.get("HF_TOKEN", "")
LOCAL_IMAGE_NAME = os.environ.get("LOCAL_IMAGE_NAME", "synapse-x")
MAX_STEPS        = 20
TEMPERATURE      = 0.0
INFERENCE_MODE   = os.environ.get("INFERENCE_MODE", "auto").lower()
TASK_NAMES       = tuple(TASK_REGISTRY.keys())
USE_LLM          = bool(HF_TOKEN) and INFERENCE_MODE not in {"baseline", "offline"} and _OPENAI_AVAILABLE

SYSTEM_PROMPT = (
    "You are an expert task scheduler inside the SYNAPSE-X decision environment.\n\n"
    "You receive a JSON observation with pending tasks, current time, and resources.\n"
    "Respond with ONLY a single valid JSON action object.\n"
    'Format: {"action_type": "execute"|"delay"|"reallocate", "task_id": <int>}\n'
)

FALLBACK_ACTION: ActionPayload = {"action_type": "delay", "task_id": 0}
BENCHMARK = "synapse-x"


def _write(line: str) -> None:
    sys.stdout.write(f"{line}\n")
    sys.stdout.flush()


def _bool(v: bool) -> str:
    return str(bool(v)).lower()


def _compact(v: object) -> str:
    return json.dumps(v, separators=(",", ":"), ensure_ascii=True)


def _safe_error(err: Optional[str]) -> str:
    if not err:
        return "null"
    return _compact(str(err).replace("\n", " ").strip())


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


def log_step(step: int, action: ActionPayload, reward: float, done: bool, error: Optional[str]) -> None:
    _write(
        f"[STEP] step={step} action={_compact(action)} reward={reward:.2f} "
        f"done={_bool(done)} error={_safe_error(error)}"
    )


def log_end(success: bool, steps: int, score: float, rewards: List[float]) -> None:
    rewards_text = ",".join(f"{r:.2f}" for r in rewards)
    _write(f"[END] success={_bool(success)} steps={steps} score={score:.3f} rewards={rewards_text}")


def _obs_to_prompt(obs: Observation) -> str:
    return json.dumps(
        {
            "time": obs.time,
            "resources": obs.resources,
            "tasks": [
                {
                    "id": t.id, "name": t.name, "priority": t.priority,
                    "risk": t.risk, "uncertainty": t.uncertainty,
                    "deadline": t.deadline, "future_risk": t.future_risk,
                    "deadline_pressure": t.deadline_pressure,
                    "resources_required": t.resources_required,
                    "completed": t.completed, "failed": t.failed,
                }
                for t in obs.tasks
            ],
        },
        indent=2,
    )


def _parse_action(text: str) -> ActionPayload:
    m = re.search(r"\{.*?\}", text, re.DOTALL)
    if m:
        try:
            return json.loads(m.group())
        except json.JSONDecodeError:
            pass
    try:
        return json.loads(text.strip())
    except json.JSONDecodeError:
        return FALLBACK_ACTION.copy()


def _coerce_action(raw: ActionPayload) -> Action:
    try:
        return Action(**raw)
    except Exception:
        return Action(**FALLBACK_ACTION)


def _baseline(obs: Observation) -> Action:
    try:
        return select_baseline_action(obs)
    except Exception:
        return Action(**FALLBACK_ACTION)


def _call_llm(client, obs: Observation) -> ActionPayload:
    completion = client.chat.completions.create(
        model=MODEL_NAME,
        messages=[
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user",   "content": _obs_to_prompt(obs)},
        ],
        temperature=TEMPERATURE,
        max_tokens=64,
        stream=False,
    )
    return _parse_action((completion.choices[0].message.content or "").strip())


def run_baseline_episode(task_name: str, seed: int = 42, verbose: bool = True) -> tuple:
    env = SynapseXEnvironment(task_config=TASK_REGISTRY[task_name], seed=seed)
    obs = env.reset()
    actions: List[ActionPayload] = []
    traces: list = []
    for step in range(MAX_STEPS):
        if obs.episode_done:
            break
        action = _baseline(obs)
        actions.append(action.model_dump())
        result = env.step(action)
        traces.append({"step": step + 1, "reward": result.reward, "done": result.done, "info": result.info})
        if verbose:
            err = (result.info or {}).get("error") or None
            log_step(step + 1, action.model_dump(), float(result.reward), bool(result.done), err)
        obs = result.observation
        if result.done:
            break
    return actions, traces


def run_llm_episode(client, task_name: str, seed: int = 42, verbose: bool = True) -> tuple:
    env = SynapseXEnvironment(task_config=TASK_REGISTRY[task_name], seed=seed)
    obs = env.reset()
    actions: List[ActionPayload] = []
    traces: list = []
    for step in range(MAX_STEPS):
        if obs.episode_done:
            break
        try:
            raw = _call_llm(client, obs)
        except Exception as exc:
            print(f"[DEBUG] LLM error step {step+1}: {exc}", file=sys.stderr)
            raw = _baseline(obs).model_dump()
        action = _coerce_action(raw)
        actions.append(action.model_dump())
        result = env.step(action)
        traces.append({"step": step + 1, "reward": result.reward, "done": result.done, "info": result.info})
        if verbose:
            err = ((result.info or {}).get("error") or None)
            log_step(step + 1, action.model_dump(), float(result.reward), bool(result.done), err)
        obs = result.observation
        if result.done:
            break
    return actions, traces


def run_episode(task_name: str, seed: int = 42, verbose: bool = False) -> float:
    """Run one episode and return the grader score. Used by report.py."""
    use_baseline = not USE_LLM
    if use_baseline:
        actions, _ = run_baseline_episode(task_name, seed=seed, verbose=verbose)
    else:
        try:
            client = OpenAI(base_url=API_BASE_URL, api_key=HF_TOKEN)
            actions, _ = run_llm_episode(client, task_name, seed=seed, verbose=verbose)
        except Exception:
            actions, _ = run_baseline_episode(task_name, seed=seed, verbose=verbose)
    return grade(task_name, actions).score


def main() -> None:
    use_baseline = not USE_LLM
    runtime_model = MODEL_NAME if not use_baseline else "baseline-fallback"

    client = None
    if not use_baseline:
        try:
            client = OpenAI(base_url=API_BASE_URL, api_key=HF_TOKEN)
        except Exception as exc:
            print(f"[DEBUG] Client init failed: {exc}", file=sys.stderr)
            use_baseline = True

    all_scores: dict = {}
    per_task: dict = {}
    started_at = time.perf_counter()

    for task_name in TASK_NAMES:
        seed = TASK_SEEDS.get(task_name, 42)
        log_start(task=task_name, env=BENCHMARK, model=runtime_model)

        if use_baseline or client is None:
            actions, traces = run_baseline_episode(task_name, seed=seed, verbose=True)
        else:
            actions, traces = run_llm_episode(client, task_name, seed=seed, verbose=True)

        result = grade(task_name, actions)
        all_scores[task_name] = result.score
        per_task[task_name] = {
            "score": result.score,
            "completion_rate": result.completion_rate,
            "efficiency": result.efficiency,
            "reward_score": result.reward_score,
            "details": result.details,
            "actions": actions,
            "step_trace": traces,
        }

        log_end(
            success=result.score >= 0.5,
            steps=len(actions),
            score=float(result.score),
            rewards=[float(t["reward"]) for t in traces],
        )

    average = sum(all_scores.values()) / len(all_scores)
    output = {
        "model":           runtime_model,
        "mode":            "llm" if not use_baseline else "baseline-fallback",
        "api_base_url":    API_BASE_URL,
        "local_image_name": LOCAL_IMAGE_NAME,
        "elapsed_seconds": round(time.perf_counter() - started_at, 4),
        "average_score":   round(average, 4),
        "scores":          all_scores,
        "tasks":           per_task,
    }
    with open(PROJECT_ROOT / "inference_results.json", "w", encoding="utf-8") as fh:
        json.dump(output, fh, indent=2)


if __name__ == "__main__":
    main()