Mvr/#17/launch docker

config/client_aux2.yaml

@@ -0,0 +1,26 @@
+huri_url: ws://localhost:8000/session
+
+topic_list: ["transcript", "question", "rag_response"]
+sample_rate: 16000
+frame_duration: 0.030
+modules:
+  mic:
+    name: mic
+    args:
+      vad_agressiveness: 3
+      silence_duration: 1.5
+      block_duration: ${frame_duration}
+  stt:
+    name: stt
+    args:
+      language: "en"
+      block_duration: ${frame_duration}
+    logging: INFO
+  tag:
+    name: tag
+    logging: INFO
+  rag:
+    name: rag
+    args:
+      language: "en"
+      tone: "formal"

config/huri.yaml

@@ -11,10 +11,25 @@ logging_config:
   enable_access_log: true
   additional_log_standard_attrs: []
 
+services:
+  qdrant:
+    port: 6333
+    image: "qdrant/qdrant:latest"
+    storage_volume: "qdrant_data"
+  ollama:
+    model: "mistral:7b"
+    image: "ollama/ollama:rocm"
+    gpu_devices: true
+    num_replicas: 1
+
 applications:
   - name: huri-app
     route_prefix: /
     import_path: src.app:app
     runtime_env: { RAY_COLOR_PREFIX=1 }
     deployments:
       - name: HuRI
+      - name: RAGHandle
+        num_replicas: 2
+      - name: OllamaService
+      - name: QdrantService

requirements.txt

@@ -11,9 +11,12 @@ numpy
 ray[serve]
 webrtcvad
 faster-whisper
+qdrant-client
+sentence-transformers
 
 
 # client
 sounddevice
 websockets
-omegaconf
+omegaconf
+

src/app.py

@@ -1,14 +1,48 @@
+from pathlib import Path
+from typing import Any
+
+import yaml
 from ray.serve import Application
 
 from src.core.huri import HuRI
 from src.modules.factory import bind_deployment_handles
 from src.modules.modules import get_modules
+from src.modules.rag.docker_services import OllamaService, QdrantService
+
+
+def load_services_config() -> Any:
+    config_path = Path(__file__).resolve().parents[1] / "config" / "huri.yaml"
+    with open(config_path) as f:
+        config = yaml.safe_load(f)
+    return config.get("services", {})
+
+
+def build_qdrant(config: dict) -> Any:
+    return QdrantService.bind(  # type: ignore[attr-defined]
+        port=config.get("port", 6333),
+        image=config.get("image", "qdrant/qdrant:latest"),
+        storage_volume=config.get("storage_volume", "qdrant_data"),
+    )
+
+
+def build_ollama(config: dict) -> Any:
+    return OllamaService.options(  # type: ignore[attr-defined]
+        num_replicas=config.get("num_replicas", 1),
+    ).bind(
+        model=config.get("model", "mistral:7b"),
+        image=config.get("image", "ollama/ollama:latest"),
+        gpu_devices=config.get("gpu_devices", False),
+    )
 
 
 def build_app() -> Application:
     modules = get_modules()
-    handles = bind_deployment_handles(modules)
+    services_config = load_services_config()
+
+    qdrant = build_qdrant(services_config.get("qdrant", {}))
+    ollama = build_ollama(services_config.get("ollama", {}))
 
+    handles = bind_deployment_handles(modules, ollama=ollama, qdrant=qdrant)
     app: Application = HuRI.bind(modules, handles)  # type: ignore[attr-defined]
     return app
 

src/client.py

@@ -12,7 +12,6 @@
 
 from src.core.dataclasses.config import ClientConfig
 
-
 USER_ID_FILE = os.path.expanduser("~/.huri_user_id")
 
 
@@ -27,6 +26,7 @@ def save_user_id(_user_id: str):
     with open(USER_ID_FILE, "w") as f:
         f.write(_user_id)
 
+
 def load_client_config(path: str) -> ClientConfig:
     with open(path) as f:
         dict_config = OmegaConf.load(f)

