feat: added LMStudio plugin

jcabrero · jcabrero · commit 022c4136bce2 · 2025-10-09T12:32:51.000+02:00
diff --git a/docker-compose.dev.yml b/docker-compose.dev.yml
@@ -31,6 +31,8 @@ services:
     depends_on:
       postgres:
         condition: service_healthy
+      nilauth-credit-server:
+        condition: service_healthy
     environment:
       - POSTGRES_DB=${POSTGRES_DB_NUC}
     volumes:
@@ -104,7 +106,11 @@ services:
     depends_on:
       nilauth-postgres:
         condition: service_healthy
-
-
+    healthcheck:
+      test: ["CMD", "wget", "--no-verbose", "--tries=1", "--spider", "http://localhost:3000/health"]
+      interval: 30s
+      retries: 3
+      start_period: 15s
+      timeout: 10s
 volumes:
   postgres_data:
diff --git a/docker/README.md b/docker/README.md
@@ -34,3 +34,17 @@ docker run -d --name etcd-server \
     --env ETCD_ADVERTISE_CLIENT_URLS=http://etcd-server:2379 \
     bitnami/etcd:latest
 ```
+
+## Announcing LMStudio Models
+
+LMStudio can run on the host at `localhost:1234` while the stack runs inside Docker. Build the announcer image and bring it up alongside the core services:
+
+```sh
+docker build -t nilai/lmstudio-announcer:latest -f docker/lmstudio-announcer.Dockerfile .
+docker compose -f docker-compose.yml \
+  -f docker-compose.dev.yml \
+  -f docker/compose/docker-compose.lmstudio.yml \
+  up -d lmstudio_announcer
+```
+
+The announcer registers every model returned by `http://host.docker.internal:1234/v1/models` in etcd so that `nilai-api` can route chat requests to LMStudio. Override defaults with environment variables such as `LMSTUDIO_MODEL_IDS`, `LMSTUDIO_TOOL_SUPPORT_MODELS`, or `LMSTUDIO_MULTIMODAL_MODELS` inside the compose override.
diff --git a/docker/compose/docker-compose.lmstudio.yml b/docker/compose/docker-compose.lmstudio.yml
@@ -0,0 +1,16 @@
+services:
+  lmstudio_announcer:
+    image: nilai/lmstudio-announcer:latest
+    container_name: nilai-lmstudio-announcer
+    restart: unless-stopped
+    depends_on:
+      etcd:
+        condition: service_healthy
+    environment:
+      - SVC_HOST=host.docker.internal
+      - SVC_PORT=1234
+      - ETCD_HOST=etcd
+      - ETCD_PORT=2379
+      - LMSTUDIO_SUPPORTED_FEATURES=chat_completion
+    extra_hosts:
+      - "host.docker.internal:host-gateway"
diff --git a/docker/lmstudio-announcer.Dockerfile b/docker/lmstudio-announcer.Dockerfile
@@ -0,0 +1,14 @@
+FROM python:3.12-slim
+
+ENV PYTHONUNBUFFERED=1 \
+    PIP_NO_CACHE_DIR=1
+
+WORKDIR /app
+
+COPY --link packages/nilai-common /app/packages/nilai-common
+COPY --link nilai-models /app/nilai-models
+
+RUN pip install --upgrade pip && \
+    pip install /app/packages/nilai-common /app/nilai-models
+
+ENTRYPOINT ["python", "-m", "nilai_models.lmstudio_announcer"]
diff --git a/nilai-models/src/nilai_models/daemon.py b/nilai-models/src/nilai_models/daemon.py
@@ -1,45 +1,44 @@
 # nilai/models/model.py
 import asyncio
-import signal
 import logging
+import signal
+
 import httpx
 
