""" inference.py — OpenEnv standard inference entry point for Repair Strategy System V2. Produces structured output required by the OpenEnv validator. Checklist Compliance: 1. Environment variables: API_BASE_URL, MODEL_NAME, HF_TOKEN, LOCAL_IMAGE_NAME 2. OpenAI client configured via these variables. 3. [START], [STEP], [END] tags strictly followed. """ from __future__ import annotations import sys import os import argparse import random import json from openai import OpenAI # Ensure the repo root is on the path sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) from env_v2 import RepairEnvV2, COMPONENTS from models_v2 import Action, ActionType, ComponentTarget from tasks_v2 import TASKS # --------------------------------------------------------------------------- # OpenEnv Configuration (Pre-Submission Checklist) # --------------------------------------------------------------------------- API_BASE_URL = os.getenv("API_BASE_URL", "https://api-inference.huggingface.co/v1/") MODEL_NAME = os.getenv("MODEL_NAME", "gpt-4o") HF_TOKEN = os.getenv("HF_TOKEN") # NO DEFAULT LOCAL_IMAGE_NAME = os.getenv("LOCAL_IMAGE_NAME") # Initialize OpenAI client # Note: If HF_TOKEN is missing, client creation might succeed but calls will fail. # We handle this by falling back to the baseline agent. client = OpenAI(base_url=API_BASE_URL, api_key=HF_TOKEN) if HF_TOKEN else None # --------------------------------------------------------------------------- # Agents # --------------------------------------------------------------------------- class InferenceAgent: """Baseline deterministic agent for reliable evaluation.""" def __init__(self): self.diagnosed_roots = set() self.inspected = set() self.fixed_roots = set() self.wait_ticks = 0 self._rng = random.Random(42) def pick_action(self, obs) -> Action: # Flaw: 20% of the time it just randomly restarts something if self._rng.random() < 0.20: target = self._rng.choice(["api", "queue", "cache"]) return Action(action_type=ActionType.restart_service, target=ComponentTarget(target)) # 1. If we know any root causes we haven't fixed yet, fix them. pending_roots = self.diagnosed_roots - self.fixed_roots if pending_roots: target = list(pending_roots)[0] self.fixed_roots.add(target) self.wait_ticks = 2 if target == "database": return Action(action_type=ActionType.repair_database, target=ComponentTarget.database) elif target == "cache": return Action(action_type=ActionType.clear_cache, target=ComponentTarget.cache) else: return Action(action_type=ActionType.restart_service, target=ComponentTarget(target)) # Cooldown period if not pending_roots and self.fixed_roots: if self.wait_ticks > 0: self.wait_ticks -= 1 return Action(action_type=ActionType.no_op, target=ComponentTarget.api) worst = max(obs.metrics.items(), key=lambda x: x[1].error_rate) if worst[1].error_rate > 0.0: if worst[0] == "database": return Action(action_type=ActionType.repair_database, target=ComponentTarget.database) elif worst[0] == "cache": return Action(action_type=ActionType.clear_cache, target=ComponentTarget.cache) else: return Action(action_type=ActionType.restart_service, target=ComponentTarget(worst[0])) return Action(action_type=ActionType.no_op, target=ComponentTarget.api) # 2. Pick most suspicious component to inspect suspicious = None highest_err = -1.0 for comp, m in obs.metrics.items(): if comp not in self.inspected and m.error_rate > highest_err: highest_err = m.error_rate suspicious = comp if suspicious and highest_err > 0.1: self.inspected.add(suspicious) action_map = { "database": ActionType.inspect_database, "api": ActionType.inspect_api, "cache": ActionType.inspect_cache, "queue": ActionType.inspect_queue } return Action(action_type=action_map[suspicious], target=ComponentTarget(suspicious)) # Fallback: just restart highest error rate worst = max(obs.metrics.items(), key=lambda x: x[1].error_rate) if worst[1].error_rate > 0.1: if worst[0] == "database": return Action(action_type=ActionType.repair_database, target=ComponentTarget.database) elif worst[0] == "cache": return Action(action_type=ActionType.clear_cache, target=ComponentTarget.cache) else: return Action(action_type=ActionType.restart_service, target=ComponentTarget(worst[0])) return Action(action_type=ActionType.no_op, target=ComponentTarget.api) def llm_agent(obs, task_config: dict, step: int) -> Action: """LLM-based agent using the OpenAI client.""" if not client: # If no client, we can't call LLM return None # Logic similar to previous version... obs_dict = obs.model_dump() if hasattr(obs, 'model_dump') else obs system_prompt = "You are an SRE agent debugging a distributed system. Respond with JSON: {'action_type': '...', 'target': '...'}" try: completion = client.chat.completions.create( model=MODEL_NAME, messages=[ {"role": "system", "content": system_prompt}, {"role": "user", "content": f"Step: {step}\\nObservation: {json.dumps(obs_dict)}"} ], response_format={"type": "json_object"} ) data = json.loads(completion.choices[0].message.content) return Action( action_type=ActionType(data.get("action_type", "no_op")), target=ComponentTarget(data.get("target", "database")) ) except Exception: return None # --------------------------------------------------------------------------- # Runner logic # --------------------------------------------------------------------------- def run_task(task_id: str): """Run a single task and print structured output for the validator.""" if task_id not in TASKS: return config = TASKS[task_id] # 1. Start tag print(f"[START] task={task_id}", flush=True) env = RepairEnvV2() obs = env.reset(task_config=config, task_id=task_id) baseline_agent = InferenceAgent() done = False step = 0 while not done: # Try LLM agent first, fallback to baseline action = llm_agent(obs, config, step) if action is None: action = baseline_agent.pick_action(obs) obs, reward, done, info = env.step(action) step += 1 # Update baseline agent state res = info.get("inspect_result") if res and res["is_likely_root_cause"] and res["confidence"] >= 0.7: baseline_agent.diagnosed_roots.add(res["component"]) # 2. Step tag print(f"[STEP] step={step} reward={reward.value:.3f}", flush=True) # 3. End tag result = env.get_episode_result() print(f"[END] task={task_id} score={result.score:.4f} steps={result.steps_taken}", flush=True) def main(): parser = argparse.ArgumentParser(description="OpenEnv inference runner") parser.add_argument("--task", help="Run specific task (easy, medium, hard, multi_root)") args = parser.parse_args() if args.task: tasks_to_run = [args.task] else: tasks_to_run = ["easy", "medium", "hard", "multi_root"] for tid in tasks_to_run: run_task(tid) if __name__ == "__main__": main()