| """The server: OpenEnv's app plus read-only JSON about the tasks on the same port. |
| |
| OpenEnv provides `/reset`, `/step`, `/state`, `/schema`, `/ws`, `/mcp`, `/health`, the Task API and its web UI at |
| `/web`. We add: |
| |
| /api/tasks task summaries with the references' costs |
| /api/tasks/{id} a public task (what the agent sees, structured) |
| /api/tasks/{id}/reference the optimum and the baselines - for people; no tool shows them |
| /api/episodes live episodes on this server (newest first) |
| /api/tasks/{id}/plans the plans the replay can show: optimum, baselines, the Corps, live policies, and |
| every model's plan in the committed runs (results/rollouts, or FLOOD_RUNS_DIR) |
| /api/tasks/{id}/replay?plan=... a plan on the true river: releases, storage, gauge flows, flood categories, cost |
| /twin/ the 3D replay of the basin (USGS terrain and hydrography) for any plan |
| |
| uvicorn flood_openenv.server:app --port 8000 |
| """ |
|
|
| from __future__ import annotations |
|
|
| import os |
| from pathlib import Path |
|
|
| from fastapi import HTTPException |
| from fastapi.responses import RedirectResponse |
| from fastapi.staticfiles import StaticFiles |
| from openenv.core.env_server.http_server import create_app |
| from openenv.core.env_server.mcp_types import CallToolAction, CallToolObservation |
|
|
| from flood_core import TaskPack, load_pack |
|
|
| try: |
| from .environment import STORE, FloodOpsEnvironment, FloodState |
| except ImportError: |
| from environment import STORE, FloodOpsEnvironment, FloodState |
|
|
| ENV_NAME = "flood_ops" |
| WEB = Path(__file__).resolve().parent / "web" |
| _UP = Path(__file__).resolve().parents |
| RUNS_DIR = Path(os.environ.get("FLOOD_RUNS_DIR") or (_UP[4] if len(_UP) > 4 else _UP[-1]) / "results" / "rollouts") |
|
|
|
|
| def mount_api(app, pack: TaskPack): |
| @app.get("/api/tasks") |
| def tasks(): |
| return [{"task_id": t["task_id"], "split": t["split"], "tier": t["tier"], "start_local": t["start_local"], |
| "steps": t["steps"], **{k: v for k, v in t["reference"].items() if k.endswith("_cost")}} |
| for t in pack.tasks] |
|
|
| def _get(tid: str) -> dict: |
| try: |
| return pack.get(tid) |
| except KeyError: |
| raise HTTPException(404, f"no task {tid}") from None |
|
|
| @app.get("/api/tasks/{tid}") |
| def task(tid: str): |
| return pack.public(_get(tid)) |
|
|
| @app.get("/api/tasks/{tid}/reference") |
| def reference(tid: str): |
| return _get(tid)["reference"] |
|
|
| @app.get("/api/episodes") |
| def episodes(n: int = 50): |
| return [e.summary() for e in STORE.recent(n)] |
|
|
| from flood_core.replay import plans as _plans |
| from flood_core.replay import replay as _replay |
|
|
| @app.get("/api/tasks/{tid}/plans") |
| def task_plans(tid: str): |
| return _plans(_get(tid), RUNS_DIR) |
|
|
| @app.get("/api/tasks/{tid}/replay") |
| def task_replay(tid: str, plan: str = "optimum"): |
| try: |
| return _replay(_get(tid), plan, RUNS_DIR) |
| except KeyError as e: |
| raise HTTPException(404, str(e)) from None |
|
|
| @app.get("/twin", include_in_schema=False) |
| def twin_slash(): |
| return RedirectResponse("/twin/") |
|
|
| app.mount("/twin/", StaticFiles(directory=WEB / "twin", html=True), name="twin") |
|
|
| @app.middleware("http") |
| async def no_stale_pages(request, call_next): |
| |
| resp = await call_next(request) |
| if request.url.path.startswith("/twin"): |
| resp.headers["Cache-Control"] = "no-cache" |
| return resp |
|
|
|
|
| def build_app(pack: TaskPack | None = None, mode: str | None = None): |
| """FLOOD_MODE=live serves live mode (forecast scenarios, block commits); the default is plan mode.""" |
| pack = pack or load_pack() |
| mode = mode or os.environ.get("FLOOD_MODE", "plan") |
| if mode == "live": |
| try: |
| from .live_env import FloodLiveEnvironment as Env |
| except ImportError: |
| from live_env import FloodLiveEnvironment as Env |
| else: |
| Env = FloodOpsEnvironment |
|
|
| def factory(): |
| return Env(pack) |
|
|
| kw = {} |
| if os.environ.get("ENABLE_WEB_INTERFACE", "").lower() == "true": |
| try: |
| from .gradio_ui import build_ui |
| except ImportError: |
| from gradio_ui import build_ui |
| kw = dict(gradio_builder=build_ui, custom_tab_name="Explore", custom_tab_primary=True, show_default_tab=True, |
| title_override=f"FloodGate v1 ({mode} mode)") |
| app = create_app(factory, CallToolAction, CallToolObservation, env_name=ENV_NAME, |
| max_concurrent_envs=int(os.environ.get("MAX_CONCURRENT_ENVS", "64")), state_cls=FloodState, **kw) |
| mount_api(app, pack) |
|
|
| @app.get("/healthz", include_in_schema=False) |
| def healthz(): |
| return {"status": "ok", "mode": mode, "tasks": {s: pack.count(s) for s in pack.splits()}} |
|
|
| if not kw: |
| @app.get("/", include_in_schema=False) |
| def root(): |
| return RedirectResponse("/twin/") |
|
|
| return app |
|
|
|
|
| _app = None |
|
|
|
|
| def __getattr__(name: str): |
| global _app |
| if name == "app": |
| if _app is None: |
| _app = build_app() |
| return _app |
| raise AttributeError(name) |
|
|
