feat: multimodal endpoint for image to text

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

@@ -0,0 +1,46 @@
+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.95
+      --max-model-len 60000
+      --max-num-batched-tokens 60000
+      --tensor-parallel-size 1
+      --enable-auto-tool-choice
+      --tool-call-parser llama3_json
+      --uvicorn-log-level warning
+    environment:
+      - SVC_HOST=gemma_27b_gpu
+      - SVC_PORT=8000
+      - ETCD_HOST=etcd
+      - ETCD_PORT=2379
+      - TOOL_SUPPORT=true
+      - 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,51 @@
+services:
+  gemma_4b_gpu:
+    image: nillion/nilai-vllm:latest
+    container_name: nilai-gemma_4b_gpu
+    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-4b-it
+      --gpu-memory-utilization 0.7
+      --max-model-len 8192
+      --max-num-batched-tokens 8192
+      --tensor-parallel-size 1
+      --enable-auto-tool-choice
+      --tool-call-parser llama3_json
+      --uvicorn-log-level warning
+      --dtype half
+    environment:
+      - SVC_HOST=gemma_4b_gpu
+      - SVC_PORT=8000
+      - ETCD_HOST=etcd
+      - ETCD_PORT=2379
+      - TOOL_SUPPORT=true
+      - MULTIMODAL_SUPPORT=true
+      - VLLM_USE_V1=1
+      - VLLM_ALLOW_LONG_MAX_MODEL_LEN=1
+      - CUDA_LAUNCH_BLOCKING=1
+    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/image_support.py

@@ -0,0 +1,92 @@
+from dataclasses import dataclass
+from typing import List, Optional, Any
+from fastapi import HTTPException
+from nilai_common import Message
+
+
+@dataclass(frozen=True)
+class MultimodalCheck:
+    has_multimodal: bool
+    error: Optional[str] = None
+
+
+def _extract_url(image_url_field: Any) -> Optional[str]:
+    """
+    Support both object-with-attr and dict-like shapes.
+    Returns the URL string or None.
+    """
+    if image_url_field is None:
+        return None
+
+    url = getattr(image_url_field, "url", None)
+    if url is not None:
+        return url
+    if isinstance(image_url_field, dict):
+        return image_url_field.get("url")
+    return None
+
+
+def multimodal_check(messages: List[Message]) -> MultimodalCheck:
+    """
+    Single-pass check:
+      - detect if any part is type=='image_url'
+      - validate that image_url.url exists and is a base64 data URL
+    Returns:
+      MultimodalCheck(has_multimodal: bool, error: Optional[str])
+    """
+    has_mm = False
+
+    for m in messages:
+        content = getattr(m, "content", None) or []
+        for item in content:
+            if getattr(item, "type", None) == "image_url":
+                has_mm = True
+                iu = getattr(item, "image_url", None)
+                url = _extract_url(iu)
+                if not url:
+                    return MultimodalCheck(
+                        True, "image_url.url is required for image_url parts"
+                    )
+                if not (url.startswith("data:image/") and ";base64," in url):
+                    return MultimodalCheck(
+                        True,
+                        "Only base64 data URLs are allowed for images (data:image/...;base64,...)",
+                    )
+
+    return MultimodalCheck(has_mm, None)
+
+
+def has_multimodal_content(
+    messages: List[Message], precomputed: Optional[MultimodalCheck] = None
+) -> bool:
+    """
+    Check if any message contains multimodal content (image_url parts).
+
+    Args:
+        messages: List of messages to check
+        precomputed: Optional precomputed result from multimodal_check() to avoid re-iterating
+
+    Returns:
+        True if any message contains image_url parts, False otherwise
+    """
+    res = precomputed or multimodal_check(messages)
+    return res.has_multimodal
+
+
+def validate_multimodal_content(
+    messages: List[Message], precomputed: Optional[MultimodalCheck] = None
+) -> None:
+    """
+    Validate that multimodal content (image_url parts) follows the required format.
+
+    Args:
+        messages: List of messages to validate
+        precomputed: Optional precomputed result from multimodal_check() to avoid re-iterating
+
+    Raises:
+        HTTPException(400): When image_url parts don't have required URL or use invalid format
+                           (only base64 data URLs are allowed: data:image/...;base64,...)
+    """
+    res = precomputed or multimodal_check(messages)
+    if res.error:
+        raise HTTPException(status_code=400, detail=res.error)

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

@@ -11,6 +11,7 @@
     Source,
     WebSearchEnhancedMessages,
     WebSearchContext,
+    MessageContentItem,
 )
 from nilai_common import Message
 
@@ -152,22 +153,29 @@ async def perform_web_search_async(query: str) -> WebSearchContext:
 async def enhance_messages_with_web_search(
     messages: List[Message], query: str
 ) -> WebSearchEnhancedMessages:
-    """Enhance a list of messages with web search context.
-
-    Args:
-        messages: List of conversation messages to enhance
-        query: Search query to retrieve web search results for
-
-    Returns:
-        WebSearchEnhancedMessages containing the original messages with web search
-        context prepended as a system message, along with source information
-    """
     ctx = await perform_web_search_async(query)
-    enhanced = [Message(role="system", content=ctx.prompt)] + messages
     query_source = Source(source="search_query", content=query)
