Improvements

inference.py

@@ -40,6 +40,31 @@ if _OPENAI_AVAILABLE and API_KEY:
     _client = OpenAI(api_key=API_KEY, base_url=API_BASE_URL)
 
 
+def infer_root_cause(state: Dict[str, Any]) -> str:
+    logs = " ".join(state.get("logs", [])).lower()
+    metrics = state.get("metrics", {})
+
+    if "timeout" in logs or metrics.get("latency", 0) > 1000:
+        return "traffic_overload"
+
+    if "connection" in logs or "db" in logs:
+        return "database_issue"
+
+    if "cache" in logs:
+        return "cache_issue"
+
+    if metrics.get("cpu", 0) > 85:
+        return "resource_exhaustion"
+
+    return "unknown"
+
+PLANS = {
+    "traffic_overload": ["scale_up", "restart_api"],
+    "database_issue": ["restart_db", "scale_up"],
+    "cache_issue": ["restart_cache", "scale_up"],
+    "resource_exhaustion": ["scale_up"],
+}
+
 # Heuristic fallback agent (used when no LLM key or on API error)
 def _fallback_agent(state: Dict[str, Any], action_history: List[str]) -> str:
     """
@@ -106,8 +131,9 @@ def _call_llm(
 
     services_degraded = [s for s, h in state["services"].items() if h == "degraded"]
     metrics = state["metrics"]
+    root_cause = infer_root_cause(state)
 
-    prompt = f"""You are an expert SRE triaging a production incident. Your goal is to resolve it in as few steps as possible.
+    prompt = f"""You are an expert SRE triaging a production incident. Your goal is to resolve it in as few steps as possible.ROOT CAUSE HYPOTHESIS: {root_cause}
 
 SYSTEM STATE (step {state['time_step']}) 
 Degraded services : {services_degraded if services_degraded else 'none'}
@@ -176,12 +202,28 @@ def run_episode(task_name: str = "simple") -> None:
 
     action_history: List[str] = []
     reward_history: List[float] = []
+    bad_actions = set()
+    plan: List[str] = []
     episode_logs:   List[Dict[str, Any]] = []
     done = False
     step = 0
 
     while not done and step < max_steps:
-        action = _call_llm(state, action_history, reward_history)
+        root_cause = infer_root_cause(state)
+
+        # generate plan if empty
+        if not plan:
+            plan = PLANS.get(root_cause, []).copy()
+
+        # choose from plan first
+        if plan:
+            action = plan.pop(0)
+        else:
+            action = _call_llm(state, action_history, reward_history)
+
+        # avoid bad actions
+        if action in bad_actions:
+            action = _fallback_agent(state, action_history)
 
         try:
             next_state, reward, done, info = env.step(action)
@@ -192,11 +234,15 @@ def run_episode(task_name: str = "simple") -> None:
 
         action_history.append(action)
         reward_history.append(reward)
+
+        if reward < -5:
+            bad_actions.add(action)
         episode_logs.append({"reward": reward, "info": info})
 
         print(
             f"[STEP] step={step} action={action} reward={round(reward, 3)} "
-            f"done={str(done).lower()} error={error}",
+            f"done={str(done).lower()} error={error} "
+            f"root_cause={root_cause} plan_remaining={len(plan)}",
             flush=True,
         )
 

tasks/task_critical.py

@@ -25,7 +25,7 @@ class CriticalEnv(DebugEnv):
 
     def reset(self) -> Dict[str, Any]:
         state = super().reset()
-
+        self.last_action = None
         self.state_data["root_cause"] = "memory_leak"
         self.state_data["fix_sequence"] = ["restart_api", "restart_db"]
 
@@ -64,15 +64,30 @@ class CriticalEnv(DebugEnv):
 
     def step(self, action: str) -> Tuple[Dict[str, Any], float, bool, Dict[str, Any]]:
         obs, reward, done, info = super().step(action)
+        if action == getattr(self, "last_action", None):
+            reward -= 5.0
+
+        self.last_action = action
 
         if info["latency"] > self.SLA_LATENCY_THRESHOLD:
             reward -= self.SLA_PENALTY
 
         if self.t > self.TIME_PRESSURE_START and not info["resolved"]:
             reward -= self.TIME_PRESSURE_PENALTY
+        if action == "scale_up":
+            reward -= 2.0  # cost penalty
+        if action == getattr(self, "last_action", None):
+            reward -= 5.0
+
 
         reward += random.uniform(-3, 3)
 
+        if info.get("resolved"):
+            info["explanation"] = "Incident resolved successfully"
+        elif reward > 0:
+            info["explanation"] = "Correct step towards resolution"
+        else:
+            info["explanation"] = "Ineffective or incorrect action"
         return obs, reward, done, info
 
 

tasks/task_multi_service.py

@@ -12,6 +12,7 @@ class MultiServiceEnv(DebugEnv):
 
     def reset(self) -> Dict[str, Any]:
         state = super().reset()
+        self.last_action = None
         self.state_data["services"]["api"] = "degraded"
         self.state_data["services"]["db"]  = "degraded"
         # Increase starting metric pressure
@@ -21,15 +22,28 @@ class MultiServiceEnv(DebugEnv):
         self.state_data["metrics"]["error_rate"] = min(
             self.state_data["metrics"]["error_rate"] * 1.2, 0.95
         )
+        self.dependencies = {
+            "api": ["db"],   # API depends on DB
+        }
         return self._obs()
 
     def step(self, action: str) -> Tuple[Dict[str, Any], float, bool, Dict[str, Any]]:
         obs, reward, done, info = super().step(action)
+        if action == getattr(self, "last_action", None):
+            reward -= 5.0
+
+        self.last_action = action
+        # Dependency impact
+        if self.state_data["services"]["db"] == "down":
+            self.state_data["services"]["api"] = "degraded"
 
         # Additional latency penalty for multi-service SLA
         if info["latency"] > 300:
             reward -= 10.0
-
+        if action == "scale_up":
+            reward -= 2.0  # cost penalty
+        
+        info["explanation"] = f"Action {action} applied to fix issue"
         return obs, reward, done, info
 
 

tasks/task_simple.py

@@ -1,6 +1,35 @@
-#Task: Simple: A single-service failure with a 2-step fix sequence.
-#Entry-level task - low penalty, generous step budget.
+from __future__ import annotations
+from typing import Dict, Any, Tuple
+import random
+
 from env.environment import DebugEnv
 
-def create_env() -> DebugEnv:
-    return DebugEnv(max_steps=15)
+
+class SimpleEnv(DebugEnv):
+
+    def reset(self) -> Dict[str, Any]:
+        state = super().reset()
+        self.last_action = None
+        # Add slight noise to logs
+        if random.random() < 0.2:
+            self.state_data["logs"].append("intermittent network glitch detected")
+
+        return self._obs()
+
+    def step(self, action: str) -> Tuple[Dict[str, Any], float, bool, Dict[str, Any]]:
+        obs, reward, done, info = super().step(action)
+        if action == getattr(self, "last_action", None):
+            reward -= 5.0
+
+        self.last_action = action
+
+        # Slight metric fluctuation
+        self.state_data["metrics"]["latency"] += random.randint(-20, 20)
+        self.state_data["metrics"]["cpu"] += random.randint(-3, 3)
+
+        info["explanation"] = f"Action {action} applied to fix issue"
+        return obs, reward, done, info
+
+
+def create_env() -> SimpleEnv:
+    return SimpleEnv(max_steps=15)