taleweave-ai/adventure/server.py

import asyncio
from collections import deque
from json import dumps, loads
from logging import getLogger
from threading import Thread
from typing import Literal
from uuid import uuid4

import websockets

from adventure.context import get_actor_agent_for_name, set_actor_agent
from adventure.models.entity import Actor, Room, World
from adventure.models.event import (
    ActionEvent,
    GameEvent,
    GenerateEvent,
    PromptEvent,
    ReplyEvent,
    ResultEvent,
    StatusEvent,
)
from adventure.player import (
    RemotePlayer,
    get_player,
    has_player,
    list_players,
    remove_player,
    set_player,
)
from adventure.state import snapshot_world, world_json

logger = getLogger(__name__)

connected = set()
recent_events = deque(maxlen=100)
last_snapshot = None


async def handler(websocket):
    id = uuid4().hex
    logger.info("Client connected, given id: %s", id)
    connected.add(websocket)

    async def next_turn(character: str, prompt: str) -> None:
        await websocket.send(
            dumps(
                {
                    "type": "prompt",
                    "id": id,
                    "character": character,
                    "prompt": prompt,
                    "actions": [],
                }
            ),
        )

    def sync_turn(event: PromptEvent) -> bool:
        player = get_player(id)
        if player and player.name == event.actor.name:
            asyncio.run(next_turn(event.actor.name, event.prompt))
            return True

        return False

    try:
        await websocket.send(dumps({"type": "id", "id": id}))

        if last_snapshot:
            await websocket.send(last_snapshot)

        for message in recent_events:
            await websocket.send(message)
    except Exception:
        logger.exception("Failed to send recent messages to new client")

    while True:
        try:
            # if this socket is attached to a character and that character's turn is active, wait for input
            message = await websocket.recv()
            logger.info(f"Received message for {id}: {message}")

            try:
                data = loads(message)
                message_type = data.get("type", None)
                if message_type == "player":
                    character_name = data["become"]
                    if has_player(character_name):
                        logger.error(f"Character {character_name} is already in use")
                        continue

                    # TODO: should this always remove?
                    remove_player(id)

                    actor, llm_agent = get_actor_agent_for_name(character_name)
                    if not actor:
                        logger.error(f"Failed to find actor {character_name}")
                        continue

                    # prevent any recursive fallback bugs
                    if isinstance(llm_agent, RemotePlayer):
                        logger.warning(
                            "patching recursive fallback for %s", character_name
                        )
                        llm_agent = llm_agent.fallback_agent

                    # player_name = data["player"]
                    player = RemotePlayer(
                        actor.name, actor.backstory, sync_turn, fallback_agent=llm_agent
                    )
                    set_player(id, player)
                    logger.info(f"Client {id} is now character {character_name}")

                    # swap out the LLM agent
                    set_actor_agent(actor.name, actor, player)

                    # notify all clients that this character is now active
                    player_event(character_name, id, "join")
                    player_list()
                elif message_type == "input":
                    player = get_player(id)
                    if player and isinstance(player, RemotePlayer):
                        logger.info(
                            "queueing input for player %s: %s", player.name, data
                        )
                        player.input_queue.put(data["input"])

            except Exception:
                logger.exception("Failed to parse message")
        except websockets.ConnectionClosedOK:
            break

    connected.remove(websocket)

    # swap out the character for the original agent when they disconnect
    player = get_player(id)
    if player and isinstance(player, RemotePlayer):
        remove_player(id)

        logger.info("Disconnecting player for %s", player.name)
        player_event(player.name, id, "leave")
        player_list()

        actor, _ = get_actor_agent_for_name(player.name)
        if actor and player.fallback_agent:
            logger.info("Restoring LLM agent for %s", player.name)
            set_actor_agent(player.name, actor, player.fallback_agent)

    logger.info("Client disconnected: %s", id)


socket_thread = None
static_thread = None


def server_json(obj):
    if isinstance(obj, Actor):
        return obj.name

    if isinstance(obj, Room):
        return obj.name

    return world_json(obj)


def send_and_append(message):
    json_message = dumps(message, default=server_json)
    recent_events.append(json_message)
    websockets.broadcast(connected, json_message)
    return json_message


def launch_server():
    global socket_thread, static_thread

    def run_sockets():
        asyncio.run(server_main())

    socket_thread = Thread(target=run_sockets, daemon=True)
    socket_thread.start()

    return [socket_thread]


async def server_main():
    async with websockets.serve(handler, "", 8001):
        logger.info("Server started")
        await asyncio.Future()  # run forever


def server_system(world: World, step: int):
    global last_snapshot
    json_state = {
        **snapshot_world(world, step),
        "type": "world",
    }
    last_snapshot = send_and_append(json_state)


def server_result(room: Room, actor: Actor, action: str):
    json_action = {
        "actor": actor,
        "result": action,
        "room": room,
        "type": "result",
    }
    send_and_append(json_action)


def server_action(room: Room, actor: Actor, message: str):
    json_input = {
        "actor": actor,
        "input": message,
        "room": room,
        "type": "action",
    }
    send_and_append(json_input)


def server_generate(event: GenerateEvent):
    json_broadcast = {
        "name": event.name,
        "type": "generate",
    }
    send_and_append(json_broadcast)


def server_event(event: GameEvent):
    if isinstance(event, GenerateEvent):
        return server_generate(event)
    elif isinstance(event, ActionEvent):
        return server_action(event.room, event.actor, event.action)
    elif isinstance(event, ReplyEvent):
        return server_action(event.room, event.actor, event.text)
    elif isinstance(event, ResultEvent):
        return server_result(event.room, event.actor, event.result)
    elif isinstance(event, StatusEvent):
        pass
    else:
        logger.warning("Unknown event type: %s", event)


def player_event(character: str, id: str, event: Literal["join", "leave"]):
    json_broadcast = {
        "type": "player",
        "character": character,
        "id": id,
        "event": event,
    }
    send_and_append(json_broadcast)


def player_list():
    players = {value: key for key, value in list_players()}
    json_broadcast = {
        "type": "players",
        "players": players,
    }
    send_and_append(json_broadcast)