server.py

import os
import shutil
import atexit
import yaml
import json
import numpy as np
import asyncio
from fastapi import FastAPI, WebSocket, APIRouter, Body
from fastapi.staticfiles import StaticFiles
from starlette.websockets import WebSocketDisconnect
from typing import List, Dict
from main import OpenLLMVTuberMain
from live2d_model import Live2dModel
from tts.stream_audio import AudioPayloadPreparer


class WebSocketServer:
    """
    WebSocketServer initializes a FastAPI application with WebSocket endpoints and a broadcast endpoint.

    Attributes:
        config (dict): Configuration dictionary.
        app (FastAPI): FastAPI application instance.
        router (APIRouter): APIRouter instance for routing.
        connected_clients (List[WebSocket]): List of connected WebSocket clients for "/client-ws".
        server_ws_clients (List[WebSocket]): List of connected WebSocket clients for "/server-ws".
    """

    def __init__(self, open_llm_vtuber_config: Dict | None = None):
        """
        Initializes the WebSocketServer with the given configuration.
        """
        self.app = FastAPI()
        self.router = APIRouter()
        self.new_connected_clients: List[WebSocket] = []
        self.connected_clients: List[WebSocket] = []
        self.server_ws_clients: List[WebSocket] = []
        self.open_llm_vtuber: OpenLLMVTuberMain | None = None
        self.open_llm_vtuber_config: Dict | None = open_llm_vtuber_config
        self._setup_routes()
        self._mount_static_files()

    def _setup_routes(self):
        """Sets up the WebSocket and broadcast routes."""

        # the connection between this server and the python backend
        @self.router.websocket("/server-ws")
        async def server_websocket_endpoint(websocket: WebSocket):
            await websocket.accept()
            self.server_ws_clients.append(websocket)
            # When a connection is established, send a specific payload to all clients connected to "/client-ws"
            control_message = {"type": "control", "text": "start-mic"}
            for client in self.connected_clients:
                await client.send_json(control_message)
            try:
                while True:
                    # Receive messages from "/server-ws" clients
                    message = await websocket.receive_text()
                    # Forward received messages to all clients connected to "/client-ws"
                    for client in self.connected_clients:
                        await client.send_text(message)
            except WebSocketDisconnect:
                self.server_ws_clients.remove(websocket)

        # the connection between this server and the frontend client
        @self.router.websocket("/legacy-ws")
        async def websocket_endpoint(websocket: WebSocket):
            await websocket.accept()
            self.connected_clients.append(websocket)
            try:
                while True:
                    # Receive messages from "/client-ws" clients
                    message = await websocket.receive_text()
                    # do some funny thing here...
                    # Forward received messages to all clients connected to "/server-ws"
                    for server_client in self.server_ws_clients:
                        await server_client.send_text(message)
            except WebSocketDisconnect:
                self.connected_clients.remove(websocket)

        # the connection between this server and the frontend client
        # The version 2 of the client-ws. Introduces breaking changes.
        # This route will initiate its own main.py instance and conversation loop
        @self.router.websocket("/client-ws")
        async def websocket_endpoint(websocket: WebSocket):
            await websocket.accept()
            await websocket.send_text(
                json.dumps({"type": "full-text", "text": "Connection established"})
            )

            self.connected_clients.append(websocket)
            print("Connection established")
            l2d = Live2dModel(self.open_llm_vtuber_config["LIVE2D_MODEL"])
            self.open_llm_vtuber = OpenLLMVTuberMain(self.open_llm_vtuber_config)
            audio_payload_preparer = AudioPayloadPreparer()

            def _play_audio_file(sentence: str | None, filepath: str | None) -> None:
                if filepath is None:
                    print("No audio to be streamed. Response is empty.")
                    return

                if sentence is None:
                    sentence = ""
                print(f">> Playing {filepath}...")
                payload, duration = audio_payload_preparer.prepare_audio_payload(
                    audio_path=filepath,
                    display_text=sentence,
                    expression_list=l2d.extract_emotion(sentence),
                )
                print("Payload send.")

                async def _send_audio():
                    await websocket.send_text(json.dumps(payload))
                    await asyncio.sleep(duration)

                new_loop = asyncio.new_event_loop()
                asyncio.set_event_loop(new_loop)
                new_loop.run_until_complete(_send_audio())
                new_loop.close()

                print("Audio played.")

            self.open_llm_vtuber.set_audio_output_func(_play_audio_file)

            await websocket.send_text(
                json.dumps({"type": "set-model", "text": l2d.model_info})
            )
            print("Model set")
            received_data_buffer = np.array([])
            # start mic
            await websocket.send_text(
                json.dumps({"type": "control", "text": "start-mic"})
            )

            conversation_task = None

            try:
                while True:
                    print(".", end="")
                    message = await websocket.receive_text()
                    data = json.loads(message)
                    # print(f"\033\n Received ws req: {data.get('type')}\033[0m\n")

                    if data.get("type") == "interrupt-signal":
                        print("Start receiving audio data from front end.")
                        if conversation_task is not None:
                            print(
                                "\033[91mLLM hadn't finish itself. Interrupting it...",
                                "heard response: \n",
                                data.get("text"),
                                "\033[0m\n",
                            )
                            self.open_llm_vtuber.interrupt(data.get("text"))
                            # conversation_task.cancel()

                    elif data.get("type") == "mic-audio-data":
                        received_data_buffer = np.append(
                            received_data_buffer,
                            np.array(
                                list(data.get("audio").values()), dtype=np.float32
                            ),
                        )
                        print("*", end="")

                    elif data.get("type") == "mic-audio-end":
                        print("Received audio data end from front end.")
                        await websocket.send_text(
                            json.dumps({"type": "full-text", "text": "Thinking..."})
                        )
                        audio = received_data_buffer
                        received_data_buffer = np.array([])

                        async def _run_conversation():
                            try:
                                await websocket.send_text(
                                    json.dumps(
                                        {
                                            "type": "control",
                                            "text": "conversation-chain-start",
                                        }
                                    )
                                )
                                await asyncio.to_thread(
                                    self.open_llm_vtuber.conversation_chain,
                                    user_input=audio,
                                )
                                await websocket.send_text(
                                    json.dumps(
                                        {
                                            "type": "control",
                                            "text": "conversation-chain-end",
                                        }
                                    )
                                )
                                print("One Conversation Loop Completed")
                            except asyncio.CancelledError:
                                print("Conversation task was cancelled.")
                            except InterruptedError as e:
                                print(f"😢Conversation was interrupted. {e}")

                        conversation_task = asyncio.create_task(_run_conversation())
                    else:
                        print("Unknown data type received.")

            except WebSocketDisconnect:
                self.connected_clients.remove(websocket)
                self.open_llm_vtuber = None

        @self.router.post("/broadcast")
        async def broadcast_message(message: str = Body(..., embed=True)):
            disconnected_clients = []
            for client in self.connected_clients:
                try:
                    await client.send_text(message)
                except WebSocketDisconnect:
                    disconnected_clients.append(client)
            for client in disconnected_clients:
                self.connected_clients.remove(client)

        self.app.include_router(self.router)

    def _mount_static_files(self):
        """Mounts static file directories."""
        self.app.mount(
            "/live2d-models",
            StaticFiles(directory="live2d-models"),
            name="live2d-models",
        )
        self.app.mount("/", StaticFiles(directory="./static", html=True), name="static")

    def run(self, host: str = "127.0.0.1", port: int = 8000, log_level: str = "info"):
        """Runs the FastAPI application using Uvicorn."""
        import uvicorn

        uvicorn.run(self.app, host=host, port=port, log_level=log_level)
        
    def clean_cache():
        cache_dir = "./cache"
        if os.path.exists(cache_dir):
            shutil.rmtree(cache_dir)
            os.makedirs(cache_dir)


if __name__ == "__main__":
    
    atexit.register(WebSocketServer.clean_cache)
    
    # Load configurations from yaml file
    with open("conf.yaml", "rb") as f:
        config = yaml.safe_load(f)

    config["LIVE2D"] = True  # make sure the live2d is enabled

    # Initialize and run the WebSocket server
    server = WebSocketServer(open_llm_vtuber_config=config)
    server.run(host=config["HOST"], port=config["PORT"])