File size: 8,995 Bytes
9c1645f
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
"""
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 = []