| import asyncio |
| import hmac |
| import logging |
| from contextlib import asynccontextmanager |
| from typing import Optional |
|
|
| from fastapi import FastAPI, HTTPException, Request, WebSocket, WebSocketDisconnect |
| from fastapi.middleware.cors import CORSMiddleware |
| from fastapi.responses import JSONResponse |
| from pydantic import BaseModel |
|
|
| from app.config import Settings |
| from app.engine import SignMotionEngine |
| from app.realtime import ( |
| ConnectionLimiter, |
| RealtimeSession, |
| Scheduler, |
| WebRTCManager, |
| WebSocketTransport, |
| ) |
|
|
| logging.basicConfig(level=logging.INFO) |
| logger = logging.getLogger("main") |
|
|
| settings = Settings.from_environment() |
|
|
|
|
| class OfferBody(BaseModel): |
| sdp: str |
| type: str |
|
|
|
|
| async def load_engine_in_background(application: FastAPI) -> None: |
| try: |
| await asyncio.to_thread(application.state.engine.load) |
| logger.info("engine ready") |
| except Exception as error: |
| logger.exception("engine failed to load") |
| application.state.load_error = type(error).__name__ + ": " + str(error) |
|
|
|
|
| @asynccontextmanager |
| async def lifespan(application: FastAPI): |
| engine = SignMotionEngine(settings) |
| scheduler = Scheduler(engine, settings) |
| limiter = ConnectionLimiter(settings.maximum_connections) |
| webrtc = WebRTCManager(settings, engine, scheduler, limiter) |
|
|
| application.state.engine = engine |
| application.state.scheduler = scheduler |
| application.state.limiter = limiter |
| application.state.webrtc = webrtc |
| application.state.load_error = None |
|
|
| await scheduler.start() |
| loader_task = asyncio.create_task(load_engine_in_background(application)) |
|
|
| yield |
|
|
| loader_task.cancel() |
| await webrtc.shutdown() |
| await scheduler.stop() |
|
|
|
|
| app = FastAPI(title="Text To Sign Motion", lifespan=lifespan) |
|
|
| app.add_middleware( |
| CORSMiddleware, |
| allow_origins=settings.allowed_origins, |
| allow_methods=["GET", "POST", "OPTIONS"], |
| allow_headers=["*"], |
| ) |
|
|
|
|
| def token_is_valid(provided_token: str) -> bool: |
| if settings.api_token == "": |
| return True |
| return hmac.compare_digest(provided_token.encode("utf-8"), settings.api_token.encode("utf-8")) |
|
|
|
|
| def extract_request_token(request: Request) -> str: |
| authorization_header = request.headers.get("authorization", "") |
| if authorization_header.lower().startswith("bearer "): |
| return authorization_header[7:].strip() |
| return request.query_params.get("token", "") |
|
|
|
|
| def require_token(request: Request) -> None: |
| if not token_is_valid(extract_request_token(request)): |
| raise HTTPException(status_code=401, detail="Invalid or missing token") |
|
|
|
|
| def origin_is_allowed(origin: Optional[str]) -> bool: |
| if "*" in settings.allowed_origins: |
| return True |
| if origin is None: |
| return True |
| return origin in settings.allowed_origins |
|
|
|
|
| @app.get("/") |
| async def root() -> dict: |
| return { |
| "service": "text-to-sign-motion", |
| "health": "/health", |
| "websocket": "/ws", |
| "webrtc_offer": "/webrtc/offer", |
| } |
|
|
|
|
| @app.get("/health") |
| async def health(request: Request) -> dict: |
| engine = request.app.state.engine |
| load_error = request.app.state.load_error |
|
|
| if engine.ready: |
| status = "ready" |
| elif load_error is not None: |
| status = "error" |
| else: |
| status = "loading" |
|
|
| return { |
| "status": status, |
| "error": load_error, |
| "engine": engine.info() if engine.ready else None, |
| "queue_depth": request.app.state.scheduler.queue.qsize(), |
| "active_connections": request.app.state.limiter.active, |
| } |
|
|
|
|
| @app.get("/ready") |
| async def ready(request: Request) -> JSONResponse: |
| if request.app.state.engine.ready: |
| return JSONResponse(status_code=200, content={"status": "ready"}) |
| return JSONResponse(status_code=503, content={"status": "not ready"}) |
|
|
|
|
| @app.get("/ice") |
| async def ice(request: Request) -> dict: |
| require_token(request) |
| return {"iceServers": settings.ice_servers} |
|
|
|
|
| @app.post("/webrtc/offer") |
| async def webrtc_offer(body: OfferBody, request: Request) -> dict: |
| require_token(request) |
|
|
| if body.type != "offer": |
| raise HTTPException(status_code=400, detail="The type must be offer") |
|
|
| try: |
| answer = await request.app.state.webrtc.handle_offer(body.sdp, body.type) |
| except ConnectionRefusedError: |
| raise HTTPException(status_code=503, detail="Too many active connections") |
|
|
| return answer |
|
|
|
|
| @app.websocket("/ws") |
| async def websocket_endpoint(websocket: WebSocket) -> None: |
| if not origin_is_allowed(websocket.headers.get("origin")): |
| await websocket.close(code=1008) |
| return |
|
|
| if not token_is_valid(websocket.query_params.get("token", "")): |
| await websocket.close(code=1008) |
| return |
|
|
| limiter = websocket.app.state.limiter |
| if not limiter.acquire(): |
| await websocket.close(code=1013) |
| return |
|
|
| await websocket.accept() |
|
|
| transport = WebSocketTransport(websocket) |
| session = RealtimeSession( |
| transport, |
| websocket.app.state.engine, |
| websocket.app.state.scheduler, |
| settings, |
| "websocket", |
| ) |
| session.start() |
|
|
| try: |
| while True: |
| incoming = await websocket.receive() |
|
|
| if incoming["type"] == "websocket.disconnect": |
| break |
|
|
| incoming_text = incoming.get("text") |
| incoming_bytes = incoming.get("bytes") |
|
|
| if incoming_text is not None: |
| session.handle_message(incoming_text) |
| elif incoming_bytes is not None: |
| session.handle_message(incoming_bytes) |
| except WebSocketDisconnect: |
| pass |
| finally: |
| await session.close() |
| limiter.release() |
|
|