-from nilai_common import (  # Model service discovery and host settings
-    SETTINGS,
+from nilai_common import (
     MODEL_SETTINGS,
-    ModelServiceDiscovery,
+    SETTINGS,
     ModelEndpoint,
     ModelMetadata,
+    ModelServiceDiscovery,
 )
 
 logger = logging.getLogger(__name__)
 
 
 async def get_metadata():
-    """Fetch model metadata from model
-    service and return as ModelMetadata object"""
+    """Fetch model metadata from model service and return as ModelMetadata object."""
     current_retries = 0
     while True:
         url = None
         try:
             url = f"http://{SETTINGS.host}:{SETTINGS.port}/v1/models"
-            # Request model metadata from localhost:8000/v1/models
             async with httpx.AsyncClient() as client:
                 response = await client.get(url)
                 response.raise_for_status()
                 response_data = response.json()
                 model_name = response_data["data"][0]["id"]
                 return ModelMetadata(
-                    id=model_name,  # Unique identifier
-                    name=model_name,  # Human-readable name
-                    version="1.0",  # Model version
+                    id=model_name,
+                    name=model_name,
+                    version="1.0",
                     description="",
-                    author="",  # Model creators
-                    license="Apache 2.0",  # Usage license
-                    source=f"https://huggingface.co/{model_name}",  # Model source
-                    supported_features=["chat_completion"],  # Capabilities
-                    tool_support=SETTINGS.tool_support,  # Tool support
-                    multimodal_support=SETTINGS.multimodal_support,  # Multimodal support
+                    author="",
+                    license="Apache 2.0",
+                    source=f"https://huggingface.co/{model_name}",
+                    supported_features=["chat_completion"],
+                    tool_support=SETTINGS.tool_support,
+                    multimodal_support=SETTINGS.multimodal_support,
                 )
 
         except Exception as e:
@@ -49,16 +48,16 @@ async def get_metadata():
                 logger.warning(f"Failed to fetch model metadata from {url}: {e}")
             current_retries += 1
             if (
-                MODEL_SETTINGS.num_retries
-                != -1  # If num_retries == -1 then we do infinite number of retries
+                MODEL_SETTINGS.num_retries != -1
                 and current_retries >= MODEL_SETTINGS.num_retries
             ):
                 raise e
             await asyncio.sleep(MODEL_SETTINGS.timeout)
 
 
 async def run_service(discovery_service, model_endpoint):
-    """Runs the model service and keeps it alive"""
+    """Register model with discovery service and keep it alive."""
+    lease = None
     try:
         logger.info(f"Registering model: {model_endpoint.metadata.id}")
         lease = await discovery_service.register_model(model_endpoint, prefix="/models")
@@ -73,50 +72,62 @@ async def run_service(discovery_service, model_endpoint):
         logger.error(f"Service error: {e}")
         raise
     finally:
-        try:
-            await discovery_service.unregister_model(model_endpoint.metadata.id)
-            logger.info(f"Model unregistered: {model_endpoint.metadata.id}")
-        except Exception as e:
-            logger.error(f"Error unregistering model: {e}")
+        if lease:
+            try:
+                await discovery_service.unregister_model(model_endpoint.metadata.id)
+                logger.info(f"Model unregistered: {model_endpoint.metadata.id}")
+            except Exception as e:
+                logger.error(f"Error unregistering model: {e}")
 
 
 async def main():
-    discovery_service = None
-    model_endpoint = None
-
-    try:
-        # Initialize discovery service
-        discovery_service = ModelServiceDiscovery(
-            host=SETTINGS.etcd_host, port=SETTINGS.etcd_port
-        )
-
-        metadata = await get_metadata()
-        model_endpoint = ModelEndpoint(
-            url=f"http://{SETTINGS.host}:{SETTINGS.port}", metadata=metadata
-        )
+    """Main entry point for model daemon."""
+    logging.basicConfig(level=logging.INFO)
 
-        # Setup signal handlers
-        loop = asyncio.get_running_loop()
-        for sig in (signal.SIGTERM, signal.SIGINT):
-            loop.add_signal_handler(sig, lambda: asyncio.create_task(shutdown()))
+    # Initialize discovery service
+    discovery_service = ModelServiceDiscovery(
+        host=SETTINGS.etcd_host, port=SETTINGS.etcd_port
+    )
 
-        # Run service
-        await run_service(discovery_service, model_endpoint)
+    # Fetch metadata and create endpoint
+    metadata = await get_metadata()
+    model_endpoint = ModelEndpoint(
+        url=f"http://{SETTINGS.host}:{SETTINGS.port}", metadata=metadata
+    )
 
-    except Exception as e:
-        logger.error(f"Failed to initialize model service: {e}")
-        raise
+    # Create service task
+    service_task = asyncio.create_task(run_service(discovery_service, model_endpoint))
 
-
-async def shutdown():
-    """Cleanup and shutdown"""
-    tasks = [t for t in asyncio.all_tasks() if t is not asyncio.current_task()]
-    [task.cancel() for task in tasks]
-    await asyncio.gather(*tasks, return_exceptions=True)
+    # Setup signal handling
+    stop_event = asyncio.Event()
     loop = asyncio.get_running_loop()
-    loop.stop()
+    for sig in (signal.SIGTERM, signal.SIGINT):
+        try:
+            loop.add_signal_handler(sig, stop_event.set)
+        except NotImplementedError:
+            # Windows doesn't support add_signal_handler
+            pass
+
+    # Wait for either shutdown signal or service completion
+    wait_task = asyncio.create_task(stop_event.wait())
+
+    done, _ = await asyncio.wait(
+        {wait_task, service_task}, return_when=asyncio.FIRST_COMPLETED
+    )
+
+    # Handle shutdown
+    if wait_task in done:
+        logger.info("Stop signal received; shutting down daemon")
+        service_task.cancel()
+        try:
+            await service_task
+        except asyncio.CancelledError:
+            pass
+    else:
+        # Service completed (possibly with error)
+        wait_task.cancel()
+        await service_task  # Re-raise any exception
 
 
 if __name__ == "__main__":
-    logging.basicConfig(level=logging.INFO)
     asyncio.run(main())
diff --git a/nilai-models/src/nilai_models/lmstudio_announcer.py b/nilai-models/src/nilai_models/lmstudio_announcer.py
