""" app.py — OpenEnv Split-Brain Collapse Environment ================================================== Serves: • FastAPI HTTP API — /reset, /step, /state, /health, /tasks • Agent endpoints — /agents, /agent/step • Static HTML UI — GET / → static/index.html All on port 7860 for Hugging Face Spaces. ARCHITECTURE: agents/__init__.py → AGENT_REGISTRY (central config) agents// → each agent's env_class + task metadata Add a new agent = new folder + one import + one dict entry """ import os from dotenv import load_dotenv from typing import Optional from fastapi import FastAPI, HTTPException, Request from fastapi.responses import JSONResponse from fastapi.staticfiles import StaticFiles from pydantic import BaseModel, Field import uvicorn from openai import OpenAI from agents import AGENT_REGISTRY from agents.split_brain.models import SplitBrainAction load_dotenv() # ── LLM Config ────────────────────────────────────────────────────────────── API_BASE_URL = os.getenv("API_BASE_URL", "https://router.huggingface.co/v1") MODEL_NAME = os.getenv("MODEL_NAME", "deepseek-ai/DeepSeek-R1-Distill-Llama-70B") SPLIT_BRAIN_MODEL = os.getenv("SPLIT_BRAIN_MODEL", "") # GRPO fine-tuned LoRA model API_KEY = os.getenv("HF_TOKEN", "").strip().strip('"').strip("'") if API_KEY: llm_client = OpenAI( base_url=API_BASE_URL, api_key=API_KEY, timeout=300.0, ) print(f"[INFO] LLM ready. Token: {API_KEY[:6]}... Model: {MODEL_NAME}") else: llm_client = None print("[WARN] HF_TOKEN not set — LLM calls will fail.") # ── Active Environment State ───────────────────────────────────────────────── active_agent_id = None active_env = None def get_or_create_env(agent_id: str): """Lazily create/switch the active environment when agent changes.""" global active_agent_id, active_env agent = AGENT_REGISTRY.get(agent_id) if not agent: raise HTTPException( status_code=400, detail=f"Unknown agent '{agent_id}'. Available: {list(AGENT_REGISTRY.keys())}", ) if active_agent_id != agent_id: active_agent_id = agent_id active_env = agent["env_class"]() return active_env # ── FastAPI App ─────────────────────────────────────────────────────────────── app = FastAPI( title="OpenEnv: Split-Brain Collapse", description=( "An OpenEnv-compliant multi-agent RL environment simulating a " "three-datacenter split-brain network partition crisis. " "AI agents acting as orchestrator, netops, and dba must collaboratively " "resolve network partitions, replication storms, and cascading deadlocks." ), version="1.0.0", ) # ══════════════════════════════════════════════════════════════════════════════ # CORE RL API (openenv validate + external agents) # ══════════════════════════════════════════════════════════════════════════════ @app.get("/health") def health(): """Health check — returns 200 OK.""" return { "status": "ok", "env": "split-brain", "version": "1.0.0", "llm_configured": bool(API_KEY), "model": MODEL_NAME, } @app.post("/reset") async def reset(request: Request): """ Reset the environment for a given task. Body: {"task": "partition_basic", "agent_id": "split_brain"} """ try: body = await request.json() except Exception: body = {} agent_id = body.get("agent_id", "split_brain") task = body.get("task", "partition_basic") env = get_or_create_env(agent_id) agent = AGENT_REGISTRY[agent_id] valid = [t["id"] for t in agent["tasks"]] if task not in valid: raise HTTPException( status_code=400, detail=f"Unknown task '{task}'. Valid: {valid}", ) obs = env.reset(task=task) return JSONResponse(content=obs.model_dump()) @app.post("/step") async def step(request: Request): """Execute one raw action in the environment.""" if active_env is None: raise HTTPException(status_code=400, detail="No environment loaded. Call /reset first.") body = await request.json() action = SplitBrainAction(**body) result = active_env.step(action) return JSONResponse(content=result.model_dump()) @app.get("/state") def state(): """Return the current observation without advancing the episode.""" if active_env is None: raise HTTPException(status_code=400, detail="No environment loaded. Call /reset first.") return JSONResponse(content=active_env.state().model_dump()) @app.get("/tasks") def list_tasks(agent_id: str = "split_brain"): """List all available tasks with difficulty metadata for a specific agent.""" if agent_id not in AGENT_REGISTRY: raise HTTPException(status_code=404, detail="Agent not found.") return {"tasks": AGENT_REGISTRY[agent_id]["tasks"]} # ══════════════════════════════════════════════════════════════════════════════ # AGENT REGISTRY API (frontend discovers agents + tasks) # ══════════════════════════════════════════════════════════════════════════════ @app.get("/agents") def list_agents(): """ Return all registered agents with their tasks. The frontend sidebar auto-populates from this endpoint. Add new agents to agents/__init__.py — no HTML changes needed. """ return [ { "id": agent["id"], "name": agent["name"], "icon": agent["icon"], "description": agent["description"], "tasks": agent["tasks"], } for agent in AGENT_REGISTRY.values() ] # ══════════════════════════════════════════════════════════════════════════════ # LLM AGENT STEP (used by the HTML UI's ▶ Agent Step button) # ══════════════════════════════════════════════════════════════════════════════ class AgentStepRequest(BaseModel): agent_id: str = Field("split_brain", description="Which agent (environment) to step") task: Optional[str] = Field(None, description="Task hint (set via /reset)") @app.post("/agent/step") def agent_step(body: AgentStepRequest): """ Run ONE LLM-powered step on the active Split-Brain environment. The environment provides its own multi-agent prompts via get_llm_prompts(). """ if not llm_client: raise HTTPException(status_code=503, detail="HF_TOKEN not configured.") env = get_or_create_env(body.agent_id) obs = env.state() # Check episode completion health = getattr(obs, "global_health", None) or getattr(obs, "health_score", 0) if health >= 1.0 or env.step_count >= env.max_steps: raise HTTPException(status_code=400, detail="Episode complete. Call /reset to start a new one.") # The split_brain env always exposes get_llm_prompts() for multi-agent routing system_prompt, user_prompt = env.get_llm_prompts() # Auto-select model: prefer fine-tuned LoRA model for cascading_deadlock active_model = MODEL_NAME if SPLIT_BRAIN_MODEL and getattr(env, "current_task", "") == "cascading_deadlock": active_model = SPLIT_BRAIN_MODEL print(f"[INFO] Using fine-tuned model: {active_model}") try: completion = llm_client.chat.completions.create( model=active_model, messages=[ {"role": "system", "content": system_prompt}, {"role": "user", "content": user_prompt}, ], temperature=0.1, max_tokens=400, ) raw_text = completion.choices[0].message.content or "" print(f"[DEBUG LLM ({active_model})] {raw_text}") except Exception as e: raise HTTPException(status_code=502, detail=f"LLM call failed: {str(e)}") action = env._parse_action(raw_text) result = env.step(action) new_obs = result.observation msg = result.info.get("message", "") return { "step": env.step_count, "agent_id": body.agent_id, "action": action.model_dump(), "reward": result.reward, "done": result.done, "message": msg, "observation": new_obs.model_dump(), "current_actor": result.info.get("current_actor", "orchestrator"), "delegation_log": result.info.get("delegation_log", []), } # ══════════════════════════════════════════════════════════════════════════════ # STATIC FILES — serves static/index.html at / # Mount AFTER all API routes so API routes take priority. # ══════════════════════════════════════════════════════════════════════════════ app.mount("/", StaticFiles(directory="static", html=True), name="static") if __name__ == "__main__": # Change "app:app" (implicit) to "app:fastapi_app" uvicorn.run(app, host="0.0.0.0", port=7860, log_level="info")