| """ |
| SRE Incident Investigation Environment — Python Client. |
| |
| EnvClient is async by default. Use .sync() for synchronous code or |
| 'async with' for async code. |
| |
| Usage (sync — simplest): |
| from client import SREEnvClient |
| from models import SREAction |
| |
| with SREEnvClient(base_url="http://localhost:8000").sync() as env: |
| result = env.reset(task_id="sre-easy-001") |
| result = env.step(SREAction(action_type="query_alerts")) |
| result = env.step(SREAction( |
| action_type="submit", |
| root_cause_service="payment-service", |
| root_cause_type="resource_exhaustion", |
| affected_services=["payment-service", "api-gateway", "order-service"], |
| severity="P2", |
| recommended_action="Increase JVM heap memory limit", |
| confidence=0.9, |
| )) |
| print("Score:", result.observation.grader_score) |
| |
| Usage (async — for training loops): |
| import asyncio |
| from client import SREEnvClient |
| from models import SREAction |
| |
| async def main(): |
| async with SREEnvClient(base_url="http://localhost:8000") as env: |
| result = await env.reset_async(task_id="sre-hard-003") |
| result = await env.step_async(SREAction(action_type="query_alerts")) |
| result = await env.step_async(SREAction( |
| action_type="submit", |
| root_cause_service="recommendation-service", |
| root_cause_type="configuration_error", |
| affected_services=["recommendation-service", "cart-service"], |
| severity="P1", |
| recommended_action="Rollback feature flag config", |
| confidence=0.9, |
| )) |
| print("Score:", result.observation.grader_score) |
| |
| asyncio.run(main()) |
| """ |
|
|
| from typing import Dict |
|
|
| from openenv.core import EnvClient |
| from openenv.core.client_types import StepResult |
|
|
| try: |
| from .models import SREAction, SREObservation, SREState |
| except ImportError: |
| from models import SREAction, SREObservation, SREState |
|
|
|
|
| class SREEnvClient(EnvClient[SREAction, SREObservation, SREState]): |
| """Typed WebSocket client for the SRE Incident Investigation environment.""" |
|
|
| def _step_payload(self, action: SREAction) -> Dict: |
| payload = {"action_type": action.action_type} |
| for field in ["service", "log_level", "time_window_minutes", "log_query", |
| "metric_name", "note", "root_cause_service", "root_cause_type", |
| "affected_services", "severity", "recommended_action", "confidence"]: |
| v = getattr(action, field, None) |
| if v is not None: |
| payload[field] = v |
| return payload |
|
|
| def _parse_result(self, payload: Dict) -> "StepResult[SREObservation]": |
| obs_data = payload.get("observation", payload) |
| observation = SREObservation( |
| action_taken=obs_data.get("action_taken", ""), |
| logs=obs_data.get("logs", []), |
| metrics=obs_data.get("metrics", []), |
| metric_name=obs_data.get("metric_name"), |
| alerts=obs_data.get("alerts", []), |
| annotation_accepted=obs_data.get("annotation_accepted", False), |
| grader_score=obs_data.get("grader_score"), |
| grader_breakdown=obs_data.get("grader_breakdown"), |
| message=obs_data.get("message", ""), |
| queries_remaining=obs_data.get("queries_remaining", 0), |
| done=payload.get("done", obs_data.get("done", False)), |
| reward=payload.get("reward", obs_data.get("reward")), |
| metadata=obs_data.get("metadata", {}), |
| ) |
| return StepResult(observation=observation, reward=payload.get("reward"), done=payload.get("done", False)) |
|
|
| def _parse_state(self, payload: Dict) -> SREState: |
| return SREState( |
| episode_id=payload.get("episode_id"), |
| task_id=payload.get("task_id", ""), |
| difficulty=payload.get("difficulty", ""), |
| step_count=payload.get("step_count", 0), |
| queries_used=payload.get("queries_used", 0), |
| max_queries=payload.get("max_queries", 12), |
| annotations=payload.get("annotations", []), |
| submitted=payload.get("submitted", False), |
| final_score=payload.get("final_score"), |
| ) |
|
|
