feat: multimodal endpoint for image to text

.github/workflows/cicd.yml

@@ -137,7 +137,7 @@ jobs:
          sed -i 's/BRAVE_SEARCH_API=.*/BRAVE_SEARCH_API=${{ secrets.BRAVE_SEARCH_API }}/' .env
 
       - name: Compose docker-compose.yml
-        run: python3 ./scripts/docker-composer.py --dev -f docker/compose/docker-compose.llama-1b-gpu.ci.yml -o development-compose.yml
+        run: python3 ./scripts/docker-composer.py --dev -f docker/compose/docker-compose.llama-1b-cpu.ci.yml -o development-compose.yml
 
       - name: GPU stack versions (non-fatal)
         shell: bash

.gitignore

@@ -179,3 +179,4 @@ private_key.key.lock
 
 development-compose.yml
 production-compose.yml
+docker/compose/docker-compose.gemma-4b-gpu.ci.yml

docker-compose.dev.yml

@@ -2,10 +2,6 @@ services:
   caddy:
     env_file:
       - .env
-    ports:
-      - "80:80"
-      - "443:443"
-      - "443:443/udp"
     volumes:
       - ./caddy/Caddyfile:/etc/caddy/Caddyfile
   api:

docker/compose/docker-compose.gemma-27b-gpu.yml

@@ -0,0 +1,45 @@
+services:
+  gemma_27b_gpu:
+    image: nillion/nilai-vllm:latest
+    deploy:
+      resources:
+        reservations:
+          devices:
+            - driver: nvidia
+              count: all
+              capabilities: [gpu]
+    ipc: host
+    ulimits:
+      memlock: -1
+      stack: 67108864
+    env_file:
+      - .env
+    restart: unless-stopped
+    depends_on:
+      etcd:
+        condition: service_healthy
+    command: >
+      --model google/gemma-3-27b-it
+      --gpu-memory-utilization 0.79
+      --max-model-len 60000
+      --max-num-batched-tokens 8192
+      --dtype bfloat16
+      --kv-cache-dtype fp8
+      --uvicorn-log-level warning
+    environment:
+      - SVC_HOST=gemma_27b_gpu
+      - SVC_PORT=8000
+      - ETCD_HOST=etcd
+      - ETCD_PORT=2379
+      - TOOL_SUPPORT=false
+      - MULTIMODAL_SUPPORT=true
+    volumes:
+      - hugging_face_models:/root/.cache/huggingface
+    healthcheck:
+      test: ["CMD", "curl", "-f", "http://localhost:8000/health"]
+      interval: 30s
+      retries: 3
+      start_period: 60s
+      timeout: 10s
+volumes:
+  hugging_face_models:

docker/compose/docker-compose.gemma-4b-gpu.ci.yml

@@ -0,0 +1,47 @@
+services:
+  gemma_4b_gpu:
+    image: nillion/nilai-vllm:latest
+    container_name: nilai-gemma_4b_gpu
+    deploy:
+      resources:
+        reservations:
+          devices:
+            - driver: nvidia
+              count: 1
+              capabilities: [gpu]
+
+    ulimits:
+      memlock: -1
+      stack: 67108864
+    env_file:
+      - .env
+    restart: unless-stopped
+    depends_on:
+      etcd:
+        condition: service_healthy
+    command: >
+      --model google/gemma-3-4b-it
+      --max-model-len 30000
+      --max-num-batched-tokens 8192
+
+      --uvicorn-log-level warning
+    environment:
+      - SVC_HOST=gemma_4b_gpu
+      - SVC_PORT=8000
+      - ETCD_HOST=etcd
+      - ETCD_PORT=2379
+      - TOOL_SUPPORT=false
+      - MULTIMODAL_SUPPORT=true
+      - CUDA_LAUNCH_BLOCKING=1
+      - VLLM_ALLOW_LONG_MAX_MODEL_LEN=1
+      - PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
+    volumes:
+      - hugging_face_models:/root/.cache/huggingface
+    healthcheck:
+      test: ["CMD", "curl", "-f", "http://localhost:8000/health"]
+      interval: 30s
+      retries: 3
+      start_period: 60s
+      timeout: 10s
+volumes:
+  hugging_face_models:

docker/compose/docker-compose.llama-8b-gpu.yml

