From a21daa24da0122fd63c0b6cbacbe8ac4df59c7d3 Mon Sep 17 00:00:00 2001 From: Pier-Jean Malandrino Date: Fri, 10 Apr 2026 20:53:24 +0200 Subject: [PATCH] feat: add embedding microservice and EmbeddingService port Closes #71 --- .env.example | 6 + CHANGELOG.md | 2 + docker-compose.dev.yml | 22 ++++ document-parser/domain/ports.py | 12 ++ document-parser/infra/embedding_client.py | 51 ++++++++ document-parser/infra/settings.py | 2 + .../tests/test_embedding_client.py | 112 ++++++++++++++++++ embedding-service/Dockerfile | 17 +++ embedding-service/main.py | 89 ++++++++++++++ embedding-service/requirements.txt | 3 + embedding-service/test_main.py | 64 ++++++++++ 11 files changed, 380 insertions(+) create mode 100644 document-parser/infra/embedding_client.py create mode 100644 document-parser/tests/test_embedding_client.py create mode 100644 embedding-service/Dockerfile create mode 100644 embedding-service/main.py create mode 100644 embedding-service/requirements.txt create mode 100644 embedding-service/test_main.py diff --git a/.env.example b/.env.example index bb9ec43..b423dce 100644 --- a/.env.example +++ b/.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 diff --git a/CHANGELOG.md b/CHANGELOG.md index 94586c3..2031ddf 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/docker-compose.dev.yml b/docker-compose.dev.yml index 4f7cee6..fc08ad2 100644 --- a/docker-compose.dev.yml +++ b/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: diff --git a/document-parser/domain/ports.py b/document-parser/domain/ports.py index 0375da5..1cd3fd5 100644 --- a/document-parser/domain/ports.py +++ b/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. diff --git a/document-parser/infra/embedding_client.py b/document-parser/infra/embedding_client.py new file mode 100644 index 0000000..c0fc861 --- /dev/null +++ b/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 diff --git a/document-parser/infra/settings.py b/document-parser/infra/settings.py index bc7dd49..f93c299 100644 --- a/document-parser/infra/settings.py +++ b/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"), diff --git a/document-parser/tests/test_embedding_client.py b/document-parser/tests/test_embedding_client.py new file mode 100644 index 0000000..0a88eb6 --- /dev/null +++ b/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 diff --git a/embedding-service/Dockerfile b/embedding-service/Dockerfile new file mode 100644 index 0000000..b7b8e04 --- /dev/null +++ b/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"] diff --git a/embedding-service/main.py b/embedding-service/main.py new file mode 100644 index 0000000..497a5a2 --- /dev/null +++ b/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(), + ) diff --git a/embedding-service/requirements.txt b/embedding-service/requirements.txt new file mode 100644 index 0000000..abe998e --- /dev/null +++ b/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 diff --git a/embedding-service/test_main.py b/embedding-service/test_main.py new file mode 100644 index 0000000..e90d01b --- /dev/null +++ b/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