Spaces:
Sleeping
Sleeping
| """ | |
| Baseline agent for V2: Inspect → Filter → Fix | |
| """ | |
| import sys | |
| import os | |
| import random | |
| sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) | |
| from env_v2 import RepairEnvV2, COMPONENTS | |
| from models_v2 import Action, ActionType, ComponentTarget, SymptomObservation | |
| class BaselineAgent: | |
| def __init__(self): | |
| self.diagnosed_roots = set() | |
| self.inspected = set() | |
| self.fixed_roots = set() | |
| self.wait_ticks = 0 | |
| self._rng = random.Random(42) # For deterministic flaws | |
| def pick_action(self, obs: SymptomObservation) -> Action: | |
| import random | |
| # Flaw: 20% of the time it just randomly restarts something, wasting a turn and cost | |
| 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 # Add cooldown ticks | |
| 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)) | |
| # If we fixed the root, clean up any remaining high-error symptoms wait... | |
| # Wait for cooldown to finish first | |
| 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 | |
| # Find component with highest error_rate among uninspected | |
| 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)) | |
| # 3. We've exhausted inspections or found nothing obvious (all healthy?). | |
| # If we have a diagnosed root from a previous step's info (we don't, agent only sees obs). | |
| # Wait, the observation doesn't contain the InspectResult! | |
| # Ah! The inspect result is in 'info', which the agent shouldn't rely on for pure observation, | |
| # but in OpenAI Gym, agents receive 'info'. Let's assume the agent can read `info["inspect_result"]` | |
| # from the previous step, or we just randomly fix degraded things. | |
| # 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 run_task(task_id: str, config: dict) -> dict: | |
| env = RepairEnvV2() | |
| obs = env.reset(config, task_id) | |
| agent = BaselineAgent() | |
| step = 0 | |
| print(f"\n=== Task: {task_id.upper()} ===") | |
| while True: | |
| action = agent.pick_action(obs) | |
| obs, reward, done, info = env.step(action) | |
| step += 1 | |
| root_list = info.get("root_causes") | |
| res = info.get("inspect_result") | |
| if res and res["is_likely_root_cause"] and res["confidence"] >= 0.7: | |
| agent.diagnosed_roots.add(res["component"]) | |
| print(f"Step {step:02d} | Action: {action.action_type}({action.target}) | Reward: {reward.value:+.2f}") | |
| for l in obs.logs: | |
| print(f" [LOG] {l.message}") | |
| if done: | |
| break | |
| result = env.get_episode_result() | |
| print(f"Success: {result.success} | Score: {result.score:.4f} | Steps: {result.steps_taken}") | |
| return result.model_dump() | |
| if __name__ == "__main__": | |
| from tasks_v2 import TASKS | |
| for tid, cfg in TASKS.items(): | |
| run_task(tid, cfg) | |