"""
Narada: FastAPI server (WebSocket primary, HTTP for debug).
Port is read from os.environ["PORT"], default 7860 for HF Spaces.
WORKERS=1 is mandatory — sessions are in-memory per WebSocket connection.
"""
from __future__ import annotations
import asyncio
import json
import logging
import os
from contextlib import asynccontextmanager
from typing import Any, Dict, Optional
import uvicorn
from fastapi import FastAPI, HTTPException, Request, WebSocket, WebSocketDisconnect
from fastapi.responses import HTMLResponse, JSONResponse
from ..graph import get_graph
from ..models import NaradaAction, NaradaObservation, NaradaState, StepResult
from .environment import NaradaEnvironment
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
@asynccontextmanager
async def lifespan(app: FastAPI):
logger.info("Narada starting — loading graph ...")
graph = get_graph() # blocks until loaded; ~5s on cold start
logger.info(
"Graph ready: %d nodes loaded",
len(graph.nodes),
)
yield
logger.info("Narada shutting down.")
app = FastAPI(
title="Narada",
description=(
"OpenEnv: LLM agent navigates a gene-disease knowledge graph to diagnose "
"rare diseases. Three task tiers: monogenic (easy), oligogenic (medium), "
"phenotype_mismatch (hard)."
),
version="1.0.0",
lifespan=lifespan,
)
# ── Health ────────────────────────────────────────────────────────────────────
@app.get("/health")
async def health() -> Dict[str, str]:
return {"status": "healthy", "version": "1.0.0", "environment": "narada"}
@app.get("/graph/subgraph")
async def subgraph(node_id: str, depth: int = 2, max_nodes: int = 60) -> Dict[str, Any]:
"""Return a JSON subgraph centred on node_id for D3.js force-graph rendering."""
graph = get_graph()
visited: dict = {}
queue = [(node_id, 0)]
while queue and len(visited) < max_nodes:
nid, d = queue.pop(0)
if nid in visited or d > depth:
continue
nd = graph.get_node(nid)
if nd is None:
continue
visited[nid] = nd
if d < depth:
for nb in graph.get_neighbors(nid)[:12]:
if nb not in visited:
queue.append((nb, d + 1))
nodes = [
{"id": nid, "label": nd.get("name", nid)[:30], "type": nd.get("type", "unknown")}
for nid, nd in visited.items()
]
node_set = {n["id"] for n in nodes}
links = [
{"source": nid, "target": nb}
for nid in visited
for nb in graph.get_neighbors(nid)
if nb in node_set and nb != nid
]
return {"nodes": nodes, "links": links, "center": node_id}
# ── OpenEnv standard endpoints ────────────────────────────────────────────────
@app.get("/metadata")
async def metadata() -> Dict[str, Any]:
return {
"name": "narada",
"version": "1.0.0",
"description": (
"LLM agent navigates a 55,000-node gene-disease knowledge graph "
"(ClinVar + HPO) to diagnose rare disease patients. Three task tiers: "
"monogenic (easy), oligogenic (medium), phenotype_mismatch (hard)."
),
"tasks": ["monogenic", "oligogenic", "phenotype_mismatch"],
"reward_range": [0.01, 0.99],
}
@app.get("/schema")
async def schema() -> Dict[str, Any]:
return {
"action": NaradaAction.model_json_schema(),
"observation": NaradaObservation.model_json_schema(),
"state": NaradaState.model_json_schema(),
}
@app.post("/mcp")
async def mcp(request: Request) -> JSONResponse:
"""Minimal MCP-style JSON-RPC bridge over the shared HTTP debug env.
Advertised tools proxy to /reset, /step, /state using the same in-memory
environment. Primary transport is still the WebSocket endpoint; this exists
so MCP clients can poke the env without needing a WS session.
"""
try:
body = await request.json()
except Exception:
body = {}
method = body.get("method", "")
req_id = body.get("id", 1)
tools = [
{
"name": "narada_reset_episode",
"description": "Reset the environment and start a new episode.",
"inputSchema": {
"type": "object",
"properties": {
"task_type": {
"type": "string",
"enum": ["monogenic", "oligogenic", "phenotype_mismatch"],
},
"seed": {"type": "integer"},
},
},
},
{
"name": "narada_step_action",
"description": "Take an action in the environment.",
"inputSchema": NaradaAction.model_json_schema(),
},
{
"name": "narada_get_state",
"description": "Return the current episode state (ground truth only when done).",
"inputSchema": {"type": "object"},
},
]
async def _tool_result(payload: Any) -> Dict[str, Any]:
return {"content": [{"type": "text", "text": json.dumps(payload, default=str)}]}
if method == "tools/list":
result: Dict[str, Any] = {"tools": tools}
elif method == "tools/call":
params = body.get("params") or {}
name = params.get("name", "")
args = params.get("arguments", {}) or {}
async with _http_lock:
try:
if name == "narada_reset_episode":
step = _http_env.reset(
task_type=args.get("task_type") or "monogenic",
seed=args.get("seed"),
)
result = await _tool_result(step.model_dump())
elif name == "narada_step_action":
action = NaradaAction.model_validate(args)
step = _http_env.step(action)
result = await _tool_result(step.model_dump())
elif name == "narada_get_state":
state = _http_env.state()
result = await _tool_result(state.model_dump())
else:
return JSONResponse(
{
"jsonrpc": "2.0",
"error": {"code": -32601, "message": f"Unknown tool: {name}"},
"id": req_id,
}
)
except Exception as e:
return JSONResponse(
{
"jsonrpc": "2.0",
"error": {"code": -32000, "message": str(e)},
"id": req_id,
}
)
else:
result = {"name": "narada", "version": "1.0.0", "tools": [t["name"] for t in tools]}
return JSONResponse({"jsonrpc": "2.0", "result": result, "id": req_id})
# ── WebSocket (primary transport) ─────────────────────────────────────────────
@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket) -> None:
"""
Persistent WebSocket session.
Protocol:
Client → Server:
{"type": "reset", "task_type": "monogenic|oligogenic|phenotype_mismatch", "seed": int|null}
{"type": "step", "action": {...}}
{"type": "state"}
Server → Client:
{"type": "observation", "data": StepResult}
{"type": "state", "data": NaradaState}
{"type": "error", "message": "..."}
"""
await websocket.accept()
env = NaradaEnvironment()
logger.info("WS session opened")
try:
while True:
raw = await websocket.receive_text()
try:
msg = json.loads(raw)
except json.JSONDecodeError as e:
await websocket.send_text(json.dumps({"type": "error", "message": f"Bad JSON: {e}"}))
continue
mtype = msg.get("type")
if mtype == "reset":
try:
result = env.reset(
task_type=msg.get("task_type", "monogenic"),
seed=msg.get("seed"),
)
await websocket.send_text(
json.dumps({"type": "observation", "data": result.model_dump()})
)
except Exception as e:
logger.exception("reset error")
await websocket.send_text(json.dumps({"type": "error", "message": str(e)}))
elif mtype == "step":
try:
action = NaradaAction.model_validate(msg.get("action", {}))
result = env.step(action)
await websocket.send_text(
json.dumps({"type": "observation", "data": result.model_dump()})
)
except Exception as e:
logger.exception("step error")
await websocket.send_text(json.dumps({"type": "error", "message": str(e)}))
elif mtype == "state":
try:
state = env.state()
await websocket.send_text(
json.dumps({"type": "state", "data": state.model_dump()})
)
except Exception as e:
logger.exception("state error")
await websocket.send_text(json.dumps({"type": "error", "message": str(e)}))
else:
await websocket.send_text(
json.dumps({"type": "error", "message": f"Unknown type: {mtype}"})
)
except WebSocketDisconnect:
logger.info("WS session closed")
except Exception as e:
logger.exception("WS unexpected error: %s", e)
try:
await websocket.send_text(json.dumps({"type": "error", "message": str(e)}))
except Exception:
pass
# ── HTTP debug endpoints ──────────────────────────────────────────────────────
_http_env = NaradaEnvironment()
_http_lock = asyncio.Lock()
@app.post("/reset", response_model=StepResult)
async def http_reset(task_type: Optional[str] = None, seed: Optional[int] = None) -> StepResult:
async with _http_lock:
return _http_env.reset(task_type=task_type or "monogenic", seed=seed)
@app.post("/step", response_model=StepResult)
async def http_step(action: NaradaAction) -> StepResult:
async with _http_lock:
try:
return _http_env.step(action)
except RuntimeError as e:
raise HTTPException(status_code=400, detail=str(e))
@app.get("/state", response_model=NaradaState)
async def http_state() -> NaradaState:
async with _http_lock:
return _http_env.state()
# ── Web UI ────────────────────────────────────────────────────────────────────
_WEB_UI = """
Narada | OpenEnv
Narada — Rare Disease Diagnosis Agent
OpenEnv | Navigate the gene-disease knowledge graph to find the causal variant
STEP
—
REWARD
—
TASK
—
STATUS
IDLE
Patient phenotypes will appear here after RESET.
[NARADA] Ready. Select a task and press RESET.
"""
@app.get("/", response_class=HTMLResponse)
async def root() -> HTMLResponse:
return HTMLResponse(content=_WEB_UI)
@app.get("/web", response_class=HTMLResponse)
async def web_ui() -> HTMLResponse:
return HTMLResponse(content=_WEB_UI)
def main() -> None:
port = int(os.environ.get("PORT", 7860))
host = os.environ.get("HOST", "0.0.0.0")
# Multiple workers would duplicate the in-memory graph and split sessions
# across processes, so the shared `_http_env` and WebSocket sessions would
# see different episodes. Pin to 1 worker and warn if the env tries to
# override it.
workers_env = int(os.environ.get("WORKERS", 1))
if workers_env != 1:
logger.warning(
"WORKERS=%d ignored; Narada requires a single worker for session "
"and graph consistency.",
workers_env,
)
uvicorn.run("narada.server.app:app", host=host, port=port, workers=1)
# Keep run() as an alias so existing callers still work
run = main
if __name__ == "__main__":
main()