142 lines
5.5 KiB
Python
142 lines
5.5 KiB
Python
"""Bridge between the game server and the AI service.
|
|
|
|
The AI service (``wizard_ai``) connects to ``/internal/ai`` via websocket.
|
|
Whenever an AI seat has to act, the server sends it the seat's observation
|
|
and waits for an action:
|
|
|
|
server -> ai: {"type": "decide", "request_id": "...", "observation": {...}}
|
|
ai -> server: {"type": "action", "request_id": "...", "action": {...}}
|
|
ai -> server: {"type": "error", "request_id": "...", "message": "..."}
|
|
|
|
Several AI service instances may connect; requests are distributed round
|
|
robin. If no instance is connected, it does not answer in time, or it answers
|
|
with an illegal action, a random legal action is played instead so a game
|
|
never gets stuck.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import itertools
|
|
import logging
|
|
import random
|
|
import uuid
|
|
from collections.abc import Callable
|
|
from typing import Any
|
|
|
|
from starlette.websockets import WebSocket, WebSocketDisconnect
|
|
|
|
from wizard_game import Action, Observation, action_from_dict
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class AIConnection:
|
|
def __init__(self, websocket: WebSocket, name: str) -> None:
|
|
self.websocket = websocket
|
|
self.name = name
|
|
self.pending: dict[str, asyncio.Future[dict[str, Any]]] = {}
|
|
self._send_lock = asyncio.Lock()
|
|
|
|
async def send(self, message: dict[str, Any]) -> None:
|
|
async with self._send_lock:
|
|
await self.websocket.send_json(message)
|
|
|
|
def fail_pending(self) -> None:
|
|
for future in self.pending.values():
|
|
if not future.done():
|
|
future.set_exception(ConnectionError("AI service disconnected"))
|
|
self.pending.clear()
|
|
|
|
|
|
class AIBridge:
|
|
def __init__(
|
|
self,
|
|
timeout: float,
|
|
rng: random.Random | None = None,
|
|
on_change: Callable[[], None] | None = None,
|
|
) -> None:
|
|
self.timeout = timeout
|
|
# Called whenever an AI service connects or disconnects.
|
|
self._on_change = on_change
|
|
self._connections: list[AIConnection] = []
|
|
self._round_robin = itertools.count()
|
|
self._rng = rng or random.Random()
|
|
|
|
@property
|
|
def available(self) -> bool:
|
|
return bool(self._connections)
|
|
|
|
async def decide(self, observation: Observation) -> Action:
|
|
"""Ask the AI service for an action; falls back to a random legal one."""
|
|
legal = observation.legal_actions
|
|
if not legal:
|
|
raise ValueError("Observation has no legal actions")
|
|
|
|
connection = self._pick_connection()
|
|
if connection is None:
|
|
return self._fallback(observation, "no AI service connected", logging.DEBUG)
|
|
|
|
request_id = uuid.uuid4().hex
|
|
future: asyncio.Future[dict[str, Any]] = asyncio.get_running_loop().create_future()
|
|
connection.pending[request_id] = future
|
|
try:
|
|
await connection.send({
|
|
"type": "decide",
|
|
"request_id": request_id,
|
|
"observation": observation.to_dict(),
|
|
})
|
|
reply = await asyncio.wait_for(future, self.timeout)
|
|
except (TimeoutError, ConnectionError, RuntimeError, WebSocketDisconnect) as exc:
|
|
return self._fallback(observation, f"AI request failed: {exc!r}")
|
|
finally:
|
|
connection.pending.pop(request_id, None)
|
|
|
|
if reply.get("type") != "action":
|
|
return self._fallback(observation, f"AI error: {reply.get('message')}")
|
|
try:
|
|
action = action_from_dict(reply.get("action", {}))
|
|
except ValueError as exc:
|
|
return self._fallback(observation, f"malformed AI action: {exc}")
|
|
if action not in legal:
|
|
return self._fallback(observation, f"illegal AI action {action}")
|
|
return action
|
|
|
|
async def serve(self, websocket: WebSocket) -> None:
|
|
"""Handle one AI service connection until it disconnects."""
|
|
connection = AIConnection(websocket, name=websocket.client.host if websocket.client else "?")
|
|
self._connections.append(connection)
|
|
logger.info("AI service connected from %s (%d connected)", connection.name, len(self._connections))
|
|
self._notify_change()
|
|
try:
|
|
while True:
|
|
message = await websocket.receive_json()
|
|
if not isinstance(message, dict):
|
|
continue
|
|
if message.get("type") == "hello":
|
|
logger.info("AI service %s: %s", connection.name, message.get("agent"))
|
|
continue
|
|
future = connection.pending.get(str(message.get("request_id")))
|
|
if future and not future.done():
|
|
future.set_result(message)
|
|
except (WebSocketDisconnect, RuntimeError, ValueError):
|
|
pass
|
|
finally:
|
|
self._connections.remove(connection)
|
|
connection.fail_pending()
|
|
logger.info("AI service %s disconnected (%d connected)", connection.name, len(self._connections))
|
|
self._notify_change()
|
|
|
|
def _notify_change(self) -> None:
|
|
if self._on_change is not None:
|
|
self._on_change()
|
|
|
|
def _pick_connection(self) -> AIConnection | None:
|
|
if not self._connections:
|
|
return None
|
|
return self._connections[next(self._round_robin) % len(self._connections)]
|
|
|
|
def _fallback(self, observation: Observation, reason: str, level: int = logging.WARNING) -> Action:
|
|
logger.log(level, "Playing random move for player %d: %s", observation.player, reason)
|
|
return self._rng.choice(observation.legal_actions)
|