From 995e891d9b29c779c1731730d4d6c68db48a36e7 Mon Sep 17 00:00:00 2001 From: Pier-Jean Malandrino Date: Fri, 10 Apr 2026 21:12:11 +0200 Subject: [PATCH] =?UTF-8?q?feat:=20orchestrated=20ingestion=20pipeline=20D?= =?UTF-8?q?ocling=20=E2=86=92=20embedding=20=E2=86=92=20OpenSearch?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Closes #72 --- document-parser/api/ingestion.py | 82 +++++++++ document-parser/api/schemas.py | 10 ++ document-parser/main.py | 28 ++- document-parser/services/ingestion_service.py | 166 ++++++++++++++++++ .../tests/test_ingestion_service.py | 150 ++++++++++++++++ 5 files changed, 435 insertions(+), 1 deletion(-) create mode 100644 document-parser/api/ingestion.py create mode 100644 document-parser/services/ingestion_service.py create mode 100644 document-parser/tests/test_ingestion_service.py diff --git a/document-parser/api/ingestion.py b/document-parser/api/ingestion.py new file mode 100644 index 0000000..2f48e7d --- /dev/null +++ b/document-parser/api/ingestion.py @@ -0,0 +1,82 @@ +"""Ingestion API router — trigger and manage vector ingestion pipeline.""" + +from __future__ import annotations + +import logging +from typing import Annotated + +from fastapi import APIRouter, Depends, HTTPException, Request + +from api.schemas import IngestionResponse, IngestionStatusResponse +from services.analysis_service import AnalysisService +from services.ingestion_service import IngestionService + +logger = logging.getLogger(__name__) +router = APIRouter(prefix="/api/ingestion", tags=["ingestion"]) + + +def _get_ingestion_service(request: Request) -> IngestionService: + svc = request.app.state.ingestion_service + if svc is None: + raise HTTPException( + status_code=503, + detail="Ingestion not available (EMBEDDING_URL and OPENSEARCH_URL required)", + ) + return svc + + +def _get_analysis_service(request: Request) -> AnalysisService: + return request.app.state.analysis_service + + +IngestionDep = Annotated[IngestionService, Depends(_get_ingestion_service)] +AnalysisDep = Annotated[AnalysisService, Depends(_get_analysis_service)] + + +@router.post("/{job_id}", response_model=IngestionResponse) +async def ingest_analysis( + job_id: str, + ingestion: IngestionDep, + analysis: AnalysisDep, +) -> IngestionResponse: + """Ingest a completed analysis into the vector index. + + Takes the chunks from an existing analysis job, embeds them, + and indexes them into OpenSearch. + """ + job = await analysis.find_by_id(job_id) + if not job: + raise HTTPException(status_code=404, detail="Analysis not found") + if job.status.value != "COMPLETED": + raise HTTPException(status_code=400, detail="Analysis is not completed") + if not job.chunks_json: + raise HTTPException(status_code=400, detail="Analysis has no chunks — run chunking first") + + try: + result = await ingestion.ingest( + doc_id=job.document_id, + filename=job.document_filename or "unknown", + chunks_json=job.chunks_json, + ) + except Exception as e: + logger.exception("Ingestion failed for job %s", job_id) + raise HTTPException(status_code=500, detail=f"Ingestion failed: {e}") from e + + return IngestionResponse( + doc_id=result.doc_id, + chunks_indexed=result.chunks_indexed, + embedding_dimension=result.embedding_dimension, + ) + + +@router.delete("/{doc_id}", status_code=204) +async def delete_ingested_document(doc_id: str, ingestion: IngestionDep) -> None: + """Delete all indexed chunks for a document.""" + await ingestion.delete_document(doc_id) + + +@router.get("/status", response_model=IngestionStatusResponse) +async def ingestion_status(request: Request) -> IngestionStatusResponse: + """Check if the ingestion pipeline is available.""" + available = request.app.state.ingestion_service is not None + return IngestionStatusResponse(available=available) diff --git a/document-parser/api/schemas.py b/document-parser/api/schemas.py index 42079a6..808b70b 100644 --- a/document-parser/api/schemas.py +++ b/document-parser/api/schemas.py @@ -180,3 +180,13 @@ class RechunkRequest(BaseModel): chunkingOptions: ChunkingOptionsRequest = Field( validation_alias=AliasChoices("chunkingOptions", "chunking_options") ) + + +class IngestionResponse(_CamelModel): + doc_id: str + chunks_indexed: int + embedding_dimension: int + + +class IngestionStatusResponse(_CamelModel): + available: bool diff --git a/document-parser/main.py b/document-parser/main.py index aa6573c..a316b34 100644 --- a/document-parser/main.py +++ b/document-parser/main.py @@ -20,6 +20,7 @@ from fastapi.middleware.cors import CORSMiddleware from api.analyses import router as analyses_router from api.documents import router as documents_router +from api.ingestion import router as ingestion_router from api.schemas import HealthResponse from infra.rate_limiter import RateLimiterMiddleware from infra.settings import settings @@ -28,6 +29,7 @@ from persistence.database import get_connection, init_db from persistence.document_repo import SqliteDocumentRepository from services.analysis_service import AnalysisConfig, AnalysisService from services.document_service import DocumentConfig, DocumentService +from services.ingestion_service import IngestionConfig, IngestionService logging.basicConfig( level=logging.INFO, @@ -87,6 +89,28 @@ def _build_analysis_service( ) +def _build_ingestion_service() -> IngestionService | None: + """Build the ingestion service — only if embedding + opensearch are configured.""" + if not settings.embedding_url or not settings.opensearch_url: + logger.info("Ingestion disabled (EMBEDDING_URL or OPENSEARCH_URL not set)") + return None + + from infra.embedding_client import EmbeddingClient + from infra.opensearch_store import OpenSearchStore + + embedding = EmbeddingClient(settings.embedding_url) + vector_store = OpenSearchStore(settings.opensearch_url) + config = IngestionConfig( + embedding_dimension=settings.embedding_dimension, + ) + logger.info( + "Ingestion enabled (embedding=%s, opensearch=%s)", + settings.embedding_url, + settings.opensearch_url, + ) + return IngestionService(embedding, vector_store, config) + + def _build_document_service( document_repo: SqliteDocumentRepository, analysis_repo: SqliteAnalysisRepository, @@ -114,6 +138,7 @@ async def lifespan(app: FastAPI) -> AsyncIterator[None]: document_repo, analysis_repo = _build_repos() app.state.analysis_service = _build_analysis_service(document_repo, analysis_repo) app.state.document_service = _build_document_service(document_repo, analysis_repo) + app.state.ingestion_service = _build_ingestion_service() logger.info("Docling Studio backend ready (engine=%s)", settings.conversion_engine) yield @@ -128,7 +153,7 @@ app.add_middleware( CORSMiddleware, allow_origins=settings.cors_origins, allow_credentials=True, - allow_methods=["GET", "POST", "DELETE", "OPTIONS"], + allow_methods=["GET", "POST", "PATCH", "DELETE", "OPTIONS"], allow_headers=["Content-Type", "Authorization"], ) if settings.rate_limit_rpm > 0: @@ -140,6 +165,7 @@ if settings.rate_limit_rpm > 0: app.include_router(documents_router) app.include_router(analyses_router) +app.include_router(ingestion_router) @app.get("/api/health", response_model=HealthResponse) diff --git a/document-parser/services/ingestion_service.py b/document-parser/services/ingestion_service.py new file mode 100644 index 0000000..ca29c1e --- /dev/null +++ b/document-parser/services/ingestion_service.py @@ -0,0 +1,166 @@ +"""Ingestion service — orchestrates Docling → embedding → OpenSearch. + +Chains the full ingestion pipeline: +1. Convert document via Docling (reuse existing analysis) +2. Chunk with selected strategy +3. Embed all chunk texts via EmbeddingService +4. Index into OpenSearch via VectorStore + +Idempotent: re-ingesting a document deletes old chunks before re-indexing. +""" + +from __future__ import annotations + +import json +import logging +from dataclasses import dataclass +from typing import TYPE_CHECKING + +from domain.vector_schema import ( + ChunkBboxEntry, + ChunkOrigin, + IndexedChunk, + build_index_mapping, +) + +if TYPE_CHECKING: + from domain.ports import EmbeddingService, VectorStore + +logger = logging.getLogger(__name__) + + +@dataclass +class IngestionConfig: + """Configuration for the ingestion pipeline.""" + + index_name: str = "docling-studio-chunks" + embedding_dimension: int = 384 + + +@dataclass +class IngestionResult: + """Result of an ingestion pipeline run.""" + + doc_id: str + chunks_indexed: int + embedding_dimension: int + + +class IngestionService: + """Orchestrates the embedding + indexing pipeline.""" + + def __init__( + self, + embedding_service: EmbeddingService, + vector_store: VectorStore, + config: IngestionConfig | None = None, + ) -> None: + self._embedding = embedding_service + self._vector_store = vector_store + self._config = config or IngestionConfig() + + async def ensure_index(self) -> None: + """Ensure the vector index exists with the correct mapping.""" + mapping = build_index_mapping(self._config.embedding_dimension) + await self._vector_store.ensure_index(self._config.index_name, mapping) + + async def ingest( + self, + doc_id: str, + filename: str, + chunks_json: str, + *, + binary_hash: str | None = None, + ) -> IngestionResult: + """Run the embedding + indexing pipeline on pre-chunked data. + + This method is idempotent: it deletes any existing chunks for the + document before re-indexing. + + Args: + doc_id: Unique document identifier. + filename: Original filename. + chunks_json: JSON-serialized list of chunk dicts (from analysis). + binary_hash: Optional hash of the source file for provenance. + + Returns: + IngestionResult with the number of chunks indexed. + """ + await self.ensure_index() + + chunks_data: list[dict] = json.loads(chunks_json) + active_chunks = [c for c in chunks_data if not c.get("deleted")] + if not active_chunks: + logger.info("No active chunks for doc %s — skipping ingestion", doc_id) + return IngestionResult(doc_id=doc_id, chunks_indexed=0, embedding_dimension=0) + + # 1. Embed all chunk texts + texts = [c["text"] for c in active_chunks] + logger.info("Embedding %d chunks for doc %s", len(texts), doc_id) + embeddings = await self._embedding.embed(texts) + + # 2. Build IndexedChunk domain objects + origin = ( + ChunkOrigin(binary_hash=binary_hash or "", filename=filename) if binary_hash else None + ) + indexed_chunks: list[IndexedChunk] = [] + for i, (chunk_data, embedding) in enumerate(zip(active_chunks, embeddings, strict=True)): + bboxes = [ + ChunkBboxEntry( + page=b["page"], + x=b["bbox"][0] if b.get("bbox") else 0, + y=b["bbox"][1] if b.get("bbox") else 0, + w=(b["bbox"][2] - b["bbox"][0]) if b.get("bbox") and len(b["bbox"]) >= 4 else 0, + h=(b["bbox"][3] - b["bbox"][1]) if b.get("bbox") and len(b["bbox"]) >= 4 else 0, + ) + for b in chunk_data.get("bboxes", []) + ] + indexed_chunks.append( + IndexedChunk( + doc_id=doc_id, + filename=filename, + content=chunk_data["text"], + embedding=embedding, + chunk_index=i, + chunk_type=chunk_data.get("chunkType", "text"), + page_number=chunk_data.get("sourcePage", 0) or 0, + bboxes=bboxes, + headings=chunk_data.get("headings", []), + origin=origin, + ) + ) + + # 3. Delete old chunks (idempotent re-indexing) + deleted = await self._vector_store.delete_document(self._config.index_name, doc_id) + if deleted: + logger.info("Deleted %d old chunks for doc %s", deleted, doc_id) + + # 4. Index new chunks + indexed = await self._vector_store.index_chunks(self._config.index_name, indexed_chunks) + logger.info("Indexed %d/%d chunks for doc %s", indexed, len(indexed_chunks), doc_id) + + return IngestionResult( + doc_id=doc_id, + chunks_indexed=indexed, + embedding_dimension=len(embeddings[0]) if embeddings else 0, + ) + + async def delete_document(self, doc_id: str) -> int: + """Remove all indexed chunks for a document.""" + return await self._vector_store.delete_document(self._config.index_name, doc_id) + + async def search( + self, + query: str, + *, + k: int = 10, + doc_id: str | None = None, + ) -> list: + """Semantic search: embed the query then find nearest chunks.""" + embeddings = await self._embedding.embed([query]) + return await self._vector_store.search_similar( + self._config.index_name, + embeddings[0], + k=k, + doc_id=doc_id, + ) diff --git a/document-parser/tests/test_ingestion_service.py b/document-parser/tests/test_ingestion_service.py new file mode 100644 index 0000000..cbea903 --- /dev/null +++ b/document-parser/tests/test_ingestion_service.py @@ -0,0 +1,150 @@ +"""Tests for the ingestion service (services.ingestion_service).""" + +from __future__ import annotations + +import json +from unittest.mock import AsyncMock + +import pytest + +from services.ingestion_service import IngestionConfig, IngestionService + + +def _make_chunks_json(count: int = 3, *, with_deleted: bool = False) -> str: + chunks = [] + for i in range(count): + chunk = { + "text": f"chunk text {i}", + "headings": [f"Heading {i}"], + "sourcePage": i + 1, + "tokenCount": 10, + "bboxes": [{"page": i + 1, "bbox": [0.0, 0.0, 100.0, 50.0]}], + } + if with_deleted and i == count - 1: + chunk["deleted"] = True + chunks.append(chunk) + return json.dumps(chunks) + + +@pytest.fixture +def mock_embedding() -> AsyncMock: + svc = AsyncMock() + svc.embed.return_value = [[0.1, 0.2, 0.3]] * 3 + return svc + + +@pytest.fixture +def mock_vector_store() -> AsyncMock: + store = AsyncMock() + store.ensure_index.return_value = None + store.delete_document.return_value = 0 + store.index_chunks.return_value = 3 + return store + + +@pytest.fixture +def service(mock_embedding: AsyncMock, mock_vector_store: AsyncMock) -> IngestionService: + return IngestionService( + embedding_service=mock_embedding, + vector_store=mock_vector_store, + config=IngestionConfig(index_name="test-idx", embedding_dimension=3), + ) + + +class TestIngest: + async def test_full_pipeline( + self, service: IngestionService, mock_embedding: AsyncMock, mock_vector_store: AsyncMock + ) -> None: + result = await service.ingest("doc-1", "test.pdf", _make_chunks_json(3)) + + assert result.doc_id == "doc-1" + assert result.chunks_indexed == 3 + mock_embedding.embed.assert_awaited_once() + texts = mock_embedding.embed.call_args[0][0] + assert len(texts) == 3 + mock_vector_store.ensure_index.assert_awaited_once() + mock_vector_store.delete_document.assert_awaited_once_with("test-idx", "doc-1") + mock_vector_store.index_chunks.assert_awaited_once() + indexed = mock_vector_store.index_chunks.call_args[0][1] + assert len(indexed) == 3 + assert indexed[0].doc_id == "doc-1" + assert indexed[0].filename == "test.pdf" + assert indexed[0].embedding == [0.1, 0.2, 0.3] + + async def test_skips_deleted_chunks( + self, service: IngestionService, mock_embedding: AsyncMock, mock_vector_store: AsyncMock + ) -> None: + mock_embedding.embed.return_value = [[0.1, 0.2, 0.3]] * 2 + mock_vector_store.index_chunks.return_value = 2 + result = await service.ingest("doc-1", "test.pdf", _make_chunks_json(3, with_deleted=True)) + + assert result.chunks_indexed == 2 + texts = mock_embedding.embed.call_args[0][0] + assert len(texts) == 2 + + async def test_empty_chunks( + self, service: IngestionService, mock_embedding: AsyncMock, mock_vector_store: AsyncMock + ) -> None: + result = await service.ingest("doc-1", "test.pdf", json.dumps([])) + assert result.chunks_indexed == 0 + mock_embedding.embed.assert_not_awaited() + + async def test_idempotent_deletes_old( + self, service: IngestionService, mock_vector_store: AsyncMock + ) -> None: + mock_vector_store.delete_document.return_value = 5 + await service.ingest("doc-1", "test.pdf", _make_chunks_json(3)) + mock_vector_store.delete_document.assert_awaited_once_with("test-idx", "doc-1") + + async def test_bbox_conversion( + self, service: IngestionService, mock_embedding: AsyncMock, mock_vector_store: AsyncMock + ) -> None: + mock_embedding.embed.return_value = [[0.1, 0.2, 0.3]] + mock_vector_store.index_chunks.return_value = 1 + await service.ingest("doc-1", "test.pdf", _make_chunks_json(1)) + indexed = mock_vector_store.index_chunks.call_args[0][1] + bbox = indexed[0].bboxes[0] + assert bbox.x == 0.0 + assert bbox.y == 0.0 + assert bbox.w == 100.0 + assert bbox.h == 50.0 + + async def test_with_binary_hash( + self, service: IngestionService, mock_vector_store: AsyncMock + ) -> None: + mock_embedding = service._embedding + mock_embedding.embed.return_value = [[0.1]] * 1 + await service.ingest("doc-1", "test.pdf", _make_chunks_json(1), binary_hash="abc123") + indexed = mock_vector_store.index_chunks.call_args[0][1] + assert indexed[0].origin is not None + assert indexed[0].origin.binary_hash == "abc123" + + +class TestDeleteDocument: + async def test_delegates_to_vector_store( + self, service: IngestionService, mock_vector_store: AsyncMock + ) -> None: + mock_vector_store.delete_document.return_value = 3 + result = await service.delete_document("doc-1") + assert result == 3 + + +class TestSearch: + async def test_embeds_and_searches( + self, service: IngestionService, mock_embedding: AsyncMock, mock_vector_store: AsyncMock + ) -> None: + mock_embedding.embed.return_value = [[0.5, 0.6, 0.7]] + mock_vector_store.search_similar.return_value = [] + await service.search("test query", k=5) + mock_embedding.embed.assert_awaited_once_with(["test query"]) + mock_vector_store.search_similar.assert_awaited_once() + + +class TestEnsureIndex: + async def test_calls_vector_store( + self, service: IngestionService, mock_vector_store: AsyncMock + ) -> None: + await service.ensure_index() + mock_vector_store.ensure_index.assert_awaited_once() + call_args = mock_vector_store.ensure_index.call_args + assert call_args[0][0] == "test-idx"