Spaces:
Sleeping
Sleeping
| """ | |
| Baseline heuristic agent for SupportOps Arena. | |
| Rule-based agent using decision tree logic (no ML required). | |
| """ | |
| import logging | |
| from typing import Any | |
| from env.state import EnvObservation, NetworkStatus, VPNStatus, AuthStatus, ServiceHealth | |
| from env.actions import ActionType | |
| from env.environment import SupportOpsArena | |
| logger = logging.getLogger(__name__) | |
| class BaselineAgent: | |
| """ | |
| Scripted heuristic baseline agent. | |
| Uses a simple decision tree strategy: | |
| 1. Gather initial evidence (logs + auth) | |
| 2. Apply rule-based diagnostics based on observations | |
| 3. Choose lowest-risk remediation | |
| 4. Escalate if uncertain after gathering evidence | |
| Expected performance: | |
| - Easy: 0.65-0.75 | |
| - Medium: 0.40-0.55 | |
| - Hard: 0.20-0.35 | |
| """ | |
| def __init__(self, seed: int | None = None): | |
| """ | |
| Initialize baseline agent. | |
| Args: | |
| seed: Random seed for reproducibility (currently unused as agent is deterministic) | |
| """ | |
| self.seed = seed | |
| self._action_history: list[str] = [] | |
| async def select_action(self, observation: EnvObservation) -> ActionType: | |
| """ | |
| Select next action based on current observation. | |
| Args: | |
| observation: Current environment observation | |
| Returns: | |
| ActionType to take | |
| """ | |
| step = observation.step_count | |
| max_steps = observation.steps_remaining + step | |
| # Phase 1: Initial diagnostic (steps 1-2) | |
| if step == 0: | |
| return ActionType.INSPECT_LOGS | |
| if step == 1: | |
| return ActionType.CHECK_AUTHENTICATION | |
| # Phase 2: Targeted diagnostics (steps 2-4) | |
| if step < 4: | |
| # Check auth status | |
| if observation.auth_status in [ | |
| AuthStatus.LOCKED, | |
| AuthStatus.EXPIRED, | |
| AuthStatus.MFA_FAIL | |
| ]: | |
| if "check_permissions" not in self._action_history: | |
| self._action_history.append("check_permissions") | |
| return ActionType.CHECK_PERMISSIONS | |
| # Check network status | |
| if observation.network_status in [NetworkStatus.DEGRADED, NetworkStatus.DOWN]: | |
| if "inspect_network" not in self._action_history: | |
| self._action_history.append("inspect_network") | |
| return ActionType.INSPECT_NETWORK | |
| # Check VPN status | |
| if observation.vpn_status in [VPNStatus.FAILED, VPNStatus.TIMEOUT]: | |
| if "query_device_status" not in self._action_history: | |
| self._action_history.append("query_device_status") | |
| return ActionType.QUERY_DEVICE_STATUS | |
| # Check for service issues | |
| has_service_issues = any( | |
| status in [ServiceHealth.DEGRADED, ServiceHealth.DOWN] | |
| for status in observation.service_health.values() | |
| ) | |
| if has_service_issues: | |
| if "run_diagnostic_script" not in self._action_history: | |
| self._action_history.append("run_diagnostic_script") | |
| return ActionType.RUN_DIAGNOSTIC | |
| # Default: query device | |
| if "query_device_status" not in self._action_history: | |
| self._action_history.append("query_device_status") | |
| return ActionType.QUERY_DEVICE_STATUS | |
| # Phase 3: Apply remediation (steps 4+) | |
| # Determine remediation based on observations | |
| remediation = self._determine_remediation(observation) | |
| if remediation: | |
| if remediation.value not in self._action_history: | |
| self._action_history.append(remediation.value) | |
| return remediation | |
| # Phase 4: Resolve or escalate | |
| # Check if we have enough evidence | |
| diagnostic_count = len(self._action_history) | |
| evidence_threshold = max_steps * 0.6 | |
| # If we've gathered enough evidence and applied a fix, resolve | |
| if diagnostic_count >= 4 and len(self._action_history) >= 5: | |
| # Check if escalation is available and we're uncertain | |
| if observation.escalation_allowed and diagnostic_count < evidence_threshold * 0.8: | |
| self._action_history.append("escalate_ticket") | |
| return ActionType.ESCALATE_TICKET | |
| # Otherwise resolve | |
| self._action_history.append("resolve_ticket") | |
| return ActionType.RESOLVE_TICKET | |
| # Need more diagnostics or stuck - escalate if possible | |
| if observation.escalation_allowed and diagnostic_count >= evidence_threshold: | |
| self._action_history.append("escalate_ticket") | |
| return ActionType.ESCALATE_TICKET | |
| # Gather more evidence | |
| if "search_internal_kb" not in self._action_history: | |
| self._action_history.append("search_internal_kb") | |
| return ActionType.SEARCH_INTERNAL_KB | |
| # Last resort: resolve | |
| self._action_history.append("resolve_ticket") | |
| return ActionType.RESOLVE_TICKET | |
| def _determine_remediation(self, observation: EnvObservation) -> ActionType | None: | |
| """ | |
| Determine appropriate remediation based on observations. | |
| Returns: | |
| Remediation action or None if more diagnostics needed | |
| """ | |
| # Auth issues → reset credentials | |
| if observation.auth_status in [ | |
| AuthStatus.LOCKED, | |
| AuthStatus.EXPIRED, | |
| AuthStatus.MFA_FAIL | |
| ]: | |
| if observation.diagnostic_steps_taken >= 2: | |
| return ActionType.RESET_CREDENTIALS | |
| # Network/DNS issues → flush DNS or reconfigure | |
| if observation.network_status in [NetworkStatus.DEGRADED, NetworkStatus.DOWN]: | |
| if observation.diagnostic_steps_taken >= 2: | |
| # Check if it's likely DNS | |
| dns_indicators = sum( | |
| 1 for log in observation.system_logs | |
| if "dns" in log.message.lower() | |
| ) | |
| if dns_indicators > 0: | |
| return ActionType.FLUSH_DNS | |
| else: | |
| return ActionType.RECONFIGURE_CLIENT | |
| # VPN issues → reconfigure client | |
| if observation.vpn_status in [VPNStatus.FAILED, VPNStatus.TIMEOUT]: | |
| if observation.diagnostic_steps_taken >= 2: | |
| return ActionType.RECONFIGURE_CLIENT | |
| # Service issues → restart service (high risk, so be cautious) | |
| has_critical_service_down = any( | |
| status == ServiceHealth.DOWN | |
| for status in observation.service_health.values() | |
| ) | |
| if has_critical_service_down: | |
| if observation.diagnostic_steps_taken >= 3: | |
| return ActionType.RESTART_SERVICE | |
| # No clear remediation yet | |
| return None | |
| async def run_episode( | |
| self, | |
| env: SupportOpsArena, | |
| task_level: str = "easy" | |
| ) -> dict[str, Any]: | |
| """ | |
| Run one complete episode. | |
| Args: | |
| env: Environment instance | |
| task_level: Difficulty level | |
| Returns: | |
| Dictionary with episode results | |
| """ | |
| # Reset tracking | |
| self._action_history = [] | |
| # Reset environment | |
| observation = await env.reset(task_level=task_level, seed=self.seed) | |
| total_reward = 0.0 | |
| done = False | |
| steps = 0 | |
| logger.info(f"Starting baseline agent episode on {task_level}") | |
| # Run episode | |
| while not done and steps < 50: # Safety limit | |
| # Select action | |
| action = await self.select_action(observation) | |
| # Take step | |
| result = await env.step(action) | |
| observation = result.observation | |
| total_reward += result.reward | |
| done = result.done | |
| steps += 1 | |
| logger.debug( | |
| f"Step {steps}: {action.value} -> reward={result.reward:.2f}, " | |
| f"done={done}" | |
| ) | |
| # Get final state | |
| final_state = await env.state() | |
| logger.info( | |
| f"Episode completed: steps={steps}, " | |
| f"reward={total_reward:.2f}, " | |
| f"score={final_state.cumulative_reward:.2f}" | |
| ) | |
| return { | |
| "steps": steps, | |
| "total_reward": total_reward, | |
| "cumulative_reward": final_state.cumulative_reward, | |
| "done": done, | |
| "task_level": task_level, | |
| } | |
| def reset(self) -> None: | |
| """Reset agent state for new episode.""" | |
| self._action_history = [] | |