+
+    if not messages or messages[-1].role != "user":
+        return WebSearchEnhancedMessages(
+            messages=messages, sources=[query_source] + ctx.sources
+        )
+
+    web_search_context = f"\n\nWeb search results:\n{ctx.prompt}"
+
+    last = messages[-1]
+    items = (
+        [MessageContentItem(type="text", text=last.content)]
+        if isinstance(last.content, str)
+        else list(last.content)
+    )
+    items.append(MessageContentItem(type="text", text=web_search_context))
+
+    enhanced_messages = list(messages)
+    enhanced_messages[-1] = Message(role="user", content=items)
+
     return WebSearchEnhancedMessages(
-        messages=enhanced,
-        sources=[query_source] + ctx.sources,
+        messages=enhanced_messages, sources=[query_source] + ctx.sources
     )
 
 

nilai-api/src/nilai_api/routers/private.py

@@ -6,6 +6,7 @@
 from nilai_api.attestation import get_attestation_report
 from nilai_api.handlers.nilrag import handle_nilrag
 from nilai_api.handlers.web_search import handle_web_search
+from nilai_api.handlers.image_support import multimodal_check
 
 from fastapi import APIRouter, Body, Depends, HTTPException, status, Request
 from fastapi.responses import StreamingResponse
@@ -211,6 +212,17 @@ async def chat_completion(
             status_code=400,
             detail="Model does not support tool usage, remove tools from request",
         )
+
+    multimodal_result = multimodal_check(req.messages)
+    if multimodal_result.has_multimodal:
+        if not endpoint.metadata.multimodal_support:
+            raise HTTPException(
+                status_code=400,
+                detail="Model does not support multimodal content, remove image inputs from request",
+            )
+        if multimodal_result.error:
+            raise HTTPException(status_code=400, detail=multimodal_result.error)
+
     model_url = endpoint.url + "/v1/"
 
     logger.info(

nilai-models/src/nilai_models/daemon.py

@@ -28,18 +28,26 @@ async def get_metadata(num_retries=30):
                 response.raise_for_status()
                 response_data = response.json()
                 model_name = response_data["data"][0]["id"]
-                return ModelMetadata(
+
+                supported_features = ["chat_completion"]
+                if SETTINGS.multimodal_support:
+                    supported_features.append("multimodal")
+
+                metadata = ModelMetadata(
                     id=model_name,  # Unique identifier
                     name=model_name,  # Human-readable name
                     version="1.0",  # Model version
                     description="",
                     author="",  # Model creators
                     license="Apache 2.0",  # Usage license
                     source=f"https://huggingface.co/{model_name}",  # Model source
-                    supported_features=["chat_completion"],  # Capabilities
+                    supported_features=supported_features,  # Capabilities
                     tool_support=SETTINGS.tool_support,  # Tool support
+                    multimodal_support=SETTINGS.multimodal_support,  # Multimodal support
                 )
 
+                return metadata
+
         except Exception as e:
             if not url:
                 logger.warning(f"Failed to build url: {e}")

packages/nilai-common/src/nilai_common/api_model.py

@@ -1,15 +1,27 @@
 import uuid
 
-from typing import Annotated, Iterable, List, Literal, Optional
+from typing import Annotated, Iterable, List, Literal, Optional, Union
 
-from openai.types.chat import ChatCompletion, ChatCompletionMessage
-from openai.types.chat.chat_completion import Choice as OpenaAIChoice
+from openai.types.chat import ChatCompletion
 from openai.types.chat import ChatCompletionToolParam
+from openai.types.chat.chat_completion import Choice as OpenaAIChoice
 from pydantic import BaseModel, Field
 
 
-class Message(ChatCompletionMessage):
-    role: Literal["system", "user", "assistant", "tool"]  # type: ignore
+class ImageUrl(BaseModel):
+    url: str
+    detail: Optional[str] = "auto"
+
+
+class MessageContentItem(BaseModel):
+    type: Literal["text", "image_url"]
+    text: Optional[str] = None
+    image_url: Optional[ImageUrl] = None
+
+
+class Message(BaseModel):
+    role: Literal["system", "user", "assistant", "tool"]
+    content: Union[str, List[MessageContentItem]]
 
 
 class Choice(OpenaAIChoice):
@@ -71,6 +83,7 @@ class ModelMetadata(BaseModel):
     source: str
     supported_features: List[str]
     tool_support: bool
+    multimodal_support: bool = False
 
 
 class ModelEndpoint(BaseModel):

packages/nilai-common/src/nilai_common/config.py

@@ -8,6 +8,7 @@ class HostSettings(BaseModel):
     etcd_host: str = "localhost"
     etcd_port: int = 2379
     tool_support: bool = False
+    multimodal_support: bool = False
     gunicorn_workers: int = 10
     attestation_host: str = "localhost"
     attestation_port: int = 8081
@@ -19,6 +20,7 @@ class HostSettings(BaseModel):
     etcd_host=str(os.getenv("ETCD_HOST", "localhost")),
     etcd_port=int(os.getenv("ETCD_PORT", 2379)),
     tool_support=bool(os.getenv("TOOL_SUPPORT", False)),
+    multimodal_support=bool(os.getenv("MULTIMODAL_SUPPORT", False)),
     gunicorn_workers=int(os.getenv("NILAI_GUNICORN_WORKERS", 10)),
     attestation_host=str(os.getenv("ATTESTATION_HOST", "localhost")),
     attestation_port=int(os.getenv("ATTESTATION_PORT", 8081)),

tests/e2e/config.py

@@ -34,6 +34,7 @@ def api_key_getter():
     ],
     "ci": [
         "meta-llama/Llama-3.2-1B-Instruct",
+        "google/gemma-3-4b-it",
     ],
 }