src/core/huri.py

@@ -7,6 +7,7 @@
 
 from src.modules.factory import Module, ModuleFactory
 from src.modules.utils.sender import Sender
+
 from .app import app
 from .dataclasses.config import ClientConfig
 from .session import Session
@@ -57,4 +58,4 @@ async def receive_loop(session: Session, ws: WebSocket):
                 print(f"Client {_user_id} disconnected")
 
         await receive_loop(self.clients[session_id], ws)
-        del self.clients[session_id]
+        del self.clients[session_id]

src/core/module.py

@@ -14,10 +14,11 @@ async def process(self, _) -> Optional[Any]:
 class ModuleWithHandle(Module):
     _handle_cls: Type[Any]
 
-    def __init__(self, _handle: handle.DeploymentHandle = None, **kwargs):
+    def __init__(self, _handle: handle.DeploymentHandle | None = None, **kwargs):
         super().__init__(**kwargs)
         self._handle = _handle
 
+
 class ModuleWithId(Module):
     def __init__(self, _user_id: str, **kwargs):
         super().__init__(**kwargs)

src/modules/factory.py

@@ -1,16 +1,14 @@
-from os import name
 from typing import Any, Dict, List, Mapping, Type
 
 from src.core.dataclasses.config import ModuleConfig
-from src.core.module import Module, ModuleWithHandle, handle, ModuleWithId
+from src.core.module import Module, ModuleWithHandle, ModuleWithId, handle
 
 
 class ModuleFactory:
     def __init__(self, handles):
         self._registry: Dict[str, Type[Module]] = {}
         self._handles = handles
 
-
     def register(self, name: str, module_cls: Type[Module]) -> None:
         if not issubclass(module_cls, Module):
             raise TypeError(f"{module_cls} must inherit from Module")
@@ -21,13 +19,9 @@ def register(self, name: str, module_cls: Type[Module]) -> None:
                 )
         self._registry[name] = module_cls
 
-
     def create(
-        self,
-        _user_id: str,
-        name: str,
-        args: Mapping[str, Any] | None = None
-        ) -> Module:
+        self, _user_id: str, name: str, args: Mapping[str, Any] | None = None
+    ) -> Module:
 
         if name not in self._registry:
             raise ValueError(f"Unknown module '{name}'")
@@ -54,7 +48,9 @@ def create_from_config(
     ) -> List[Module]:
         modules: List[Module] = []
         for _, module_config in module_configs.items():
-            modules.append(self.create(_user_id, module_config.name, module_config.args))
+            modules.append(
+                self.create(_user_id, module_config.name, module_config.args)
+            )
         if modules == []:
             raise Exception
 
@@ -63,6 +59,7 @@ def create_from_config(
 
 def bind_deployment_handles(
     modules: Dict[str, Type[Module]],
+    **service_handles,
 ) -> Dict[str, handle.DeploymentHandle]:
     handles: Dict[str, handle.DeploymentHandle] = {}
     for name, module_cls in modules.items():
@@ -71,7 +68,15 @@ def bind_deployment_handles(
 
         if not hasattr(module_cls, "_handle_cls"):
             raise TypeError(f"{module_cls.__name__} must define _handle_cls")
+
         handle_cls = module_cls._handle_cls
-        handles[name] = handle_cls.bind()
+
+        if name == "rag" and service_handles:
+            handles[name] = handle_cls.bind(
+                ollama_handle=service_handles.get("ollama"),
+                qdrant_handle=service_handles.get("qdrant"),
+            )
+        else:
+            handles[name] = handle_cls.bind()
 
     return handles

src/modules/modules.py

@@ -1,9 +1,9 @@
 from typing import Dict, Type
 
+from src.modules.rag.rag import RAG
 from src.modules.speech_to_text.record_speech import MIC
 from src.modules.speech_to_text.speech_to_text import STT
 from src.modules.speech_to_text.text_aggregator import TAG
-from src.modules.rag.rag import RAG
 
 from .factory import Module