tasks/task_multi_service.py

# Task: Multi-Service: Forces two services into a degraded state simultaneously.
# Agent must correctly sequence actions across both services.
# Tighter step budget increases difficulty.

from __future__ import annotations
from typing import Dict, Any, Tuple

from env.environment import DebugEnv

class MultiServiceEnv(DebugEnv):
    """Both API and DB start degraded regardless of sampled root cause."""

    def reset(self) -> Dict[str, Any]:
        state = super().reset()
        self.state_data["services"]["api"] = "degraded"
        self.state_data["services"]["db"]  = "degraded"
        # Increase starting metric pressure
        self.state_data["metrics"]["latency"]    = min(
            self.state_data["metrics"]["latency"] * 1.3, 600
        )
        self.state_data["metrics"]["error_rate"] = min(
            self.state_data["metrics"]["error_rate"] * 1.2, 0.95
        )
        return self._obs()

    def step(self, action: str) -> Tuple[Dict[str, Any], float, bool, Dict[str, Any]]:
        obs, reward, done, info = super().step(action)

        # Additional latency penalty for multi-service SLA
        if info["latency"] > 300:
            reward -= 10.0

        return obs, reward, done, info


def create_env() -> MultiServiceEnv:
    return MultiServiceEnv(max_steps=12)