76 lines
3.1 KiB
Python
76 lines
3.1 KiB
Python
"""Production mode: serve decisions to the game server over a websocket.
|
|
|
|
The service connects to the server's ``/internal/ai`` endpoint and answers
|
|
``decide`` requests (see ``wizard_server/app/ai_bridge.py``). It reconnects
|
|
automatically if the connection drops.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import json
|
|
import logging
|
|
from typing import Any
|
|
|
|
from websockets.asyncio.client import ClientConnection, connect
|
|
from websockets.exceptions import ConnectionClosed, InvalidHandshake, InvalidStatus
|
|
|
|
from wizard_game import Observation
|
|
|
|
from ..agents import Agent
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
MAX_BACKOFF = 30.0
|
|
|
|
|
|
class AIService:
|
|
def __init__(self, server_url: str, agent: Agent, secret: str = "") -> None:
|
|
self.server_url = server_url
|
|
self.agent = agent
|
|
self.secret = secret
|
|
# Inference runs in worker threads; the lock keeps agents that are not
|
|
# thread safe (e.g. stateful models) consistent.
|
|
self._agent_lock = asyncio.Lock()
|
|
|
|
async def run_forever(self) -> None:
|
|
backoff = 1.0
|
|
while True:
|
|
try:
|
|
await self._run_once()
|
|
backoff = 1.0
|
|
except (OSError, ConnectionClosed, InvalidHandshake, TimeoutError) as exc:
|
|
if isinstance(exc, InvalidStatus) and exc.response.status_code in (401, 403):
|
|
logger.error("Server rejected the connection, check AI_SHARED_SECRET")
|
|
logger.warning("Connection to %s lost (%s), retrying in %.0fs", self.server_url, exc, backoff)
|
|
await asyncio.sleep(backoff)
|
|
backoff = min(backoff * 2, MAX_BACKOFF)
|
|
|
|
async def _run_once(self) -> None:
|
|
headers = {"Authorization": f"Bearer {self.secret}"} if self.secret else {}
|
|
async with connect(self.server_url, additional_headers=headers, open_timeout=10) as websocket:
|
|
logger.info("Connected to %s as '%s' agent", self.server_url, self.agent.name)
|
|
await websocket.send(json.dumps({"type": "hello", "agent": self.agent.name}))
|
|
tasks: set[asyncio.Task] = set()
|
|
async for raw in websocket:
|
|
message = json.loads(raw)
|
|
if message.get("type") == "decide":
|
|
task = asyncio.create_task(self._decide(websocket, message))
|
|
tasks.add(task)
|
|
task.add_done_callback(tasks.discard)
|
|
|
|
async def _decide(self, websocket: ClientConnection, message: dict[str, Any]) -> None:
|
|
request_id = message.get("request_id")
|
|
try:
|
|
observation = Observation.from_dict(message["observation"])
|
|
async with self._agent_lock:
|
|
action = await asyncio.to_thread(self.agent.act, observation)
|
|
reply = {"type": "action", "request_id": request_id, "action": action.to_dict()}
|
|
except Exception as exc: # report instead of crashing; the server falls back
|
|
logger.exception("Failed to decide request %s", request_id)
|
|
reply = {"type": "error", "request_id": request_id, "message": str(exc)}
|
|
try:
|
|
await websocket.send(json.dumps(reply))
|
|
except ConnectionClosed:
|
|
pass
|