inference.py

"""
SYNAPSE-X inference.py  —  Official submission entrypoint
==========================================================

MANDATORY REQUIREMENTS MET:
  - Named `inference.py` and placed in project root
  - Reads API_BASE_URL, MODEL_NAME, HF_TOKEN from environment
  - Uses OpenAI client for all LLM calls when HF_TOKEN is available
  - Emits exact [START] / [STEP] / [END] stdout format per spec
  - Runs ALL tasks (easy, medium, hard, triage) so graders cover 4 tasks
  - Every task score is in [0.0, 1.0]
  - Completes well under the 20-minute runtime limit
  - Falls back cleanly to deterministic baseline when no HF_TOKEN
  - Writes inference_results.json to project root

STDOUT FORMAT (per spec):
  [START] task=<n> env=synapse-x model=<model>
  [STEP]  step=<n> action=<json> reward=<0.00> done=<true|false> error=<msg|null>
  [END]   success=<true|false> steps=<n> score=<0.000> rewards=<r1,r2,...>
"""

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

# ---------------------------------------------------------------------------
# Path bootstrap
# ---------------------------------------------------------------------------
PROJECT_ROOT = Path(__file__).resolve().parent
if str(PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(PROJECT_ROOT))

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

# ---------------------------------------------------------------------------
# Internal imports
# ---------------------------------------------------------------------------
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

# ---------------------------------------------------------------------------
# Environment configuration  (hackathon-mandated variable names)
# ---------------------------------------------------------------------------
API_BASE_URL     = os.getenv("API_BASE_URL",  "https://router.huggingface.co/v1")
MODEL_NAME       = os.getenv("MODEL_NAME",    "meta-llama/Meta-Llama-3-8B-Instruct")
HF_TOKEN         = os.getenv("HF_TOKEN") or os.getenv("API_KEY") or ""
LOCAL_IMAGE_NAME = os.getenv("LOCAL_IMAGE_NAME", "synapse-x")

# All four tasks evaluated in order
TASK_NAMES: tuple = ("easy", "medium", "hard", "triage")
ENV_NAME          = "synapse-x"
MAX_STEPS         = 20
TEMPERATURE       = 0.0
MAX_TOKENS        = 64
SUCCESS_THRESHOLD = 0.5

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"
    "Each task has: id, name, priority, risk, uncertainty, deadline, future_risk,\n"
    "deadline_pressure, resources_required, dependencies, completed, failed.\n\n"
    "Respond with ONLY a single valid JSON action object -- no prose, no markdown.\n"
    'Format: {"action_type": "execute"|"delay"|"reallocate", "task_id": <int>}\n\n'
    "Strategy:\n"
    "- execute  high-priority tasks whose dependencies are satisfied and risk < 0.7\n"
    "- delay    tasks with risk > 0.7 and uncertainty > 0.5 when deadline permits\n"
    "- reallocate when resources are too low for the best feasible task\n"
)

FALLBACK_ACTION: ActionPayload = {"action_type": "delay", "task_id": 0}
USE_LLM: bool = bool(HF_TOKEN) and _OPENAI_AVAILABLE


# ===========================================================================
# Logging helpers  (exact format required by hackathon spec)
# ===========================================================================

def _bool(v: bool) -> str:
    return "true" if v else "false"


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, model: str) -> None:
    print(f"[START] task={task} env={ENV_NAME} model={model}", flush=True)


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


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


# ===========================================================================
# Baseline (deterministic, no API required)
# ===========================================================================

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


# ===========================================================================
# LLM policy helpers
# ===========================================================================

def _obs_to_prompt(obs: Observation) -> str:
    data = {
        "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,
                "dependencies": t.dependencies,
                "completed": t.completed,
                "failed": t.failed,
            }
            for t in obs.tasks
        ],
    }
    return json.dumps(data, indent=2)


def _parse_action(text: str) -> ActionPayload:
    match = re.search(r"\{.*?\}", text, re.DOTALL)
    if match:
        try:
            return json.loads(match.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 _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=MAX_TOKENS,
        stream=False,
    )
    text = (completion.choices[0].message.content or "").strip()
    return _parse_action(text)


