add vision embedding

Dockerfile

@@ -24,7 +24,7 @@ RUN pip install --no-cache-dir "youtube-transcript-api>=1.2.4"
 
 COPY --chown=appuser:appuser . .
 
-RUN mkdir -p /app/models && python3 -c "from huggingface_hub import snapshot_download; snapshot_download(repo_id='ibm-granite/granite-embedding-small-english-r2', local_dir='/app/models/bge-384'); snapshot_download(repo_id='nomic-ai/modernbert-embed-base', local_dir='/app/models/bge-768'); snapshot_download(repo_id='lightonai/modernbert-embed-large', local_dir='/app/models/bge-1024')" && chown -R appuser:appuser /app/models
+RUN mkdir -p /app/models && python3 -c "from huggingface_hub import snapshot_download; snapshot_download(repo_id='ibm-granite/granite-embedding-small-english-r2', local_dir='/app/models/bge-384'); snapshot_download(repo_id='nomic-ai/modernbert-embed-base', local_dir='/app/models/bge-768'); snapshot_download(repo_id='lightonai/modernbert-embed-large', local_dir='/app/models/bge-1024'); snapshot_download(repo_id='nomic-ai/nomic-embed-vision-v1.5', local_dir='/app/models/vision')" && chown -R appuser:appuser /app/models
 
 RUN mkdir -p /app/logs && \
     chown -R appuser:appuser /app/logs

app/api/server.py

@@ -40,7 +40,9 @@ async def lifespan(app: FastAPI):
     _logger.info("Initializing embedding service (loading all models)...")
     loop = asyncio.get_running_loop()
     await loop.run_in_executor(None, _embedding_service.load_all_models)
-    _logger.info("Embedding service initialized with dims: %s", _embedding_service.loaded_dimensions)
+    await loop.run_in_executor(None, _embedding_service.load_vision_model)
+    _logger.info("Embedding service initialized with dims: %s, vision=%s",
+                 _embedding_service.loaded_dimensions, _embedding_service._vision_loaded)
 
     asyncio.create_task(_self_ping())
     yield

app/api/v1/embeddings.py

@@ -2,20 +2,54 @@ from __future__ import annotations
 
 import asyncio
 import concurrent.futures
+import io
 import os
 import time
+from typing import Annotated, List, Optional
 
-from fastapi import APIRouter, Depends, HTTPException
+import httpx
+from fastapi import APIRouter, Depends, File, Form, HTTPException, UploadFile, status
+from PIL import Image
 
 from app.api.deps import require_auth, get_embeddings_service
+from app.config import get_settings
 from app.core.logger import get_logger
-from app.models.schemas import EmbeddingItem, EmbeddingRequest, EmbeddingResponse
+from app.models.schemas import EmbeddingItem, EmbeddingRequest, EmbeddingResponse, VisionUrlRequest
 from app.services.embeddings_service import EmbeddingService
 
 router = APIRouter()
 _logger = get_logger(__name__)
+_settings = get_settings()
 _MAX_WORKERS = min(32, (os.cpu_count() or 1) + 4)
 _thread_pool = concurrent.futures.ThreadPoolExecutor(max_workers=_MAX_WORKERS)
