agent/worker.py

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'])])