Spaces:
Running
Running
| """ | |
| Oceanus 3D Mission Control — FastAPI + WebSocket backend | |
| Serves the 3D HTML frontend and streams live environment state. | |
| Run: uvicorn dashboard.server:app --host 0.0.0.0 --port 8000 --reload | |
| """ | |
| import sys, os, json, asyncio, re | |
| from typing import Dict, List, Optional, Set | |
| from pathlib import Path | |
| sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) | |
| from fastapi import FastAPI, WebSocket, WebSocketDisconnect | |
| from fastapi.responses import HTMLResponse, FileResponse | |
| from fastapi.staticfiles import StaticFiles | |
| from fastapi.middleware.cors import CORSMiddleware | |
| from oceanus.models import OceanusEnv | |
| from oceanus.adversary import AdversaryAgent | |
| from oceanus.runner import OceanusRunner, MockASVAgent, MockPolicyAgent | |
| app = FastAPI(title="Oceanus Mission Control", version="3.0.0") | |
| app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"], allow_credentials=True) | |
| # Serve static JS files (Three.js, Chart.js, OrbitControls) | |
| _static_dir = Path(__file__).parent / "static" | |
| if _static_dir.exists(): | |
| app.mount("/static", StaticFiles(directory=str(_static_dir)), name="static") | |
| # ── Groq LLM Agent ──────────────────────────────────────────────────────────── | |
| class GroqLLMAgent: | |
| """Live LLM agent powered by Groq API (llama3-8b-8192).""" | |
| def __init__(self, agent_id: str, api_key: str): | |
| self.agent_id = agent_id | |
| self.api_key = api_key | |
| self._client = None | |
| self._step = 0 | |
| self._last_action = None | |
| self._last_reasoning = None | |
| def _get_client(self): | |
| if self._client is None: | |
| try: | |
| from groq import Groq | |
| if not self.api_key: | |
| sim.warnings.append("LLM failed: No Groq API Key provided. Falling back to mock agents.") | |
| return None | |
| self._client = Groq(api_key=self.api_key) | |
| except ImportError: | |
| sim.warnings.append("LLM failed: 'groq' package not installed. Falling back to mock agents.") | |
| return None | |
| except Exception as e: | |
| sim.warnings.append(f"LLM failed: {e}. Falling back to mock agents.") | |
| return None | |
| return self._client | |
| def act(self, prompt: str, obs: Dict) -> str: | |
| self._step += 1 | |
| client = self._get_client() | |
| if client is None: | |
| return self._fallback(obs) | |
| try: | |
| resp = client.chat.completions.create( | |
| model="llama3-8b-8192", | |
| messages=[{"role": "user", "content": prompt}], | |
| max_tokens=80, | |
| temperature=0.3, | |
| timeout=3.0, | |
| ) | |
| raw = resp.choices[0].message.content.strip() | |
| self._last_action = raw | |
| # Extract reasoning (text before JSON) | |
| import re | |
| match = re.search(r'\{[^{}]*\}', raw, re.DOTALL) | |
| if match: | |
| self._last_reasoning = raw[:match.start()].strip() or "Analyzing environment..." | |
| return raw | |
| except Exception as e: | |
| if "rate limit" in str(e).lower() or "429" in str(e): | |
| sim.warnings.append("LLM failed: Rate limit exceeded. Falling back to mock agents.") | |
| else: | |
| sim.warnings.append(f"LLM failed: API Error. Falling back to mock agents.") | |
| return self._fallback(obs) | |
| def _fallback(self, obs: Dict) -> str: | |
| """Heuristic fallback if Groq unavailable.""" | |
| if self.agent_id.startswith("ASV"): | |
| sonar = obs.get("sonar_3x3", []) | |
| for dr in range(3): | |
| for dc in range(3): | |
| if len(sonar) > dr and len(sonar[dr]) > dc: | |
| if sonar[dr][dc] == "net": | |
| if dr == 1 and dc == 1: | |
| return '{"intent": "clean"}' | |
| dirs = {(0,0):"stay",(0,1):"east",(0,-1):"west",(1,0):"south",(-1,0):"north"} | |
| d = dirs.get((dr-1, dc-1), "north") | |
| return json.dumps({"intent": "move", "direction": d}) | |
| dirs = ["north","south","east","west"] | |
| return json.dumps({"intent": "move", "direction": dirs[self._step % 4]}) | |
| else: | |
| status = obs.get("current_policy_status", "") | |
| if status == "Treaty Proposed" and self.agent_id == "Fleet_Manager": | |
| return '{"intent": "accept_treaty", "target": "Port_Authority", "content": "Agreed. Fleet will comply."}' | |
| if status == "No Tagging Mandate" and self.agent_id == "Port_Authority": | |
| return '{"intent": "propose_treaty", "target": "Fleet_Manager", "content": "50% subsidy on tracking tags."}' | |
| inbox = obs.get("inbox", []) | |
| if inbox: | |
| sender = inbox[0].get("from", "Unknown") | |
| return json.dumps({"intent": "reply_email", "target": sender, "content": "We are addressing your concerns immediately."}) | |
| return '{"intent": "reply_email", "target": "General_Inbox", "content": "Monitoring situation."}' | |
| # ── Global simulation state ──────────────────────────────────────────────────── | |
| class SimState: | |
| def __init__(self): | |
| self.env: Optional[OceanusEnv] = None | |
| self.adversary: Optional[AdversaryAgent] = None | |
| self.runner: Optional[OceanusRunner] = None | |
| self.obs_all: Optional[Dict] = None | |
| self.running: bool = False | |
| self.done: bool = False | |
| self.total_reward: float = 0.0 | |
| self.reward_history: List[float] = [] | |
| self.bio_history: List[float] = [] | |
| self.net_history: List[int] = [] | |
| self.event_log: List[Dict] = [] | |
| self.episode_count: int = 0 | |
| self.agent_mode: str = "mock" # "mock" or "llm" | |
| self.groq_key: str = "" | |
| self.agent_thoughts: Dict[str, Dict] = {} # last action + reasoning per agent | |
| self.warnings: List[str] = [] | |
| self.cfg = { | |
| "seed": 42, "max_steps": 120, "chaos_interval": 25, | |
| "chaos_enabled": True, "step_delay": 0.08, "difficulty": "medium" | |
| } | |
| def reset(self, seed=42, max_steps=120, chaos_interval=25, agent_mode="mock", groq_key=""): | |
| self.env = OceanusEnv(seed=seed, max_steps=max_steps) | |
| self.adversary = AdversaryAgent(inject_interval=chaos_interval, seed=seed) | |
| self.agent_mode = agent_mode | |
| self.groq_key = groq_key | |
| if agent_mode == "llm": | |
| agents = { | |
| aid: GroqLLMAgent(aid, groq_key) | |
| for aid in ["ASV-1", "ASV-2", "ASV-3", "ASV-4", "Port_Authority", "Fleet_Manager"] | |
| } | |
| self.runner = OceanusRunner(self.env, self.adversary, agents=agents, verbose=False) | |
| else: | |
| self.runner = OceanusRunner(self.env, self.adversary, use_mock=True, verbose=False) | |
| self.obs_all = self.env.reset() | |
| self.running = False | |
| self.done = False | |
| self.total_reward = 0.0 | |
| self.reward_history = [] | |
| self.bio_history = [] | |
| self.net_history = [] | |
| self.event_log = [] | |
| self.agent_thoughts = {} | |
| self.episode_count += 1 | |
| def get_frame(self) -> Dict: | |
| if not self.env or not self.env.state: | |
| return {} | |
| s = self.env.state | |
| return { | |
| "type": "state", | |
| "step": s.step_count, | |
| "max_steps": s.max_steps, | |
| "grid": s.grid.tolist(), | |
| "asvs": { | |
| aid: { | |
| "row": a.row, "col": a.col, "battery": a.battery, | |
| "nets_cleaned": a.nets_cleaned, | |
| "on_net": bool(s.grid[a.row, a.col] > 0), | |
| "sector": s.get_sector(a.row, a.col), | |
| } | |
| for aid, a in s.asvs.items() | |
| }, | |
| "ghost_nets": [{"row": n.row, "col": n.col, "density": n.density} for n in s.ghost_nets], | |
| "biodiversity": round(s.biodiversity_index, 2), | |
| "active_nets": len(s.ghost_nets), | |
| "total_cleaned": s.total_cleaned, | |
| "treaty_status": s.treaty_status, | |
| "net_spawn_rate": s.net_spawn_rate, | |
| "wind_vector": list(s.wind_vector), | |
| "total_reward": round(self.total_reward, 1), | |
| "reward_history": self.reward_history[-60:], | |
| "bio_history": self.bio_history[-60:], | |
| "net_history": self.net_history[-60:], | |
| "event_log": self.event_log[-20:], | |
| "running": self.running, | |
| "done": self.done, | |
| "episode": self.episode_count, | |
| "schema_version": s.schema_version, | |
| "agent_mode": self.agent_mode, | |
| "agent_thoughts": self.agent_thoughts, | |
| } | |
| sim = SimState() | |
| sim.reset() | |
| # ── WebSocket connection manager ─────────────────────────────────────────────── | |
| class ConnectionManager: | |
| def __init__(self): | |
| self.active: Set[WebSocket] = set() | |
| async def connect(self, ws: WebSocket): | |
| await ws.accept() | |
| self.active.add(ws) | |
| def disconnect(self, ws: WebSocket): | |
| self.active.discard(ws) | |
| async def broadcast(self, data: Dict): | |
| msg = json.dumps(data) | |
| dead = set() | |
| for ws in self.active: | |
| try: | |
| await ws.send_text(msg) | |
| except Exception: | |
| dead.add(ws) | |
| self.active -= dead | |
| manager = ConnectionManager() | |
| # ── Background simulation loop ───────────────────────────────────────────────── | |
| async def simulation_loop(): | |
| while True: | |
| if hasattr(sim, "warnings") and sim.warnings: | |
| for w in sim.warnings: | |
| if w not in [log.get("msg") for log in sim.event_log]: | |
| sim.event_log.append({"type": "chaos", "msg": w}) | |
| await manager.broadcast({"type": "notification", "msg": w, "level": "warning"}) | |
| sim.warnings.clear() | |
| if sim.running and not sim.done and sim.obs_all and sim.runner: | |
| actions = {} | |
| for agent_id, agent in sim.runner.agents.items(): | |
| if agent_id in sim.obs_all: | |
| ao = sim.obs_all[agent_id] | |
| raw = agent.act(ao["prompt"], ao["observation"]) | |
| actions[agent_id] = raw | |
| # Capture thought bubble for LLM agents | |
| if isinstance(agent, GroqLLMAgent): | |
| import re as _re | |
| match = _re.search(r'\{[^{}]*\}', raw, _re.DOTALL) | |
| action_json = match.group() if match else raw | |
| reasoning = raw[:match.start()].strip() if match else "Processing..." | |
| sim.agent_thoughts[agent_id] = { | |
| "action": action_json, | |
| "reasoning": reasoning or "Analyzing sonar data...", | |
| "is_llm": True, | |
| } | |
| else: | |
| # Parse mock agent action for display | |
| import re as _re | |
| match = _re.search(r'"intent"\s*:\s*"([^"]+)"', raw) | |
| intent = match.group(1) if match else "scan" | |
| sim.agent_thoughts[agent_id] = { | |
| "action": intent.upper(), | |
| "reasoning": "", | |
| "is_llm": False, | |
| } | |
| obs_all, rewards, done, info = sim.env.step(actions) | |
| sim.obs_all = obs_all | |
| sim.done = done | |
| step_reward = rewards.get("__total__", 0.0) | |
| sim.total_reward += step_reward | |
| sim.reward_history.append(round(step_reward, 2)) | |
| sim.bio_history.append(round(info["biodiversity"], 1)) | |
| sim.net_history.append(info["active_nets"]) | |
| events = [] | |
| for c in info.get("cleaned", []): | |
| events.append({"type": "clean", "msg": f"CLEANED: {c}"}) | |
| for te in info.get("treaty_events", []): | |
| events.append({"type": "treaty", "msg": f"TREATY: {te}"}) | |
| for b in info.get("broadcasts", []): | |
| events.append({"type": "broadcast", "msg": f"BROADCAST: {b}"}) | |
| if sim.cfg["chaos_enabled"] and sim.adversary.should_inject(info["step"]): | |
| for c in sim.adversary.inject(sim.env.state): | |
| events.append({"type": "chaos", "msg": f"CHAOS: {c}"}) | |
| events.append({"type": "step", "msg": f"[{info['step']:3d}] Nets:{info['active_nets']} | Bio:{info['biodiversity']:.1f}% | R:{step_reward:+.1f}"}) | |
| sim.event_log.extend(events) | |
| if done: | |
| sim.running = False | |
| summary = sim.env.get_episode_summary() | |
| sim.event_log.append({"type": "done", "msg": f"EPISODE COMPLETE | Bio:{summary['biodiversity_final']:.1f}% | Cleaned:{summary['total_cleaned']} | Treaty:{summary['treaty_status']}"}) | |
| await manager.broadcast(sim.get_frame()) | |
| await asyncio.sleep(sim.cfg["step_delay"]) | |
| else: | |
| await asyncio.sleep(0.1) | |
| async def startup(): | |
| asyncio.create_task(simulation_loop()) | |
| # ── WebSocket endpoint ───────────────────────────────────────────────────────── | |
| async def websocket_endpoint(websocket: WebSocket): | |
| await manager.connect(websocket) | |
| await websocket.send_text(json.dumps(sim.get_frame())) | |
| try: | |
| while True: | |
| data = await websocket.receive_text() | |
| msg = json.loads(data) | |
| cmd = msg.get("cmd") | |
| if cmd == "start": | |
| difficulty = msg.get("difficulty", "medium") | |
| chaos_map = {"easy": 40, "medium": 25, "hard": 12} | |
| chaos_interval = chaos_map.get(difficulty, 25) | |
| agent_mode = msg.get("agent_mode", "mock") | |
| groq_key = msg.get("groq_key", "") | |
| # Fallback to HF Space Secret if available | |
| env_key = os.getenv("GROQ_API_KEY") | |
| if not groq_key and env_key: | |
| groq_key = env_key | |
| agent_mode = "llm" | |
| sim.cfg["chaos_enabled"] = True | |
| sim.cfg["step_delay"] = msg.get("step_delay", 0.08) | |
| sim.cfg["difficulty"] = difficulty | |
| sim.reset( | |
| seed=msg.get("seed", 42), | |
| max_steps=msg.get("max_steps", 120), | |
| chaos_interval=chaos_interval, | |
| agent_mode=agent_mode, | |
| groq_key=groq_key, | |
| ) | |
| sim.running = True | |
| await manager.broadcast({"type": "control", "status": "started", "agent_mode": agent_mode}) | |
| elif cmd == "stop": | |
| sim.running = False | |
| await manager.broadcast({"type": "control", "status": "stopped"}) | |
| elif cmd == "reset": | |
| sim.reset() | |
| await manager.broadcast(sim.get_frame()) | |
| elif cmd == "get_state": | |
| await websocket.send_text(json.dumps(sim.get_frame())) | |
| elif cmd == "get_replay": | |
| mode = msg.get("mode", "trained") | |
| base_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) | |
| path = os.path.join(base_dir, "data", f"{mode}_episode.json") | |
| if os.path.exists(path): | |
| with open(path) as f: | |
| replay = json.load(f) | |
| await websocket.send_text(json.dumps({"type": "replay_data", "data": replay})) | |
| else: | |
| await websocket.send_text(json.dumps({"type": "error", "msg": f"No replay data for mode '{mode}'."})) | |
| except WebSocketDisconnect: | |
| manager.disconnect(websocket) | |
| # ── REST endpoints ───────────────────────────────────────────────────────────── | |
| def api_state(): | |
| return sim.get_frame() | |
| def api_replay(mode: str): | |
| base_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) | |
| path = os.path.join(base_dir, "data", f"{mode}_episode.json") | |
| if os.path.exists(path): | |
| with open(path) as f: | |
| return json.load(f) | |
| return {"error": f"No replay data for '{mode}'"} | |
| def health(): | |
| return {"status": "ok", "environment": "Oceanus-v3", "connected_clients": len(manager.active), "agent_mode": sim.agent_mode} | |
| async def serve_ui(): | |
| html_path = Path(__file__).parent / "index.html" | |
| if html_path.exists(): | |
| return HTMLResponse(html_path.read_text(encoding="utf-8")) | |
| return HTMLResponse("<h1>index.html not found.</h1>") | |