first prototype

wizard_server/app/ai_bridge.py

@@ -0,0 +1,141 @@
+"""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)