inference.py

"""
DebugOps inference script
=
Runs the LLM agent (or heuristic fallback) against all three tasks and
emits structured stdout logs in the exact required format:

    [START] task=<name> env=debugops model=<model>
    [STEP]  step=<n> action=<a> reward=<r> done=<bool> error=<null|err>
    [END]   success=<bool> steps=<n> score=<s> rewards=<comma-list>

Environment variables

API_BASE_URL   LLM endpoint  (default: HuggingFace router)
MODEL_NAME     Model ID      (default: Qwen/Qwen2.5-72B-Instruct)
HF_TOKEN       API key (also checked as OPENAI_API_KEY or API_KEY)
"""
from __future__ import annotations
import os
import sys
from typing import Any, Dict, List

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

from grader.grader import evaluate_episode

# Configuration
API_KEY      = os.getenv("HF_TOKEN") or os.getenv("OPENAI_API_KEY") or os.getenv("API_KEY")
API_BASE_URL = os.getenv("API_BASE_URL", "https://router.huggingface.co/v1")
MODEL_NAME   = os.getenv("MODEL_NAME",   "Qwen/Qwen2.5-72B-Instruct")

VALID_ACTIONS = ["restart_api", "restart_db", "restart_cache", "scale_up", "noop"]

# OpenAI client (spec requires OpenAI client for all LLM calls)
_client: Any = None
if _OPENAI_AVAILABLE and API_KEY:
    _client = OpenAI(api_key=API_KEY, base_url=API_BASE_URL)


# Heuristic fallback agent (used when no LLM key or on API error)
def _fallback_agent(state: Dict[str, Any], action_history: List[str]) -> str:
    """
    Rule-based agent driven by log keywords and fix_progress from state.
    Uses fix_progress from the observation to know exactly which step to issue next.
    """
    logs = " ".join(state.get("logs", [])).lower()
    metrics = state.get("metrics", {})
    progress = state.get("fix_progress", 0)

    # Known fix sequences keyed by log signal (ordered by specificity)
    _log_to_seq = [
        ("timeout",     ["scale_up",      "restart_api"]),
        ("upstream",    ["scale_up",      "restart_api"]),
        ("heap",        ["restart_api",   "restart_db"]),
        ("oom",         ["restart_api",   "restart_db"]),
        ("memory",      ["restart_api",   "restart_db"]),
        ("db pool",     ["restart_db",    "scale_up"]),
        ("connections", ["restart_db",    "scale_up"]),
        ("cache miss",  ["restart_cache", "scale_up"]),
        ("cache",       ["restart_cache", "scale_up"]),
        ("db",          ["restart_db",    "scale_up"]),
    ]

    # Identify sequence from logs
    seq = None
    for keyword, candidate_seq in _log_to_seq:
        if keyword in logs:
            seq = candidate_seq
            break

    if seq and progress < len(seq):
        # Issue the next step in the correct sequence
        return seq[progress]

    # Metric-based fallback
    if metrics.get("cpu", 0) > 80 and "scale_up" not in action_history:
        return "scale_up"
    if metrics.get("error_rate", 0) > 0.5 and "restart_api" not in action_history:
        return "restart_api"

    # Try any action not yet used twice
    for a in ["restart_api", "restart_db", "restart_cache", "scale_up"]:
        if action_history.count(a) < 2:
            return a

    return "noop"

# LLM call (OpenAI client as required)
def _call_llm(
    state: Dict[str, Any],
    action_history: List[str],
    reward_history: List[float],
) -> str:
    if _client is None:
        return _fallback_agent(state, action_history)

    # Build a human-readable step history with outcome signals
    step_lines = []
    for i, (a, r) in enumerate(zip(action_history, reward_history)):
        outcome = "✓ progress made" if r > 20 else ("✗ wrong / no effect" if r < -5 else "~ neutral")
        step_lines.append(f"  step {i}: {a:18s} reward={r:>8.1f}  [{outcome}]")
    history_block = "\n".join(step_lines) if step_lines else "  (none yet)"

    services_degraded = [s for s, h in state["services"].items() if h == "degraded"]
    metrics = state["metrics"]

    prompt = f"""You are an expert SRE triaging a production incident. Your goal is to resolve it in as few steps as possible.

SYSTEM STATE (step {state['time_step']}) 
Degraded services : {services_degraded if services_degraded else 'none'}
Metrics           : latency={metrics.get('latency', 0):.0f}ms  error_rate={metrics.get('error_rate', 0):.2%}  cpu={metrics.get('cpu', 0):.0f}%
Metric trend      : {state.get('metric_trend', 'unknown')}
Fix progress      : {state.get('fix_progress', 0)} step(s) completed correctly so far

SYSTEM LOGS 
{chr(10).join('  ' + l for l in state['logs'])}

ACTION HISTORY & OUTCOMES 
{history_block}

INSTRUCTIONS 
Root causes have multi-step fix sequences that MUST be performed in order.
A positive reward means the last action was a correct step — continue the sequence.
A negative reward means the last action was wrong — try something different.
Do NOT repeat an action that already got a negative reward.
If fix_progress increased after your last action, continue to the NEXT step in the sequence.

Choose ONE action from: restart_api, restart_db, restart_cache, scale_up, noop

Respond with ONLY the action name."""

    try:
        response = _client.chat.completions.create(
            model=MODEL_NAME,
            messages=[{"role": "user", "content": prompt}],
            temperature=0,
            max_tokens=20,
            timeout=15,
        )
        action = response.choices[0].message.content.strip().lower()
        action = action.split()[0] if action else "noop"
        # strip any punctuation
        action = "".join(c for c in action if c.isalnum() or c == "_")

        if action not in VALID_ACTIONS:
            return _fallback_agent(state, action_history)
        return action

    except Exception:
        return _fallback_agent(state, action_history)


def run_episode(task_name: str = "simple") -> None:
    if task_name == "simple":
        from tasks.task_simple import create_env
        max_steps = 15
    elif task_name == "multi_service":
        from tasks.task_multi_service import create_env
        max_steps = 12
    elif task_name == "critical":
        from tasks.task_critical import create_env
        max_steps = 10
    else:
        raise ValueError(f"Unknown task: {task_name!r}")

    env = create_env()
    state = env.reset()

    print(
        f"[START] task={task_name} env=debugops model={MODEL_NAME}",
        flush=True,
    )

    action_history: List[str] = []
    reward_history: List[float] = []
    episode_logs:   List[Dict[str, Any]] = []
    done = False
    step = 0

    while not done and step < max_steps:
        action = _call_llm(state, action_history, reward_history)

        try:
            next_state, reward, done, info = env.step(action)
            error = "null"
        except Exception as exc:
            next_state, reward, done, info = state, 0.0, True, {}
            error = type(exc).__name__

        action_history.append(action)
        reward_history.append(reward)
        episode_logs.append({"reward": reward, "info": info})

        print(
            f"[STEP] step={step} action={action} reward={round(reward, 3)} "
            f"done={str(done).lower()} error={error}",
            flush=True,
        )

        # Early exit on explicit success flag
        if info.get("success", False):
            done = True

        state = next_state
        step += 1

    result = evaluate_episode(episode_logs, max_steps=max_steps)
    score   = result["score"]                    # already in [0, 1]
    success = result["resolved"]

    print(
        f"[END] success={str(success).lower()} steps={step} "
        f"score={round(score, 3)} "
        f"rewards={','.join(str(round(e['reward'], 3)) for e in episode_logs)}",
        flush=True,
    )


if __name__ == "__main__":
    for task in ["simple", "multi_service", "critical"]:
        run_episode(task)