Merge pull request #155 from scub-france/feature/embedding-service

feat: embedding microservice (sentence-transformers)

.env.example

@@ -40,5 +40,11 @@
 # OpenSearch URL (used by docker-compose.dev.yml, auto-set to service name)
 # OPENSEARCH_URL=http://opensearch:9200
 
+# Embedding service URL (used by docker-compose.dev.yml, auto-set to service name)
+# EMBEDDING_URL=http://embedding:8001
+
+# Embedding model (default: all-MiniLM-L6-v2, used by the embedding service)
+# EMBEDDING_MODEL=all-MiniLM-L6-v2
+
 # Embedding vector dimension (default: 384 for Granite Embedding 30M / all-MiniLM-L6-v2)
 # EMBEDDING_DIMENSION=384

CHANGELOG.md

@@ -14,6 +14,8 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/), and this
 - Vector index metadata schema: `IndexedChunk` domain model, OpenSearch mapping builder, configurable embedding dimension
 - `VectorStore` port (Protocol): `ensure_index`, `index_chunks`, `search_similar`, `get_chunks`, `delete_document`
 - OpenSearch adapter (`OpenSearchStore`): kNN vector search, full-text search, bulk indexing, document CRUD
+- Embedding microservice (`embedding-service/`): sentence-transformers REST API with batch processing and Dockerfile
+- `EmbeddingService` port and `EmbeddingClient` HTTP adapter for remote embedding generation
 
 ### Fixed
 

docker-compose.dev.yml

@@ -38,6 +38,25 @@ services:
       opensearch:
         condition: service_healthy
 
+  # --- Embedding service (sentence-transformers) ---
+  embedding:
+    build:
+      context: ./embedding-service
+    ports:
+      - "8001:8001"
+    environment:
+      EMBEDDING_MODEL: ${EMBEDDING_MODEL:-all-MiniLM-L6-v2}
+      EMBEDDING_BATCH_SIZE: ${EMBEDDING_BATCH_SIZE:-64}
+    healthcheck:
+      test: ["CMD-SHELL", "curl -sf http://localhost:8001/health || exit 1"]
+      interval: 10s
+      timeout: 5s
+      retries: 10
+    deploy:
+      resources:
+        limits:
+          memory: 2g
+
   # --- Backend (FastAPI with hot-reload) ---
   document-parser:
     build:
@@ -57,10 +76,13 @@ services:
       MAX_FILE_SIZE_MB: ${MAX_FILE_SIZE_MB:-50}
       BATCH_PAGE_SIZE: ${BATCH_PAGE_SIZE:-0}
       OPENSEARCH_URL: http://opensearch:9200
+      EMBEDDING_URL: http://embedding:8001
     command: ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "8000", "--reload"]
     depends_on:
       opensearch:
         condition: service_healthy
+      embedding:
+        condition: service_healthy
     deploy:
       resources:
         limits:

document-parser/domain/ports.py

@@ -82,6 +82,18 @@ class AnalysisRepository(Protocol):
     async def delete_by_document(self, document_id: str) -> int: ...
 
 
