Spaces:
Sleeping
Sleeping
| """ | |
| 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() | |