@@ -20,7 +20,7 @@ services:
         condition: service_healthy
     command: >
       --model meta-llama/Llama-3.1-8B-Instruct
-      --gpu-memory-utilization 0.21
+      --gpu-memory-utilization 0.20
       --max-model-len 10000
       --max-num-batched-tokens 10000
       --tensor-parallel-size 1

docker/compose/docker-compose.qwen-2b-gpu.ci.yml

@@ -0,0 +1,64 @@
+version: "3.8"
+
+services:
+  c:
+    image: nillion/nilai-vllm:latest
+    container_name: qwen2vl_2b_gpu
+    deploy:
+      resources:
+        reservations:
+          devices:
+            - driver: nvidia
+              count: 1
+              capabilities: [gpu]
+    ulimits:
+      memlock: -1
+      stack: 67108864
+    env_file:
+      - .env
+    restart: unless-stopped
+    depends_on:
+      etcd:
+        condition: service_healthy
+    command:
+      [
+        "--model", "Qwen/Qwen2-VL-2B-Instruct-AWQ",
+        "--model-impl", "vllm",
+        "--tensor-parallel-size", "1",
+        "--trust-remote-code",
+        "--quantization", "awq",
+
+        "--max-model-len", "1280",
+        "--max-num-batched-tokens", "1280",
+        "--max-num-seqs", "1",
+
+        "--gpu-memory-utilization", "0.75",
+        "--swap-space", "8",
+        "--uvicorn-log-level", "warning",
+
+        "--limit-mm-per-prompt", "{\"image\":1,\"video\":0}",
+        "--skip-mm-profiling",
+        "--enforce-eager"
+      ]
+
+    environment:
+      SVC_HOST: qwen2vl_2b_gpu
+      SVC_PORT: "8000"
+      ETCD_HOST: etcd
+      ETCD_PORT: "2379"
+      TOOL_SUPPORT: "true"
+      MULTIMODAL_SUPPORT: "true"
+      CUDA_LAUNCH_BLOCKING: "1"
+      VLLM_ALLOW_LONG_MAX_MODEL_LEN: "1"
+      PYTORCH_CUDA_ALLOC_CONF: "expandable_segments:True"
+    volumes:
+      - hugging_face_models:/root/.cache/huggingface
+    healthcheck:
+      test: ["CMD", "curl", "-f", "http://localhost:8000/health"]
+      interval: 30s
+      retries: 3
+      start_period: 60s
+      timeout: 10s
+
+volumes:
+  hugging_face_models:

nilai-api/src/nilai_api/handlers/nilrag.py

@@ -1,11 +1,11 @@
 import logging
+from typing import Union
 
 import nilrag
 
-from nilai_common import ChatRequest, Message
+from nilai_common import ChatRequest, MessageAdapter
 from fastapi import HTTPException, status
 from sentence_transformers import SentenceTransformer
-from typing import Union
 
 logger = logging.getLogger(__name__)
 
@@ -63,13 +63,9 @@ async def handle_nilrag(req: ChatRequest):
 
         # Get user query
         logger.debug("Extracting user query")
-        query = None
-        for message in req.messages:
-            if message.role == "user":
-                query = message.content
-                break
+        query = req.get_last_user_query()
 
-        if query is None:
+        if not query:
             raise HTTPException(status_code=400, detail="No user query found")
 
         # Get number of chunks to include
@@ -85,20 +81,25 @@ async def handle_nilrag(req: ChatRequest):
         relevant_context = f"\n\nRelevant Context:\n{formatted_results}"
 
         # Step 4: Update system message
-        for message in req.messages:
+        for message in req.adapted_messages:
             if message.role == "system":
-                if message.content is None:
+                content = message.content
+                if content is None:
                     raise HTTPException(
                         status_code=status.HTTP_400_BAD_REQUEST,
                         detail="system message is empty",
                     )
-                message.content += (
-                    relevant_context  # Append the context to the system message
-                )
+
+                if isinstance(content, str):
+                    message.content = content + relevant_context
+                elif isinstance(content, list):
+                    content.append({"type": "text", "text": relevant_context})
                 break
         else:
             # If no system message exists, add one
-            req.messages.insert(0, Message(role="system", content=relevant_context))
+            req.messages.insert(
+                0, MessageAdapter.new_message(role="system", content=relevant_context)
+            )
 
         logger.debug(f"System message updated with relevant context:\n {req.messages}")