# ===========================================================================
# Single-episode runner
# ===========================================================================

def run_episode(
    task_name: str,
    seed: int,
    runtime_model: str,
    client=None,
) -> tuple:
    """
    Run one episode for `task_name`.
    Emits [START] ... [STEP]* lines.
    Returns (actions_taken, rewards, steps_taken).
    The caller emits [END].
    """
    env = SynapseXEnvironment(task_config=TASK_REGISTRY[task_name], seed=seed)
    obs: Observation = env.reset()

    log_start(task=task_name, model=runtime_model)

    actions_taken: List[ActionPayload] = []
    rewards: List[float] = []
    steps_taken = 0

    try:
        for step_idx in range(MAX_STEPS):
            if obs.episode_done:
                break

            # Choose action: LLM with baseline fallback, or pure baseline
            if client is not None:
                try:
                    raw = _call_llm(client, obs)
                    action = _coerce_action(raw)
                except Exception as exc:
                    print(f"[DEBUG] LLM error step {step_idx+1}: {exc}", file=sys.stderr, flush=True)
                    action = baseline_action(obs)
            else:
                action = baseline_action(obs)

            result = env.step(action)

            actions_taken.append(action.model_dump())
            rewards.append(float(result.reward))
            steps_taken = step_idx + 1

            error_val: Optional[str] = None
            if isinstance(result.info, dict):
                raw_err = result.info.get("error")
                error_val = str(raw_err) if raw_err else None

            log_step(
                step=steps_taken,
                action=action.model_dump(),
                reward=float(result.reward),
                done=bool(result.done),
                error=error_val,
            )

            obs = result.observation
            if result.done:
                break

    except Exception as exc:
        print(f"[DEBUG] Episode exception ({task_name}): {exc}", file=sys.stderr, flush=True)

    return actions_taken, rewards, steps_taken


# ===========================================================================
# Main  —  iterate all tasks, produce full START/STEP*/END per task
# ===========================================================================

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

    # Build OpenAI client only when token is available
    client = None
    if USE_LLM:
        try:
            client = OpenAI(base_url=API_BASE_URL, api_key=HF_TOKEN)
        except Exception as exc:
            print(f"[DEBUG] OpenAI client init failed: {exc}", file=sys.stderr, flush=True)
            client = None

    all_results: dict = {}
    started_at = time.perf_counter()

    for task_name in TASK_NAMES:
        seed = TASK_SEEDS.get(task_name, 42)

        actions, rewards, steps = run_episode(
            task_name=task_name,
            seed=seed,
            runtime_model=runtime_model,
            client=client,
        )

        # Grade — deterministic, always returns score in [0.0, 1.0]
        grade_result = grade(task_name, actions)
        score   = float(grade_result.score)
        success = score >= SUCCESS_THRESHOLD

        log_end(success=success, steps=steps, score=score, rewards=rewards)

        all_results[task_name] = {
            "score":           score,
            "completion_rate": grade_result.completion_rate,
            "efficiency":      grade_result.efficiency,
            "reward_score":    grade_result.reward_score,
            "details":         grade_result.details,
            "actions":         actions,
            "rewards":         rewards,
            "steps":           steps,
            "success":         success,
        }

    elapsed       = round(time.perf_counter() - started_at, 3)
    average_score = round(sum(v["score"] for v in all_results.values()) / len(all_results), 4)

    output = {
        "model":           runtime_model,
        "mode":            "llm" if client is not None else "baseline-fallback",
        "api_base_url":    API_BASE_URL,
        "local_image_name": LOCAL_IMAGE_NAME,
        "used_hf_token":   bool(HF_TOKEN),
        "elapsed_seconds": elapsed,
        "average_score":   average_score,
        "tasks":           all_results,
    }

    results_path = PROJECT_ROOT / "inference_results.json"
    with open(results_path, "w", encoding="utf-8") as fh:
        json.dump(output, fh, indent=2)


if __name__ == "__main__":
    main()