inference.py

"""
inference.py — OpenEnv-compliant inference script for xsecure.

Required env vars:
    HF_TOKEN       Hugging Face / API key
    API_BASE_URL   LLM endpoint
    MODEL_NAME     Model identifier
"""

from __future__ import annotations

import asyncio
import json
import os
import re
import sys
from typing import Dict, List, Optional, Any

from dotenv import load_dotenv
from openai import OpenAI

from client import IncidentResponseEnv, StepResult
from models import IncidentAction, IncidentObservation

# Load .env for local dev
#load_dotenv()

# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------
"""
API_KEY      = os.getenv("API_KEY", "")
API_BASE_URL = os.getenv("API_BASE_URL")
MODEL_NAME   = os.getenv("MODEL_NAME")
ENV_URL      = os.getenv("ENV_URL", "http://localhost:7860")
BENCHMARK    = "xsecure"
MAX_STEPS    = 20
"""

API_BASE_URL = os.environ["API_BASE_URL"]
MODEL_NAME   = os.environ.get("MODEL_NAME",  "gpt-4o-mini")
API_KEY      = os.environ.get("API_KEY", "")
ENV_URL      = os.environ.get("ENV_URL", "http://localhost:7860")

BENCHMARK    = "xsecure"
MAX_STEPS    = 20

if not API_KEY:
    print("ERROR: HF_TOKEN is not set.", file=sys.stderr)
    sys.exit(1)

# Use AsyncOpenAI to prevent blocking the event loop
llm = OpenAI(base_url=API_BASE_URL, api_key=API_KEY)

# ---------------------------------------------------------------------------
# Mandatory stdout loggers (Fixed spacing to match spec)
# ---------------------------------------------------------------------------

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: Optional[str]) -> None:
    # Spec requires double space after [STEP] for some parsers
    print(
        f"[STEP]  step={step} action={action} reward={reward:.2f} "
        f"done={str(done).lower()} error={error or 'null'}",
        flush=True,
    )

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

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

SYSTEM_PROMPT = """\
You are an expert cybersecurity incident responder AI agent.
Your goal is to investigate logs and alerts, identify the threat, and mitigate it.

## Available Actions (one per step):
- analyze_log(log_id)
- trace_user(user_id)
- block_ip(ip_address)
- disable_account(user_id)
- restart_service(service)
- ignore

## Response Format (STRICT JSON):
{"action_type": "analyze_log", "target": "L001"}
"""

# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def _format_observation(obs: IncidentObservation) -> str:
    # Use dot notation as expected by the environment models
    """logs_txt     = "\n".join(f"  [{l.log_id}] {l.timestamp} — {l.message}" for l in obs.logs)
    alerts_txt   = "\n".join(f"  [{a.severity.upper()}] {a.message}" for a in obs.alerts)
    services_txt = "\n".join(f"  {s.name}: {s.status}" for s in obs.services)
    """
    logs_txt     = "\n".join(f"  [{l['log_id']}] {l['timestamp']} — {l['message']}" for l in obs.logs)
    alerts_txt   = "\n".join(f"  [{a['severity'].upper()}] {a.message}" for a in obs.alerts)
    services_txt = "\n".join(f"  {s['name']}: {s['status']}" for s in obs.services)
    
    return (
        f"=== Incident Dashboard (Step {obs.step_count}) ===\n\n"
        f"LOGS:\n{logs_txt}\n\n"
        f"ALERTS:\n{alerts_txt}\n\n"
        f"SERVICES:\n{services_txt}\n\n"
        f"ACTIVE USERS: {', '.join(obs.active_users)}\n\n"
        f"Last action result: {obs.last_action_result}"
    )

def _parse_action(text: str) -> IncidentAction:
    """Extract JSON action with filtering for extra fields to avoid Pydantic errors."""
    try:
        # 1. Try direct or markdown-wrapped JSON
        pattern = re.search(r"(\{.*?\})", text.strip().replace("\n", " "), re.DOTALL)
        if pattern:
            data = json.loads(pattern.group(1))
            # Only pass fields known to IncidentAction
            valid_keys = {"action_type", "target"}
            filtered = {k: v for k, v in data.items() if k in valid_keys}
            return IncidentAction(**filtered)
    except Exception:
        pass
    return IncidentAction(action_type="ignore", target="")

def _get_action(conversation: List[Dict], obs: IncidentObservation) -> IncidentAction:
    conversation.append({"role": "user", "content": _format_observation(obs)})
    
    response = llm.chat.completions.create(
        model=MODEL_NAME,
        messages=[{"role": "system", "content": SYSTEM_PROMPT}] + conversation,
        max_tokens=256,
        temperature=0.0,
    )
    
    text = response.choices[0].message.content or ""
    conversation.append({"role": "assistant", "content": text})
    return _parse_action(text)

# ---------------------------------------------------------------------------
# Episode runner
# ---------------------------------------------------------------------------

TASK_NAMES = {1: "brute-force-easy", 2: "suspicious-login-medium", 3: "multi-stage-apt-hard"}

def run_episode(task_id: int) -> None:
    task_name = TASK_NAMES.get(task_id, f"task-{task_id}")
    rewards: List[float] = []
    steps_taken = 0
    success = False
    score = 0.0
    conversation: List[Dict] = []

    log_start(task=task_name, env=BENCHMARK, model=MODEL_NAME)

    try:
        with IncidentResponseEnv(base_url=ENV_URL) as env:
            obs = env.reset(task_id=task_id)

            for step in range(1, MAX_STEPS + 1):
                # Now awaited correctly
                action = _get_action(conversation, obs)
                result = env.step(action)

                rewards.append(result.reward)
                steps_taken = step
                obs = result.observation

                log_step(
                    step=step,
                    action=f"{action.action_type}({action.target!r})",
                    reward=result.reward,
                    done=result.done,
                    error=None,
                )

                if result.done:
                    info = result.info or {}
                    # Robust score parsing
                    raw_score = info.get("final_score", 0.0)
                    score = min(max(float(raw_score or 0.0), 0.0), 1.0)
                    success = bool(info.get("success", False))
                    break
    except Exception as e:
        print(f"ERROR: Episode failed: {e}", file=sys.stderr)
    finally:
        log_end(success=success, steps=steps_taken, score=score, rewards=rewards)

# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------

def main():
    task_ids_str = os.getenv("TASK_IDS", "1,2,3")
    task_ids = [int(t.strip()) for t in task_ids_str.split(",") if t.strip()]
    for task_id in task_ids:
        run_episode(task_id)

if __name__ == "__main__":
    main()