from __future__ import annotations import io, uuid from starlette.applications import Starlette from starlette.requests import Request from starlette.responses import JSONResponse from starlette.routing import Route from PIL import Image from nitrogen_runtime import runtime SESSIONS={} def error(e,status=400): return JSONResponse({'error':str(e),'type':type(e).__name__},status_code=status) async def health(_): return JSONResponse({'ok':True,'runtime':runtime.info(),'sessions':len(SESSIONS)}) async def create(request): try: b=await request.json(); s=await runtime.create_session(b.get('title'),b.get('game_override'),int(b.get('context',1))) sid=uuid.uuid4().hex[:12]; SESSIONS[sid]=s; match,score=runtime.match_game(b.get('title')) return JSONResponse({'session_id':sid,'selected_game':s.selected_game,'match_score':score,'runtime':runtime.info()}) except Exception as e:return error(e,500) async def info(request): sid=request.path_params['sid']; s=SESSIONS.get(sid) if not s:return error(KeyError(sid),404) return JSONResponse({'session_id':sid,'selected_game':s.selected_game}) async def predict(request): try: sid=request.path_params['sid']; s=SESSIONS[sid]; raw=await request.body(); img=Image.open(io.BytesIO(raw)).convert('RGB'); return JSONResponse(await runtime.predict(s,img)) except Exception as e:return error(e,500) async def reset(request): try: SESSIONS[request.path_params['sid']].reset(); return JSONResponse({'ok':True}) except Exception as e:return error(e) async def close(request): SESSIONS.pop(request.path_params['sid'],None); return JSONResponse({'ok':True}) app=Starlette(routes=[Route('/health',health),Route('/sessions',create,methods=['POST']),Route('/sessions/{sid}',info),Route('/sessions/{sid}/predict',predict,methods=['POST']),Route('/sessions/{sid}/reset',reset,methods=['POST']),Route('/sessions/{sid}',close,methods=['DELETE'])])