Spaces:
Sleeping
Sleeping
Download api_server.py from aadi-joshi/omninpc: direct link, hf CLI and curl.
- Browser
- Download file 17.4 kB
-
https://huggingface.co/spaces/aadi-joshi/omninpc/resolve/main/api_server.py
- Command line
-
hf download hf://spaces/aadi-joshi/omninpc/api_server.py
-
curl -L -o api_server.py https://huggingface.co/spaces/aadi-joshi/omninpc/resolve/main/api_server.py
17.4 kB
| """OmniNPC FastAPI Server โ the backend for the HTML / Three.js demo. | |
| Routes: | |
| POST /infer โ JPEG-base64 frame โ Nano action JSON (perception) | |
| POST /infer_with_thinking โ same as /infer but returns the raw thinking trace | |
| POST /super โ request a fresh strategic directive (Super agent) | |
| POST /dialogue โ when two NPCs are near each other, get a dialogue line | |
| GET /state โ current global state snapshot (for the HUD) | |
| GET /metrics โ latency / call counts | |
| GET /health โ health check | |
| GET /config โ frontend-safe config (model names, intervals) | |
| CORS is open so the static HTML page can call this from any origin (file:// included). | |
| On Hugging Face Spaces this server runs alongside the static site. | |
| Requires NVIDIA_NIM_API_KEY in the environment. | |
| """ | |
| from __future__ import annotations | |
| import argparse | |
| import asyncio | |
| import base64 | |
| import csv | |
| import logging | |
| import os | |
| import threading | |
| import time | |
| from collections import Counter, deque | |
| from datetime import datetime | |
| from pathlib import Path | |
| from typing import Optional | |
| from dotenv import load_dotenv | |
| load_dotenv() | |
| logging.basicConfig( | |
| level=logging.INFO, | |
| format="%(asctime)s [%(levelname)s] %(name)s: %(message)s", | |
| ) | |
| logger = logging.getLogger("api_server") | |
| try: | |
| from fastapi import FastAPI, HTTPException | |
| from fastapi.middleware.cors import CORSMiddleware | |
| from fastapi.responses import JSONResponse, FileResponse | |
| from fastapi.staticfiles import StaticFiles | |
| from pydantic import BaseModel | |
| except ImportError: | |
| logger.error("FastAPI not installed. Run: pip install fastapi uvicorn") | |
| raise | |
| from config import load_config | |
| from inference.nim_client import NIMClient | |
| from inference.prompt_templates import ( | |
| build_nano_system_prompt, | |
| build_nano_user_message, | |
| SUPER_SYSTEM_PROMPT, | |
| build_super_user_message, | |
| ) | |
| # โโโ Request / Response models โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ | |
| class InferRequest(BaseModel): | |
| image_b64: str # raw base64 (no data: prefix) | |
| audio_text: str = "" # transcribed voice command (or "") | |
| npc_role: str = "hostile" # "hostile" | "ally" | "wanderer" | |
| directive: str = "" # current Super directive (frontend-cached) | |
| class SuperRequest(BaseModel): | |
| scene_summary: str | |
| mission_context: str = "Compete with the player and other NPCs to collect coins." | |
| recent_actions: list[str] = [] | |
| class DialogueRequest(BaseModel): | |
| npc_a: str # role of NPC A | |
| npc_b: str # role of NPC B | |
| scene_summary: str = "" | |
| last_event: str = "" # e.g., "ally collected a coin" | |
| class CommandRequest(BaseModel): | |
| text: str | |
| target_npc: str = "hostile" | |
| # โโโ Session logger โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ | |
| class _SessionLogger: | |
| """Lightweight CSV logger - one row per /infer call.""" | |
| _COLUMNS = ["timestamp", "npc_role", "action", "direction", | |
| "confidence", "total_latency_ms", "say"] | |
| def __init__(self, log_dir: str) -> None: | |
| self._lock = threading.Lock() | |
| Path(log_dir).mkdir(parents=True, exist_ok=True) | |
| ts = datetime.now().strftime("%Y%m%d_%H%M%S") | |
| self._path = Path(log_dir) / f"session_{ts}.csv" | |
| with open(self._path, "w", newline="", encoding="utf-8") as f: | |
| csv.DictWriter(f, self._COLUMNS).writeheader() | |
| logger.info("Session log: %s", self._path) | |
| def log(self, npc_role: str, action: str, direction: str, | |
| confidence: float, latency_ms: float, say: str) -> None: | |
| row = { | |
| "timestamp": datetime.now().isoformat(), | |
| "npc_role": npc_role, | |
| "action": action, | |
| "direction": direction, | |
| "confidence": f"{confidence:.3f}", | |
| "total_latency_ms": f"{latency_ms:.1f}", | |
| "say": (say or "")[:50], | |
| } | |
| with self._lock: | |
| with open(self._path, "a", newline="", encoding="utf-8") as f: | |
| csv.DictWriter(f, self._COLUMNS).writerow(row) | |
| # โโโ App state โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ | |
| class ServerState: | |
| def __init__(self) -> None: | |
| self.config = load_config() | |
| self.api_key = os.getenv("NVIDIA_NIM_API_KEY", "") | |
| self.nim = NIMClient(api_key=self.api_key, base_url=self.config.nvidia.base_url) | |
| self.last_directive = "Compete with the player to collect the most coins." | |
| self.start_time = time.time() | |
| self.total_infer_calls = 0 | |
| self.total_super_calls = 0 | |
| self.total_dialogue_calls = 0 | |
| self.recent_latencies: deque[float] = deque(maxlen=50) | |
| self.recent_actions: deque[dict] = deque(maxlen=50) | |
| self.recent_thinks: deque[str] = deque(maxlen=20) | |
| self.action_histogram: Counter = Counter() | |
| self.session_log = _SessionLogger(self.config.metrics.log_dir) | |
| state = ServerState() | |
| # โโโ App + CORS โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ | |
| app = FastAPI( | |
| title="OmniNPC API", | |
| description="Hierarchical Nemotron Super+Nano backend for the OmniNPC HTML demo.", | |
| version="2.0.0", | |
| ) | |
| app.add_middleware( | |
| CORSMiddleware, | |
| allow_origins=["*"], | |
| allow_credentials=False, | |
| allow_methods=["*"], | |
| allow_headers=["*"], | |
| ) | |
| # โโโ Helpers โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ | |
| import re | |
| import json | |
| def _strip_thinking(raw: str) -> tuple[str, str]: | |
| """Return (think_block, cleaned_response).""" | |
| think_match = re.search(r"<think>(.*?)</think>", raw, flags=re.DOTALL) | |
| think = think_match.group(1).strip() if think_match else "" | |
| cleaned = re.sub(r"<think>.*?</think>", "", raw, flags=re.DOTALL).strip() | |
| cleaned = re.sub(r"```[a-z]*", "", cleaned).replace("```", "").strip() | |
| return think, cleaned | |
| def _extract_json(cleaned: str) -> dict: | |
| """Pull the first {...} JSON object out of a model response.""" | |
| match = re.search(r"\{.*\}", cleaned, re.DOTALL) | |
| if not match: | |
| return {} | |
| try: | |
| return json.loads(match.group(0)) | |
| except json.JSONDecodeError: | |
| return {} | |
| def _role_persona(role: str) -> str: | |
| """Return persona text injected into Nano's system prompt for each NPC role.""" | |
| role = role.lower() | |
| if role == "hostile": | |
| return ("You are the HOSTILE NPC. Your goal: collect coins faster than the player. " | |
| "If the player gets within 3m, evade or sprint away. Be competitive.") | |
| if role == "ally": | |
| return ("You are the ALLY NPC. You team up with the player. Stay near the player " | |
| "(within 4-6m), help collect coins, and warn about the hostile NPC.") | |
| if role == "wanderer": | |
| return ("You are the WANDERER NPC. You do your own thing โ explore, collect coins, " | |
| "occasionally dance. Friendly but independent.") | |
| return "" | |
| # โโโ Routes โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ | |
| async def health(): | |
| return { | |
| "status": "ok", | |
| "version": "2.0.0", | |
| "uptime_s": round(time.time() - state.start_time, 1), | |
| "api_key_set": bool(state.api_key), | |
| "nano_model": state.config.nvidia.nano_omni_model, | |
| "super_model": state.config.nvidia.super_model, | |
| } | |
| async def get_config(): | |
| return { | |
| "nano_model": state.config.nvidia.nano_omni_model, | |
| "super_model": state.config.nvidia.super_model, | |
| "perception_interval_ms": int(state.config.orchestrator.perception_interval * 1000), | |
| "strategy_interval_ms": int(state.config.orchestrator.strategy_interval * 1000), | |
| } | |
| async def metrics(): | |
| lats = list(state.recent_latencies) | |
| avg = sum(lats) / len(lats) if lats else 0.0 | |
| p95 = sorted(lats)[int(len(lats) * 0.95)] if len(lats) >= 5 else (max(lats) if lats else 0.0) | |
| uptime_s = time.time() - state.start_time | |
| rate_hz = round(state.total_infer_calls / uptime_s, 3) if uptime_s > 0 else 0.0 | |
| return { | |
| "total_infer_calls": state.total_infer_calls, | |
| "total_super_calls": state.total_super_calls, | |
| "total_dialogue_calls": state.total_dialogue_calls, | |
| "avg_latency_ms": round(avg, 1), | |
| "p95_latency_ms": round(p95, 1), | |
| "recent_latencies_ms": [round(x, 1) for x in lats[-20:]], | |
| "current_directive": state.last_directive, | |
| "inference_rate_hz": rate_hz, | |
| "total_input_tokens": state.nim.total_input_tokens, | |
| "total_output_tokens": state.nim.total_output_tokens, | |
| "total_nim_errors": state.nim.total_errors, | |
| "action_histogram": dict(state.action_histogram.most_common()), | |
| } | |
| async def state_snapshot(): | |
| return { | |
| "directive": state.last_directive, | |
| "recent_actions": list(state.recent_actions)[-10:], | |
| "recent_thinks": list(state.recent_thinks)[-5:], | |
| } | |
| async def infer(req: InferRequest): | |
| """Run Nano perception on a single frame. Returns action JSON.""" | |
| if not state.api_key: | |
| raise HTTPException(500, "NVIDIA_NIM_API_KEY not configured on server") | |
| if not req.image_b64: | |
| raise HTTPException(400, "image_b64 missing") | |
| persona = _role_persona(req.npc_role) | |
| directive = req.directive or state.last_directive | |
| full_directive = f"{directive} {persona}".strip() | |
| system_prompt = build_nano_system_prompt(full_directive) | |
| audio_ctx = ( | |
| f"Voice command heard: '{req.audio_text}'. Respond accordingly." | |
| if req.audio_text else "No voice command. React to visual input only." | |
| ) | |
| user_text = build_nano_user_message(audio_ctx) | |
| content = [ | |
| {"type": "text", "text": user_text}, | |
| {"type": "image_url", | |
| "image_url": {"url": f"data:image/jpeg;base64,{req.image_b64}"}}, | |
| ] | |
| messages = [ | |
| {"role": "system", "content": system_prompt}, | |
| {"role": "user", "content": content}, | |
| ] | |
| t0 = time.perf_counter() | |
| raw = await state.nim.complete( | |
| model=state.config.nvidia.nano_omni_model, | |
| messages=messages, | |
| max_tokens=state.config.nvidia.max_tokens, | |
| temperature=0.1, | |
| ) | |
| latency_ms = (time.perf_counter() - t0) * 1000 | |
| state.recent_latencies.append(latency_ms) | |
| state.total_infer_calls += 1 | |
| if not raw: | |
| err = state.nim.last_error or "unknown error" | |
| return {"action": "idle", "magnitude": 0, "confidence": 0, | |
| "say": "", | |
| "scene_description": f"NIM error: {err}", | |
| "thinking": "", | |
| "error": err, | |
| "latency_ms": round(latency_ms, 1), | |
| "npc_role": req.npc_role} | |
| think, cleaned = _strip_thinking(raw) | |
| data = _extract_json(cleaned) | |
| response = { | |
| "action": data.get("action", "idle"), | |
| "target": data.get("target", ""), | |
| "direction": data.get("direction", "none"), | |
| "magnitude": float(data.get("magnitude", 0.5)), | |
| "confidence": float(data.get("confidence", 0.5)), | |
| "say": data.get("say", ""), | |
| "scene_description": data.get("scene_description", ""), | |
| "thinking": think[:500], | |
| "latency_ms": round(latency_ms, 1), | |
| "npc_role": req.npc_role, | |
| "ts": time.time(), | |
| } | |
| state.action_histogram[response["action"]] += 1 | |
| state.recent_actions.append({k: response[k] for k in ("action", "say", "npc_role", "ts")}) | |
| if think: | |
| state.recent_thinks.append(think[:200]) | |
| state.session_log.log( | |
| npc_role=req.npc_role, | |
| action=response["action"], | |
| direction=response["direction"], | |
| confidence=response["confidence"], | |
| latency_ms=latency_ms, | |
| say=response["say"], | |
| ) | |
| return response | |
| async def super_directive(req: SuperRequest): | |
| """Request a fresh strategic directive from the Super model.""" | |
| if not state.api_key: | |
| raise HTTPException(500, "NVIDIA_NIM_API_KEY not configured on server") | |
| user_msg = build_super_user_message( | |
| req.scene_summary, req.mission_context, req.recent_actions | |
| ) | |
| messages = [ | |
| {"role": "system", "content": SUPER_SYSTEM_PROMPT}, | |
| {"role": "user", "content": user_msg}, | |
| ] | |
| t0 = time.perf_counter() | |
| raw = await state.nim.complete( | |
| model=state.config.nvidia.super_model, | |
| messages=messages, | |
| max_tokens=2048, | |
| temperature=0.3, | |
| ) | |
| latency_ms = (time.perf_counter() - t0) * 1000 | |
| state.total_super_calls += 1 | |
| think, cleaned = _strip_thinking(raw) | |
| data = _extract_json(cleaned) | |
| directive = data.get("directive", "Compete to collect coins.") | |
| state.last_directive = directive | |
| if think: | |
| state.recent_thinks.append(f"[SUPER] {think[:400]}") | |
| return { | |
| "directive": directive, | |
| "priority": data.get("priority", "explore"), | |
| "reasoning": data.get("reasoning", ""), | |
| "thinking": think[:800], | |
| "latency_ms": round(latency_ms, 1), | |
| } | |
| async def dialogue(req: DialogueRequest): | |
| """Generate a single line of dialogue between two NPCs.""" | |
| if not state.api_key: | |
| raise HTTPException(500, "NVIDIA_NIM_API_KEY not configured on server") | |
| sys_prompt = ( | |
| "You write one short line of dialogue spoken between two NPC characters " | |
| "in a coin-collecting game. Output ONLY a JSON object: " | |
| '{"speaker":"<a|b>","text":"<one line, max 12 words>"}.' | |
| ) | |
| user_prompt = ( | |
| f"NPC A is the {req.npc_a}. NPC B is the {req.npc_b}. " | |
| f"Scene: {req.scene_summary or 'they meet near the coins'}. " | |
| f"Recent event: {req.last_event or 'none'}. Output the JSON now." | |
| ) | |
| messages = [ | |
| {"role": "system", "content": sys_prompt}, | |
| {"role": "user", "content": user_prompt}, | |
| ] | |
| t0 = time.perf_counter() | |
| raw = await state.nim.complete( | |
| model=state.config.nvidia.super_model, | |
| messages=messages, | |
| max_tokens=120, | |
| temperature=0.7, | |
| ) | |
| latency_ms = (time.perf_counter() - t0) * 1000 | |
| state.total_dialogue_calls += 1 | |
| _, cleaned = _strip_thinking(raw) | |
| data = _extract_json(cleaned) | |
| return { | |
| "speaker": data.get("speaker", "a"), | |
| "text": data.get("text", "..."), | |
| "latency_ms": round(latency_ms, 1), | |
| } | |
| # โโโ Startup: pre-warm NIM connection โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ | |
| async def _prewarm_nim(): | |
| """Establish TLS + TCP to NIM before first user request (eliminates cold-start latency).""" | |
| if not state.api_key: | |
| logger.warning("Skipping NIM pre-warm โ no API key.") | |
| return | |
| logger.info("Pre-warming NIM connection...") | |
| try: | |
| await state.nim.complete( | |
| model=state.config.nvidia.nano_omni_model, | |
| messages=[{"role": "user", "content": "ping"}], | |
| max_tokens=5, | |
| temperature=0.1, | |
| ) | |
| logger.info("NIM pre-warm complete โ first real request will be fast.") | |
| except Exception as exc: | |
| logger.warning("NIM pre-warm failed (non-fatal): %s", exc) | |
| # โโโ Static HTML hosting (so HF Spaces serves both API and game) โโโโโโโโโโโโ | |
| # Mount at "/" so index.html's `./game.js` reference resolves correctly. | |
| # Explicit API routes (above) take precedence over the static catch-all. | |
| _HTML_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "html_game") | |
| if os.path.isdir(_HTML_DIR): | |
| app.mount("/", StaticFiles(directory=_HTML_DIR, html=True), name="game") | |
| # โโโ Entry point โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ | |
| def main() -> None: | |
| parser = argparse.ArgumentParser(description="OmniNPC API Server") | |
| parser.add_argument("--port", type=int, default=int(os.getenv("PORT", "8000"))) | |
| parser.add_argument("--host", type=str, default=os.getenv("HOST", "0.0.0.0")) | |
| args = parser.parse_args() | |
| if not state.api_key: | |
| logger.warning("NVIDIA_NIM_API_KEY not set โ /infer routes will return 500.") | |
| else: | |
| logger.info("API key OK โ Nano=%s Super=%s", | |
| state.config.nvidia.nano_omni_model, | |
| state.config.nvidia.super_model) | |
| import uvicorn | |
| uvicorn.run(app, host=args.host, port=args.port, log_level="info") | |
| if __name__ == "__main__": | |
| main() | |