Update src/open_llm_vtuber/server.py

src/open_llm_vtuber/server.py

@@ -1,210 +1,211 @@
-"""
-Open-LLM-VTuber Server
-========================
-This module contains the WebSocket server for Open-LLM-VTuber, which handles
-the WebSocket connections, serves static files, and manages the web tool.
-It uses FastAPI for the server and Starlette for static file serving.
-"""
-
-import os
-import shutil
-
-from fastapi import FastAPI
-from starlette.middleware.cors import CORSMiddleware
-from starlette.responses import Response
-from starlette.staticfiles import StaticFiles as StarletteStaticFiles
-
-from .routes import init_client_ws_route, init_webtool_routes, init_proxy_route
-from .service_context import ServiceContext
-from .config_manager.utils import Config
-from .openllm_vtuber_main import OpenLLMVTuberMain
-
-# Create a custom StaticFiles class that adds CORS headers
-class CORSStaticFiles(StarletteStaticFiles):
-    """
-    Static files handler that adds CORS headers to all responses.
-    Needed because Starlette StaticFiles might bypass standard middleware.
-    """
-
-    async def get_response(self, path: str, scope):
-        response = await super().get_response(path, scope)
-
-        # Add CORS headers to all responses
-        response.headers["Access-Control-Allow-Origin"] = "*"
-        response.headers["Access-Control-Allow-Methods"] = "GET, OPTIONS"
-        response.headers["Access-Control-Allow-Headers"] = "*"
-
-        if path.endswith(".js"):
-            response.headers["Content-Type"] = "application/javascript"
-
-        return response
-
-
-class AvatarStaticFiles(CORSStaticFiles):
-    """
-    Avatar files handler with security restrictions and CORS headers
-    """
-
-    async def get_response(self, path: str, scope):
-        allowed_extensions = (".jpg", ".jpeg", ".png", ".gif", ".svg")
-        if not any(path.lower().endswith(ext) for ext in allowed_extensions):
-            return Response("Forbidden file type", status_code=403)
-        response = await super().get_response(path, scope)
-        return response
-
-
-class WebSocketServer:
-    """
-    API server for Open-LLM-VTuber. This contains the websocket endpoint for the client, hosts the web tool, and serves static files.
-
-    Creates and configures a FastAPI app, registers all routes
-    (WebSocket, web tools, proxy) and mounts static assets with CORS.
-
-    Args:
-        config (Config): Application configuration containing system settings.
-        default_context_cache (ServiceContext, optional):
-            Pre‑initialized service context for sessions' service context to reference to.
-            **If omitted, `initialize()` method needs to be called to load service context.**
-
-    Notes:
-        - If default_context_cache is omitted, call `await initialize()` to load service context cache.
-        - Use `clean_cache()` to clear and recreate the local cache directory.
-    """
-
-    def __init__(self, config: Config, default_context_cache: ServiceContext = None):
-        self.app = FastAPI(title="Open-LLM-VTuber Server")  # Added title for clarity
-        self.config = config
-        self.vtuber_main = None
-        self.is_ready = False
-        self.default_context_cache = (
-            default_context_cache or ServiceContext()
-        )  # Use provided context or initialize a new empty one waiting to be loaded
-        # It will be populated during the initialize method call
-
-        # Add global CORS middleware
-        self.app.add_middleware(
-            CORSMiddleware,
-            allow_origins=["*"],
-            allow_credentials=True,
-            allow_methods=["*"],
-            allow_headers=["*"],
-        )
-
-        
-        # The context will be populated during the initialize step
-        self.app.include_router(
-            init_client_ws_route(default_context_cache=self.default_context_cache, server_instance=self)
-        )
-
-        # Initialize and include proxy routes if proxy is enabled
-        system_config = config.system_config
-        if hasattr(system_config, "enable_proxy") and system_config.enable_proxy:
-            # Construct the server URL for the proxy
-            host = system_config.host
-            port = system_config.port
-            server_url = f"ws://{host}:{port}/client-ws"
-            self.app.include_router(
-                init_proxy_route(server_url=server_url),
-            )
-
-        # Mount cache directory first (to ensure audio file access)
-        if not os.path.exists("cache"):
-            os.makedirs("cache")
-        if not os.path.exists("sing"):
-            os.makedirs("sing/original", exist_ok=True)
-        self.app.mount(
-            "/cache",
-            CORSStaticFiles(directory="cache"),
-            name="cache",
-        )
-
-        # Mount static files with CORS-enabled handlers
-        self.app.mount(
-            "/live2d-models",
-            CORSStaticFiles(directory="live2d-models"),
-            name="live2d-models",
-        )
-        self.app.mount(
-            "/bg",
-            CORSStaticFiles(directory="backgrounds"),
-            name="backgrounds",
-        )
-        self.app.mount(
-            "/avatars",
-            AvatarStaticFiles(directory="avatars"),
-            name="avatars",
-        )
-        self.app.mount(
-            "/sing/tracks", 
-            CORSStaticFiles(directory="sing/tracks"),
-            name="sing_tracks"
-        )
-        
-        # Mount web tool directory separately from frontend
-        self.app.mount(
-            "/web-tool",
-            CORSStaticFiles(directory="web_tool", html=True),
-            name="web_tool",
-        )
-
-        # Mount main frontend last (as catch-all)
-        self.app.mount(
-            "/",
-            CORSStaticFiles(directory="frontend", html=True),
-            name="frontend",
-        )
-        self.is_ready = False
-
-    async def initialize(self):
-        """Asynchronously load the service context and VTuber Main logic."""
-        import asyncio
-        from loguru import logger
-        import traceback
-        
-        await self.default_context_cache.load_from_config(self.config)
-        
-        try:
-            configs_dict = self.config.model_dump()
-        except:
-            configs_dict = vars(self.config)
-
-        # --- ĐOẠN SỬA LỖI: Tự động map persona_prompt sang system_prompt ---
-        if isinstance(configs_dict, dict):
-            char_cfg = configs_dict.get('character_config', {})
-            
-            # 1. Nếu system_prompt trống nhưng persona_prompt có dữ liệu, thì lấy persona_prompt
-            if char_cfg.get('system_prompt') is None:
-                char_cfg['system_prompt'] = char_cfg.get('persona_prompt', "")
-            
-            # 2. Đảm bảo các chuỗi quan trọng không bao giờ là None để tránh lỗi +=
-            keys_to_fix = [
-                'system_prompt', 'personality', 'instruction', 
-                'system_with_tools', 'name', 'background'
-            ]
-            for key in keys_to_fix:
-                if key in char_cfg and char_cfg[key] is None:
-                    char_cfg[key] = ""
-                # Kiểm tra cả ở cấp ngoài nếu có
-                if key in configs_dict and configs_dict[key] is None:
-                    configs_dict[key] = ""
-        # -----------------------------------------------------------------
-            
-        try:
-            self.vtuber_main = OpenLLMVTuberMain(
-                configs=configs_dict, 
-                loop=asyncio.get_running_loop()
-            )
-            self.is_ready = True 
-            logger.info("✅ VTuber Main Logic initialized and Ready.")
-        except Exception as e:
-            self.is_ready = False 
-            logger.error(f"❌ Failed to initialize VTuber Main: {e}")
-            logger.error(traceback.format_exc())
-
-    @staticmethod
-    def clean_cache():
-        """Clean the cache directory by removing and recreating it."""
-        cache_dir = "cache"
-        if os.path.exists(cache_dir):
-            shutil.rmtree(cache_dir)
-            os.makedirs(cache_dir)
+"""
+Open-LLM-VTuber Server
+========================
+This module contains the WebSocket server for Open-LLM-VTuber, which handles
+the WebSocket connections, serves static files, and manages the web tool.
+It uses FastAPI for the server and Starlette for static file serving.
+"""
+
+import os
+import shutil
+
+from fastapi import FastAPI
+from starlette.middleware.cors import CORSMiddleware
+from starlette.responses import Response
+from starlette.staticfiles import StaticFiles as StarletteStaticFiles
+
+from .routes import init_client_ws_route, init_webtool_routes, init_proxy_route
+from .service_context import ServiceContext
+from .config_manager.utils import Config
+from .openllm_vtuber_main import OpenLLMVTuberMain
+
+# Create a custom StaticFiles class that adds CORS headers
+class CORSStaticFiles(StarletteStaticFiles):
+    """
+    Static files handler that adds CORS headers to all responses.
+    Needed because Starlette StaticFiles might bypass standard middleware.
+    """
+
+    async def get_response(self, path: str, scope):
+        response = await super().get_response(path, scope)
+
+        # Add CORS headers to all responses
+        response.headers["Access-Control-Allow-Origin"] = "*"
+        response.headers["Access-Control-Allow-Methods"] = "GET, OPTIONS"
+        response.headers["Access-Control-Allow-Headers"] = "*"
+
+        if path.endswith(".js"):
+            response.headers["Content-Type"] = "application/javascript"
+
+        return response
+
+
+class AvatarStaticFiles(CORSStaticFiles):
+    """
+    Avatar files handler with security restrictions and CORS headers
+    """
+
+    async def get_response(self, path: str, scope):
+        allowed_extensions = (".jpg", ".jpeg", ".png", ".gif", ".svg")
+        if not any(path.lower().endswith(ext) for ext in allowed_extensions):
+            return Response("Forbidden file type", status_code=403)
+        response = await super().get_response(path, scope)
+        return response
+
+
+class WebSocketServer:
+    """
+    API server for Open-LLM-VTuber. This contains the websocket endpoint for the client, hosts the web tool, and serves static files.
+
+    Creates and configures a FastAPI app, registers all routes
+    (WebSocket, web tools, proxy) and mounts static assets with CORS.
+
+    Args:
+        config (Config): Application configuration containing system settings.
+        default_context_cache (ServiceContext, optional):
+            Pre‑initialized service context for sessions' service context to reference to.
+            **If omitted, `initialize()` method needs to be called to load service context.**
+
+    Notes:
+        - If default_context_cache is omitted, call `await initialize()` to load service context cache.
+        - Use `clean_cache()` to clear and recreate the local cache directory.
+    """
+
+    def __init__(self, config: Config, default_context_cache: ServiceContext = None):
+        self.app = FastAPI(title="Open-LLM-VTuber Server")  # Added title for clarity
+        self.config = config
+        self.vtuber_main = None
+        self.is_ready = False
+        self.default_context_cache = (
+            default_context_cache or ServiceContext()
+        )  # Use provided context or initialize a new empty one waiting to be loaded
+        # It will be populated during the initialize method call
+
+        # Add global CORS middleware
+        self.app.add_middleware(
+            CORSMiddleware,
+            allow_origins=["*"],
+            allow_credentials=True,
+            allow_methods=["*"],
+            allow_headers=["*"],
+        )
+
+        
+        # The context will be populated during the initialize step
+        self.app.include_router(
+            init_client_ws_route(default_context_cache=self.default_context_cache, server_instance=self)
+        )
+
+        # Initialize and include proxy routes if proxy is enabled
+        system_config = config.system_config
+        if hasattr(system_config, "enable_proxy") and system_config.enable_proxy:
+            # Construct the server URL for the proxy
+            host = system_config.host
+            port = system_config.port
+            server_url = f"ws://{host}:{port}/client-ws"
+            self.app.include_router(
+                init_proxy_route(server_url=server_url),
+            )
+
+        # Mount cache directory first (to ensure audio file access)
+        if not os.path.exists("cache"):
+            os.makedirs("cache")
+        # Always ensure all required sing subdirectories exist
+        os.makedirs("sing/original", exist_ok=True)
+        os.makedirs("sing/tracks", exist_ok=True)
+        self.app.mount(
+            "/cache",
+            CORSStaticFiles(directory="cache"),
+            name="cache",
+        )
+
+        # Mount static files with CORS-enabled handlers
+        self.app.mount(
+            "/live2d-models",
+            CORSStaticFiles(directory="live2d-models"),
+            name="live2d-models",
+        )
+        self.app.mount(
+            "/bg",
+            CORSStaticFiles(directory="backgrounds"),
+            name="backgrounds",
+        )
+        self.app.mount(
+            "/avatars",
+            AvatarStaticFiles(directory="avatars"),
+            name="avatars",
+        )
+        self.app.mount(
+            "/sing/tracks",
+            CORSStaticFiles(directory="sing/tracks"),
+            name="sing_tracks"
+        )
+        
+        # Mount web tool directory separately from frontend
+        self.app.mount(
+            "/web-tool",
+            CORSStaticFiles(directory="web_tool", html=True),
+            name="web_tool",
+        )
+
+        # Mount main frontend last (as catch-all)
+        self.app.mount(
+            "/",
+            CORSStaticFiles(directory="frontend", html=True),
+            name="frontend",
+        )
+        self.is_ready = False
+
+    async def initialize(self):
+        """Asynchronously load the service context and VTuber Main logic."""
+        import asyncio
+        from loguru import logger
+        import traceback
+        
+        await self.default_context_cache.load_from_config(self.config)
+        
+        try:
+            configs_dict = self.config.model_dump()
+        except:
+            configs_dict = vars(self.config)
+
+        # --- ĐOẠN SỬA LỖI: Tự động map persona_prompt sang system_prompt ---
+        if isinstance(configs_dict, dict):
+            char_cfg = configs_dict.get('character_config', {})
+            
+            # 1. Nếu system_prompt trống nhưng persona_prompt có dữ liệu, thì lấy persona_prompt
+            if char_cfg.get('system_prompt') is None:
+                char_cfg['system_prompt'] = char_cfg.get('persona_prompt', "")
+            
+            # 2. Đảm bảo các chuỗi quan trọng không bao giờ là None để tránh lỗi +=
+            keys_to_fix = [
+                'system_prompt', 'personality', 'instruction', 
+                'system_with_tools', 'name', 'background'
+            ]
+            for key in keys_to_fix:
+                if key in char_cfg and char_cfg[key] is None:
+                    char_cfg[key] = ""
+                # Kiểm tra cả ở cấp ngoài nếu có
+                if key in configs_dict and configs_dict[key] is None:
+                    configs_dict[key] = ""
+        # -----------------------------------------------------------------
+            
+        try:
+            self.vtuber_main = OpenLLMVTuberMain(
+                configs=configs_dict, 
+                loop=asyncio.get_running_loop()
+            )
+            self.is_ready = True 
+            logger.info("✅ VTuber Main Logic initialized and Ready.")
+        except Exception as e:
+            self.is_ready = False 
+            logger.error(f"❌ Failed to initialize VTuber Main: {e}")
+            logger.error(traceback.format_exc())
+
+    @staticmethod
+    def clean_cache():
+        """Clean the cache directory by removing and recreating it."""
+        cache_dir = "cache"
+        if os.path.exists(cache_dir):
+            shutil.rmtree(cache_dir)
+            os.makedirs(cache_dir)