"""Baseline and LLM inference runner for the incident triage environment.""" import json import os import re import time from openai import OpenAI from env.environment import DevOpsEnv from env.models import VALID_ACTIONS from graders.grader import compute_score SYSTEM_PROMPT = f"""You are the primary on-call SRE for a SaaS platform. Your job is to investigate production incidents, apply safe mitigations, communicate clearly, and resolve the incident only when the system is actually stable. Valid actions (pick EXACTLY one): {json.dumps(VALID_ACTIONS)} Guidance: - inspect first when the root cause is unclear - avoid noisy actions like repeated restarts or no_op - use rollback_auth_deploy for bad auth deploys - use scale_db_cluster and shift_traffic_canary for DB saturation - use post_status_update before resolve_incident when customer impact is ongoing Return EXACTLY one action name from the list above. Do not explain. Do not include markdown. Do not include multiple actions.""" BASELINE_PLANS = { "easy": [ "acknowledge_incident", "inspect_deploy_history", "rollback_auth_deploy", "resolve_incident", ], "medium": [ "acknowledge_incident", "inspect_db_metrics", "scale_db_cluster", "shift_traffic_canary", "post_status_update", "resolve_incident", ], "hard": [ "acknowledge_incident", "inspect_auth_logs", "inspect_deploy_history", "inspect_db_metrics", "rollback_auth_deploy", "scale_db_cluster", "shift_traffic_canary", "post_status_update", "resolve_incident", ], } FALLBACK_ACTIONS = { "easy": [ "inspect_auth_logs", "post_status_update", "restart_auth_service", "inspect_db_metrics", "no_op", ], "medium": [ "inspect_auth_logs", "inspect_deploy_history", "flush_cache", "restart_auth_service", "no_op", ], "hard": [ "restart_auth_service", "no_op", ], } def simple_policy(state: dict) -> str: """Deterministic, reproducible baseline policy.""" all_actions_taken = state.get("all_actions_taken", []) plan = BASELINE_PLANS[state["task"]] for action in plan: if action not in state["recent_actions"] and action not in all_actions_taken: return action if state["incident_phase"] == "monitoring" and "resolve_incident" not in all_actions_taken: return "resolve_incident" for action in FALLBACK_ACTIONS[state["task"]]: if action == "no_op": return action if action not in all_actions_taken: return action return "no_op" def extract_action(text: str) -> str | None: normalized = text.strip().lower() normalized = normalized.replace("`", " ") normalized = normalized.replace("\n", " ") if normalized in VALID_ACTIONS: return normalized for valid in VALID_ACTIONS: if re.search(rf"\b{re.escape(valid)}\b", normalized): return valid return None def extract_response_text(response) -> str: choice = response.choices[0] message = getattr(choice, "message", None) if message is None: return "" content = getattr(message, "content", None) if isinstance(content, str): return content if isinstance(content, list): parts: list[str] = [] for item in content: if isinstance(item, str): parts.append(item) elif isinstance(item, dict): text_value = item.get("text") or item.get("content") or "" if text_value: parts.append(str(text_value)) else: text_value = getattr(item, "text", None) or getattr(item, "content", None) if text_value: parts.append(str(text_value)) if parts: return " ".join(parts) for attr in ("reasoning_content", "refusal", "audio"): value = getattr(message, attr, None) if isinstance(value, str) and value.strip(): return value return "" def format_api_error(exc: Exception) -> str: status = getattr(exc, "status_code", None) body = getattr(exc, "body", None) if isinstance(body, dict): message = body.get("error", {}).get("message") or body.get("message") if message: compact = " ".join(str(message).split())[:80] return f"{type(exc).__name__}:{status}:{compact}" if status else f"{type(exc).__name__}:{compact}" return type(exc).__name__ def llm_policy(state: dict, client: OpenAI, model: str) -> tuple[str | None, str | None]: last_error = None for attempt in range(1, 3): try: response = client.chat.completions.create( model=model, messages=[ {"role": "system", "content": SYSTEM_PROMPT}, {"role": "user", "content": json.dumps(state)}, ], max_tokens=24, temperature=0.0, response_format={"type": "text"}, extra_body={"provider": "auto"}, ) raw_action = extract_response_text(response).strip() action = extract_action(raw_action) if action in VALID_ACTIONS: return action, None compact = " ".join(raw_action.split())[:80] or "empty_response" return None, f"llm_invalid_action:{compact}" except Exception as exc: last_error = format_api_error(exc) transient = "APIConnectionError" in last_error or "APIStatusError" in last_error if attempt < 2 and transient: time.sleep(1.0) continue return None, f"llm_fallback:{last_error}" return None, f"llm_fallback:{last_error or 'unknown_error'}" def run_task(task_name: str, client: OpenAI | None, model: str | None) -> tuple[float, bool]: env = DevOpsEnv(task=task_name) state = env.reset() use_llm = client is not None and model is not None print(f"[START] task={task_name} env=devops model={model if use_llm else 'baseline'}") rewards: list[str] = [] success = False score = 0.0 for step_num in range(1, env.max_steps + 1): state["all_actions_taken"] = list(env._state["actions_taken"]) action = None fallback_reason = None if use_llm: action, fallback_reason = llm_policy(state, client, model) if action is None: action = simple_policy(state) state, reward, done, info = env.step(action) rewards.append(f"{reward:.2f}") error_msg = fallback_reason or info.get("error", "null") print( f"[STEP] step={step_num} action={action} reward={reward:.2f} " f"done={str(done).lower()} error={error_msg}" ) if done: break score, _ = compute_score(task_name, env._state) success = env._state["resolved"] print( f"[END] success={str(success).lower()} steps={step_num} " f"score={score:.2f} rewards={','.join(rewards)}" ) return score, success if __name__ == "__main__": api_base = os.environ.get("API_BASE_URL", "https://router.huggingface.co/v1") model_name = os.environ.get("MODEL_NAME", "meta-llama/Llama-3.1-8B-Instruct") api_key = ( os.environ.get("OPENAI_API_KEY") or os.environ.get("API_KEY") or os.environ.get("HF_TOKEN") ) client = None if api_base and model_name: client = OpenAI(base_url=api_base, api_key=api_key or "no-key") for task in ["easy", "medium", "hard"]: run_task(task, client, model_name)