feat: setup method in ragbits-chat (#586)

Author: mhordynski

packages/ragbits-chat/CHANGELOG.md

@@ -2,6 +2,8 @@
 
 ## Unreleased
 
+- Add setup method to ChatInterface (#586)
+
 ## 0.19.1 (2025-05-27)
 
 ### Changed

packages/ragbits-chat/src/ragbits/chat/api.py

@@ -2,6 +2,7 @@
 import json
 import logging
 from collections.abc import AsyncGenerator
+from contextlib import asynccontextmanager
 from pathlib import Path
 from typing import Any, Literal
 
@@ -59,11 +60,17 @@ def __init__(
             cors_origins: List of allowed CORS origins. If None, defaults to common development origins.
             ui_build_dir: Path to a custom UI build directory. If None, uses the default package UI.
         """
-        self.app = FastAPI()
         self.chat_interface: ChatInterface = self._load_chat_interface(chat_interface)
         self.dist_dir = Path(ui_build_dir) if ui_build_dir else Path(__file__).parent / "ui-build"
         self.cors_origins = cors_origins or []
 
+        @asynccontextmanager
+        async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]:
+            await self.chat_interface.setup()
+            yield
+
+        self.app = FastAPI(lifespan=lifespan)
+
         self.configure_app()
         self.setup_routes()
         self.setup_exception_handlers()

packages/ragbits-chat/src/ragbits/chat/interface/_interface.py

@@ -150,6 +150,17 @@ def verify_state(state: dict[str, Any], signature: str) -> bool:
         expected_signature = ChatInterface._sign_state(state)
         return hmac.compare_digest(expected_signature, signature)
 
+    async def setup(self) -> None:  # noqa: B027
+        """
+        Setup the chat interface.
+
+        This method is called after the chat interface is initialized and before the chat method is called.
+        It is used to setup the chat interface, such as loading the model or initializing the vector store.
+
+        This method is optional and can be overridden by subclasses.
+        """
+        pass
+
     @abstractmethod
     async def chat(
         self,