+_MAX_VISION_ITEMS = 5
+_MAX_IMAGE_BYTES = 15 * 1024 * 1024
+
+
+def _validate_image(raw: bytes, source: str) -> Image.Image:
+    if len(raw) > _MAX_IMAGE_BYTES:
+        raise ValueError(f"Image {source} exceeds 15 MB limit")
+    try:
+        img = Image.open(io.BytesIO(raw))
+        img.load()
+        if img.mode != "RGB":
+            img = img.convert("RGB")
+        return img
+    except Exception as exc:
+        raise ValueError(f"Invalid image {source}: {exc}")
+
+
+async def _download_image(url: str) -> bytes:
+    try:
+        async with httpx.AsyncClient(timeout=30.0, follow_redirects=True) as client:
+            resp = await client.get(url)
+            resp.raise_for_status()
+            ctype = resp.headers.get("content-type", "")
+            if not ctype.startswith("image/"):
+                raise ValueError(f"URL {url} returned non-image Content-Type: {ctype}")
+            return resp.content
+    except httpx.HTTPError as exc:
+        raise ValueError(f"Failed to download {url}: {exc}")
 
 
 @router.post(
@@ -59,11 +93,7 @@ async def create_embeddings(
             failed_count=len(body.content),
             error_message=str(exc),
             results=[
-                EmbeddingItem(
-                    success=False,
-                    time_ms=0,
-                    error_message=str(exc),
-                )
+                EmbeddingItem(success=False, time_ms=0, error_message=str(exc))
                 for _ in body.content
             ],
         )
@@ -89,3 +119,184 @@ async def create_embeddings(
         failed_count=0,
         results=results,
     )
+
+
+@router.post(
+    "/embeddings/vision/file",
+    response_model=EmbeddingResponse,
+    summary="Generate embeddings from uploaded images",
+)
+async def create_vision_embeddings_file(
+    files: Annotated[List[UploadFile], File(description="Image files to embed (max 5)")],
+    token: str = Depends(require_auth),
+    embedding_service: EmbeddingService = Depends(get_embeddings_service),
+) -> EmbeddingResponse:
+    if not files:
+        raise HTTPException(status_code=400, detail={"success": False, "message": "No files provided."})
+    if len(files) > _MAX_VISION_ITEMS:
+        raise HTTPException(status_code=400, detail={"success": False, "message": f"Maximum {_MAX_VISION_ITEMS} images per request."})
+
+    if not embedding_service._vision_loaded:
+        raise HTTPException(status_code=503, detail={"success": False, "message": "Vision model not loaded."})
+
+    _logger.info("Vision embedding file request: files=%s", len(files))
+
+    dim = embedding_service.vision_dimension
+    start = time.perf_counter()
+    images: List[Image.Image] = []
+    item_results: List[EmbeddingItem] = []
+
+    for f in files:
+        t0 = time.perf_counter()
+        try:
+            raw = await f.read()
+            img = await asyncio.get_running_loop().run_in_executor(_thread_pool, _validate_image, raw, f.filename or "unknown")
+            images.append(img)
+        except Exception as exc:
+            elapsed = (time.perf_counter() - t0) * 1000
+            item_results.append(EmbeddingItem(
+                success=False,
+                time_ms=round(elapsed, 3),
+                error_message=str(exc),
+            ))
+
+    if not images:
+        total_ms = (time.perf_counter() - start) * 1000
+        return EmbeddingResponse(
+            success=False,
+            time_ms=round(total_ms, 3),
+            success_count=0,
+            failed_count=len(item_results),
+            error_message="No valid images could be processed.",
+            results=item_results,
+        )
+
+    try:
+        loop = asyncio.get_running_loop()
+        vectors = await loop.run_in_executor(
+            _thread_pool,
+            embedding_service.generate_image_embedding,
+            images,
+        )
+    except Exception as exc:
+        elapsed = (time.perf_counter() - start) * 1000
+        _logger.error("Vision embedding error: %s", exc)
+        for _ in range(len(images) - len(item_results)):
+            item_results.append(EmbeddingItem(success=False, time_ms=0, error_message=str(exc)))
+        return EmbeddingResponse(
+            success=False,
+            time_ms=round(elapsed, 3),
+            success_count=0,
+            failed_count=len(item_results),
+            error_message=str(exc),
+            results=item_results,
+        )
+
+    total_ms = (time.perf_counter() - start) * 1000
+    for i, vec in enumerate(vectors):
+        item_results.append(EmbeddingItem(
+            success=True,
+            time_ms=round(total_ms / len(vectors), 3),
+            embeddings=vec,
+            dimension=dim,
+        ))
+
+    success_count = sum(1 for r in item_results if r.success)
+    failed_count = len(item_results) - success_count
+    _logger.info("Vision embedding success: items=%s, success=%s, failed=%s, total_ms=%s",
+                 len(item_results), success_count, failed_count, round(total_ms, 3))
+    return EmbeddingResponse(
+        success=failed_count == 0,
+        time_ms=round(total_ms, 3),
+        success_count=success_count,
+        failed_count=failed_count,
+        results=item_results,
+    )
+
+
+@router.post(
+    "/embeddings/vision/url",
+    response_model=EmbeddingResponse,
+    summary="Generate embeddings from image URLs",
+)
+async def create_vision_embeddings_url(
+    body: VisionUrlRequest,
+    token: str = Depends(require_auth),
+    embedding_service: EmbeddingService = Depends(get_embeddings_service),
+) -> EmbeddingResponse:
+    if not embedding_service._vision_loaded:
+        raise HTTPException(status_code=503, detail={"success": False, "message": "Vision model not loaded."})
+
+    _logger.info("Vision embedding URL request: urls=%s", len(body.urls))
+
+    dim = embedding_service.vision_dimension
+    start = time.perf_counter()
+    images: List[Image.Image] = []
+    item_results: List[EmbeddingItem] = []
+
+    for url in body.urls:
+        t0 = time.perf_counter()
+        try:
+            raw = await _download_image(url)
+            img = await asyncio.get_running_loop().run_in_executor(_thread_pool, _validate_image, raw, url)
+            images.append(img)
+        except Exception as exc:
+            elapsed = (time.perf_counter() - t0) * 1000
+            item_results.append(EmbeddingItem(
+                success=False,
+                time_ms=round(elapsed, 3),
+                error_message=str(exc),
+            ))
+
+    if not images:
+        total_ms = (time.perf_counter() - start) * 1000
+        return EmbeddingResponse(
+            success=False,
+            time_ms=round(total_ms, 3),
+            success_count=0,
+            failed_count=len(item_results),
+            error_message="No valid images could be downloaded.",
+            results=item_results,
+        )
+
+    try:
+        loop = asyncio.get_running_loop()
+        vectors = await loop.run_in_executor(
+            _thread_pool,
+            embedding_service.generate_image_embedding,
+            images,
+        )
+    except Exception as exc:
+        elapsed = (time.perf_counter() - start) * 1000
+        _logger.error("Vision embedding error: %s", exc)
+        for _ in range(len(images) - len(item_results)):
+            item_results.append(EmbeddingItem(success=False, time_ms=0, error_message=str(exc)))
+        return EmbeddingResponse(
+            success=False,
+            time_ms=round(elapsed, 3),
+            success_count=0,
+            failed_count=len(item_results),
+            error_message=str(exc),
+            results=item_results,
+        )
+
+    total_ms = (time.perf_counter() - start) * 1000
+    for i, vec in enumerate(vectors):
+        item_results.append(EmbeddingItem(
+            success=True,
+            time_ms=round(total_ms / len(vectors), 3),
+            embeddings=vec,
+            dimension=dim,
+        ))
+
+    success_count = sum(1 for r in item_results if r.success)
+    failed_count = len(item_results) - success_count
+    _logger.info("Vision embedding success: items=%s, success=%s, failed=%s, total_ms=%s",
+                 len(item_results), success_count, failed_count, round(total_ms, 3))
+    return EmbeddingResponse(
+        success=failed_count == 0,
+        time_ms=round(total_ms, 3),
+        success_count=success_count,
+        failed_count=failed_count,
+        results=item_results,
+    )

app/models/__init__.py

@@ -15,6 +15,7 @@ from app.models.schemas import (
     SpacyLabelsResponse,
     SupportedFormatsResponse,
     UrlRequest,
+    VisionUrlRequest,
 )
 
 __all__ = [
@@ -33,4 +34,5 @@ __all__ = [
     "InfoResponse",
     "SupportedFormatsResponse",
     "SpacyLabelsResponse",
+    "VisionUrlRequest",
 ]

app/models/schemas.py

@@ -254,3 +254,15 @@ class EmbeddingResponse(BaseModel):
     failed_count: int
     error_message: Optional[str] = None
     results: List[EmbeddingItem]
+
+
+class VisionUrlRequest(BaseModel):
+    urls: List[str] = Field(..., min_length=1, max_length=5, description="Array of image URLs to embed (max 5)")
+
+    @field_validator("urls")
+    @classmethod
+    def validate_urls(cls, v: List[str]) -> List[str]:
+        for url in v:
+            if not url.startswith(("http://", "https://")):
+                raise ValueError(f"Invalid URL scheme: {url}")
+        return v

app/services/embeddings_service.py

@@ -5,6 +5,7 @@ import os
 from typing import Dict, List, Optional
 
 import numpy as np
+from PIL import Image
 from sentence_transformers import SentenceTransformer
 
 _logger = logging.getLogger(__name__)
@@ -15,6 +16,9 @@ _MODEL_MAP: Dict[int, str] = {
     1024: "lightonai/modernbert-embed-large",
 }
 
+_VISION_MODEL_NAME = "nomic-ai/nomic-embed-vision-v1.5"
+_VISION_DIMENSION = 768
+
 
 class EmbeddingService:
     def __init__(self, models_dir: Optional[str] = None) -> None:
@@ -29,6 +33,8 @@ class EmbeddingService:
             self._device = "cpu"
 
         self._loaded_dimensions: List[int] = []
+        self._vision_model: Optional[SentenceTransformer] = None
+        self._vision_loaded = False
 
     def load_model(self, dimension: int) -> None:
         if dimension in self._models:
@@ -53,6 +59,21 @@ class EmbeddingService:
         for dim in _MODEL_MAP:
             self.load_model(dim)
 
+    def load_vision_model(self) -> None:
+        if self._vision_loaded:
+            return
+        local_path = os.path.join(self._models_dir, "vision")
+        source = local_path if os.path.isdir(local_path) else _VISION_MODEL_NAME
+        _logger.info("Loading vision embedding model from %s", source)
+        self._vision_model = SentenceTransformer(
+            source,
+            device=self._device,
+            trust_remote_code=True,
+        )
+        self._vision_model.eval()
+        self._vision_loaded = True
+        _logger.info("Loaded vision embedding model (device=%s)", self._device)
+
     def generate_embedding(self, text: List[str], dimension: int) -> List[List[float]]:
         if dimension not in self._models:
             raise ValueError(f"Model for dimension {dimension} not loaded")
@@ -65,9 +86,23 @@ class EmbeddingService:
         )
         return result.tolist()
 
+    def generate_image_embedding(self, images: List[Image.Image]) -> List[List[float]]:
+        if not self._vision_loaded or self._vision_model is None:
+            raise ValueError("Vision model not loaded")
+        result: np.ndarray = self._vision_model.encode(
+            images,
+            convert_to_numpy=True,
+            show_progress_bar=False,
+        )
+        return result.tolist()
+
     @property
     def loaded_dimensions(self) -> List[int]:
         return list(self._loaded_dimensions)
 
     def is_loaded(self, dimension: int) -> bool:
         return dimension in self._models
+
+    @property
+    def vision_dimension(self) -> int:
+        return _VISION_DIMENSION