+@runtime_checkable
+class EmbeddingService(Protocol):
+    """Port for text-to-vector embedding.
+
+    Implementations may call a local model, a remote microservice, etc.
+    """
+
+    async def embed(self, texts: list[str]) -> list[list[float]]:
+        """Generate embedding vectors for a batch of texts."""
+        ...
+
+
 @runtime_checkable
 class VectorStore(Protocol):
     """Port for vector storage and retrieval.

document-parser/infra/embedding_client.py

@@ -0,0 +1,51 @@
+"""HTTP client adapter for the embedding microservice.
+
+Satisfies the ``EmbeddingService`` Protocol defined in ``domain.ports``.
+Calls the embedding-service REST API (POST /embed).
+"""
+
+from __future__ import annotations
+
+import logging
+
+import httpx
+
+logger = logging.getLogger(__name__)
+
+# Maximum texts per request to avoid payload / memory issues on the server.
+_MAX_BATCH = 256
+
+
+class EmbeddingClient:
+    """Remote embedding adapter backed by the embedding-service microservice.
+
+    Args:
+        base_url: Embedding service URL (e.g. ``http://localhost:8001``).
+        timeout: HTTP request timeout in seconds.
+    """
+
+    def __init__(self, base_url: str, *, timeout: float = 120.0) -> None:
+        self._base_url = base_url.rstrip("/")
+        self._timeout = timeout
+
+    async def embed(self, texts: list[str]) -> list[list[float]]:
+        """Generate embeddings by calling the remote service.
+
+        Automatically splits large batches into sub-batches of ``_MAX_BATCH``.
+        """
+        if not texts:
+            return []
+
+        all_embeddings: list[list[float]] = []
+        async with httpx.AsyncClient(timeout=self._timeout) as client:
+            for start in range(0, len(texts), _MAX_BATCH):
+                batch = texts[start : start + _MAX_BATCH]
+                resp = await client.post(
+                    f"{self._base_url}/embed",
+                    json={"texts": batch},
+                )
+                resp.raise_for_status()
+                data = resp.json()
+                all_embeddings.extend(data["embeddings"])
+
+        return all_embeddings

document-parser/infra/settings.py

@@ -24,6 +24,7 @@ class Settings:
     rate_limit_rpm: int = 100  # requests per minute per IP (0 = disabled)
     batch_page_size: int = 0  # 0 = disabled, > 0 = pages per batch
     opensearch_url: str = ""  # empty = disabled
+    embedding_url: str = ""  # empty = disabled (e.g. http://localhost:8001)
     embedding_dimension: int = 384  # Granite Embedding 30M / all-MiniLM-L6-v2
     upload_dir: str = "./uploads"
     db_path: str = "./data/docling_studio.db"
@@ -95,6 +96,7 @@ class Settings:
             rate_limit_rpm=int(os.environ.get("RATE_LIMIT_RPM", "100")),
             batch_page_size=int(os.environ.get("BATCH_PAGE_SIZE", "0")),
             opensearch_url=os.environ.get("OPENSEARCH_URL", ""),
+            embedding_url=os.environ.get("EMBEDDING_URL", ""),
             embedding_dimension=int(os.environ.get("EMBEDDING_DIMENSION", "384")),
             upload_dir=os.environ.get("UPLOAD_DIR", "./uploads"),
             db_path=os.environ.get("DB_PATH", "./data/docling_studio.db"),

document-parser/tests/test_embedding_client.py

@@ -0,0 +1,112 @@
+"""Tests for the embedding client adapter (infra.embedding_client).
+
+Mock httpx to validate adapter logic without running the embedding service.
+"""
+
+from __future__ import annotations
+
+from unittest.mock import AsyncMock, MagicMock, patch
+
+from domain.ports import EmbeddingService
+from infra.embedding_client import _MAX_BATCH, EmbeddingClient
+
+# -- Protocol satisfaction -----------------------------------------------------
+
+
+class TestProtocolSatisfaction:
+    def test_satisfies_embedding_service_protocol(self) -> None:
+        client = EmbeddingClient("http://localhost:8001")
+        assert isinstance(client, EmbeddingService)
+
+
+# -- embed ---------------------------------------------------------------------
+
+
+class TestEmbed:
+    async def test_returns_empty_for_empty_input(self) -> None:
+        client = EmbeddingClient("http://localhost:8001")
+        result = await client.embed([])
+        assert result == []
+
+    async def test_calls_service_and_returns_embeddings(self) -> None:
+        client = EmbeddingClient("http://localhost:8001")
+        mock_response = MagicMock()
+        mock_response.raise_for_status = MagicMock()
+        mock_response.json.return_value = {
+            "embeddings": [[0.1, 0.2], [0.3, 0.4]],
+            "model": "all-MiniLM-L6-v2",
+            "dimension": 2,
+        }
+
+        mock_http_client = AsyncMock()
+        mock_http_client.post.return_value = mock_response
+        mock_http_client.__aenter__ = AsyncMock(return_value=mock_http_client)
+        mock_http_client.__aexit__ = AsyncMock(return_value=False)
+
+        with patch("infra.embedding_client.httpx.AsyncClient", return_value=mock_http_client):
+            result = await client.embed(["hello", "world"])
+
+        assert result == [[0.1, 0.2], [0.3, 0.4]]
+        mock_http_client.post.assert_awaited_once_with(
+            "http://localhost:8001/embed",
+            json={"texts": ["hello", "world"]},
+        )
+
+    async def test_strips_trailing_slash_from_base_url(self) -> None:
+        client = EmbeddingClient("http://localhost:8001/")
+        mock_response = MagicMock()
+        mock_response.raise_for_status = MagicMock()
+        mock_response.json.return_value = {"embeddings": [[0.1]], "model": "m", "dimension": 1}
+
+        mock_http_client = AsyncMock()
+        mock_http_client.post.return_value = mock_response
+        mock_http_client.__aenter__ = AsyncMock(return_value=mock_http_client)
+        mock_http_client.__aexit__ = AsyncMock(return_value=False)
+
+        with patch("infra.embedding_client.httpx.AsyncClient", return_value=mock_http_client):
+            await client.embed(["test"])
+
+        mock_http_client.post.assert_awaited_once_with(
+            "http://localhost:8001/embed",
+            json={"texts": ["test"]},
+        )
+
+    async def test_splits_large_batches(self) -> None:
+        client = EmbeddingClient("http://localhost:8001")
+        texts = [f"text_{i}" for i in range(_MAX_BATCH + 10)]
+
+        call_count = 0
+
+        def make_response(batch_size: int) -> MagicMock:
+            resp = MagicMock()
+            resp.raise_for_status = MagicMock()
+            resp.json.return_value = {
+                "embeddings": [[0.1]] * batch_size,
+                "model": "m",
+                "dimension": 1,
+            }
+            return resp
+
+        async def mock_post(url: str, json: dict) -> MagicMock:
+            nonlocal call_count
+            call_count += 1
+            return make_response(len(json["texts"]))
+
+        mock_http_client = AsyncMock()
+        mock_http_client.post = mock_post
+        mock_http_client.__aenter__ = AsyncMock(return_value=mock_http_client)
+        mock_http_client.__aexit__ = AsyncMock(return_value=False)
+
+        with patch("infra.embedding_client.httpx.AsyncClient", return_value=mock_http_client):
+            result = await client.embed(texts)
+
+        assert len(result) == _MAX_BATCH + 10
+        assert call_count == 2  # _MAX_BATCH + 10 remaining
+
+
+# -- max batch constant --------------------------------------------------------
+
+
+class TestMaxBatch:
+    def test_max_batch_is_256(self) -> None:
+        assert _MAX_BATCH == 256

embedding-service/Dockerfile

@@ -0,0 +1,17 @@
+FROM python:3.12-slim
+
+WORKDIR /app
+
+# Install dependencies first (cache layer)
+COPY requirements.txt .
+RUN pip install --no-cache-dir -r requirements.txt
+
+# Pre-download default model into the image
+ARG EMBEDDING_MODEL=all-MiniLM-L6-v2
+RUN python -c "from sentence_transformers import SentenceTransformer; SentenceTransformer('${EMBEDDING_MODEL}')"
+
+COPY main.py .
+
+EXPOSE 8001
+
+CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "8001"]

embedding-service/main.py

@@ -0,0 +1,89 @@
+"""Embedding microservice — exposes sentence-transformers models via REST API.
+
+POST /embed  {"texts": ["...", "..."]}  →  {"embeddings": [[...], [...]], "model": "...", "dimension": N}
+GET  /health                            →  {"status": "ok", "model": "...", "dimension": N}
+"""
+
+from __future__ import annotations
+
+import logging
+import os
+import time
+
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel, Field
+from sentence_transformers import SentenceTransformer
+
+logger = logging.getLogger(__name__)
+
+MODEL_NAME = os.environ.get("EMBEDDING_MODEL", "all-MiniLM-L6-v2")
+BATCH_SIZE = int(os.environ.get("EMBEDDING_BATCH_SIZE", "64"))
+
+app = FastAPI(title="Docling Studio — Embedding Service", version="0.4.0")
+
+# Load model at startup (downloaded / cached in HF cache dir)
+model: SentenceTransformer | None = None
+
+
+@app.on_event("startup")
+async def _load_model() -> None:
+    global model  # noqa: PLW0603
+    logger.info("Loading sentence-transformers model '%s' …", MODEL_NAME)
+    t0 = time.monotonic()
+    model = SentenceTransformer(MODEL_NAME)
+    elapsed = time.monotonic() - t0
+    dim = model.get_sentence_embedding_dimension()
+    logger.info("Model loaded in %.1fs — dimension=%d", elapsed, dim)
+
+
+# -- Schemas -------------------------------------------------------------------
+
+
+class EmbedRequest(BaseModel):
+    texts: list[str] = Field(..., min_length=1, description="Texts to embed")
+
+
+class EmbedResponse(BaseModel):
+    embeddings: list[list[float]]
+    model: str
+    dimension: int
+
+
+class HealthResponse(BaseModel):
+    status: str
+    model: str
+    dimension: int
+
+
+# -- Endpoints -----------------------------------------------------------------
+
+
+@app.post("/embed", response_model=EmbedResponse)
+async def embed(request: EmbedRequest) -> EmbedResponse:
+    """Generate embeddings for a batch of texts."""
+    if model is None:
+        raise HTTPException(status_code=503, detail="Model not loaded yet")
+
+    vectors = model.encode(
+        request.texts,
+        batch_size=BATCH_SIZE,
+        show_progress_bar=False,
+        normalize_embeddings=True,
+    )
+    return EmbedResponse(
+        embeddings=vectors.tolist(),
+        model=MODEL_NAME,
+        dimension=model.get_sentence_embedding_dimension(),
+    )
+
+
+@app.get("/health", response_model=HealthResponse)
+async def health() -> HealthResponse:
+    """Health check — verifies the model is loaded."""
+    if model is None:
+        raise HTTPException(status_code=503, detail="Model not loaded yet")
+    return HealthResponse(
+        status="ok",
+        model=MODEL_NAME,
+        dimension=model.get_sentence_embedding_dimension(),
+    )

embedding-service/requirements.txt

@@ -0,0 +1,3 @@
+fastapi>=0.115.0,<1.0.0
+uvicorn[standard]>=0.32.0,<1.0.0
+sentence-transformers>=3.0.0,<4.0.0

embedding-service/test_main.py

@@ -0,0 +1,64 @@
+"""Tests for the embedding microservice API."""
+
+from __future__ import annotations
+
+from unittest.mock import MagicMock, patch
+
+import numpy as np
+import pytest
+from fastapi.testclient import TestClient
+
+import main
+
+
+@pytest.fixture(autouse=True)
+def _mock_model() -> None:
+    """Inject a mock SentenceTransformer model for all tests."""
+    mock = MagicMock()
+    mock.get_sentence_embedding_dimension.return_value = 3
+    mock.encode.return_value = np.array([[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]])
+    main.model = mock
+    yield
+    main.model = None
+
+
+@pytest.fixture
+def client() -> TestClient:
+    return TestClient(main.app)
+
+
+class TestEmbed:
+    def test_embed_returns_vectors(self, client: TestClient) -> None:
+        resp = client.post("/embed", json={"texts": ["hello", "world"]})
+        assert resp.status_code == 200
+        data = resp.json()
+        assert len(data["embeddings"]) == 2
+        assert data["dimension"] == 3
+        assert data["model"] == main.MODEL_NAME
+
+    def test_embed_empty_texts_rejected(self, client: TestClient) -> None:
+        resp = client.post("/embed", json={"texts": []})
+        assert resp.status_code == 422
+
+    def test_embed_missing_texts(self, client: TestClient) -> None:
+        resp = client.post("/embed", json={})
+        assert resp.status_code == 422
+
+    def test_embed_model_not_loaded(self, client: TestClient) -> None:
+        main.model = None
+        resp = client.post("/embed", json={"texts": ["test"]})
+        assert resp.status_code == 503
+
+
+class TestHealth:
+    def test_health_ok(self, client: TestClient) -> None:
+        resp = client.get("/health")
+        assert resp.status_code == 200
+        data = resp.json()
+        assert data["status"] == "ok"
+        assert data["dimension"] == 3
+
+    def test_health_model_not_loaded(self, client: TestClient) -> None:
+        main.model = None
+        resp = client.get("/health")
+        assert resp.status_code == 503