danvics/docling-studio
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 1

feat(#72): orchestrated ingestion pipeline — Docling → embedding → OpenSearch

Browse source 
Add IngestionService chaining analysis chunks → EmbeddingClient → OpenSearchStore.
Idempotent: existing doc chunks deleted before re-indexing. REST API:
  POST /api/ingestion/{jobId}, DELETE /api/ingestion/{docId}, GET /api/ingestion/status.
Wired in lifespan when OPENSEARCH_URL + EMBEDDING_URL are set. 25 new tests.
This commit is contained in:
Pier-Jean Malandrino 2026-04-10 21:45:52 +02:00
parent c49708990c
commit 4c3870bf3e
5 changed files with 759 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

120
document-parser/api/ingestion.py Normal file
View file

@ -0,0 +1,120 @@
"""Ingestion API — REST endpoints for the embedding → OpenSearch pipeline.
Routes:
POST /api/ingestion/{job_id} — Trigger ingestion for a completed analysis
DELETE /api/ingestion/{doc_id} — Remove indexed chunks for a document
GET /api/ingestion/status — Check whether the ingestion pipeline is available
"""
from __future__ import annotations
import logging
from fastapi import APIRouter, HTTPException, Request, status
from pydantic import BaseModel
logger = logging.getLogger(__name__)
router = APIRouter(prefix="/api/ingestion", tags=["ingestion"])
# ---------------------------------------------------------------------------
# Response schemas
# ---------------------------------------------------------------------------
class IngestionResponse(BaseModel):
doc_id: str
chunks_indexed: int
embedding_dimension: int
model_config = {"populate_by_name": True}
class IngestionStatusResponse(BaseModel):
available: bool
reason: str = ""
# ---------------------------------------------------------------------------
# Dependency helpers
# ---------------------------------------------------------------------------
def _get_ingestion_service(request: Request):
svc = getattr(request.app.state, "ingestion_service", None)
if svc is None:
raise HTTPException(
status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
detail="Ingestion pipeline not available (OpenSearch or embedding service not configured).",
)
return svc
# ---------------------------------------------------------------------------
# Endpoints
# ---------------------------------------------------------------------------
@router.get("/status", response_model=IngestionStatusResponse)
async def get_ingestion_status(request: Request) -> IngestionStatusResponse:
"""Return whether the ingestion pipeline (OpenSearch + embedding) is available."""
svc = getattr(request.app.state, "ingestion_service", None)
if svc is None:
return IngestionStatusResponse(
available=False,
reason="OpenSearch or embedding service not configured",
)
return IngestionStatusResponse(available=True)
@router.post("/{job_id}", response_model=IngestionResponse, status_code=status.HTTP_200_OK)
async def ingest_analysis(job_id: str, request: Request) -> IngestionResponse:
"""Run the full ingestion pipeline for a completed analysis job.
Chains: loaded chunks → embedding → OpenSearch indexing.
Idempotent: re-ingesting a document replaces existing indexed chunks.
"""
svc = _get_ingestion_service(request)
try:
from services.ingestion_service import IngestionError
result = await svc.ingest(job_id)
except IngestionError as exc:
logger.warning("Ingestion failed for job %s: %s", job_id, exc)
raise HTTPException(
status_code=status.HTTP_422_UNPROCESSABLE_ENTITY, detail=str(exc)
) from exc
except Exception as exc:
logger.exception("Unexpected error during ingestion of job %s", job_id)
raise HTTPException(
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
detail=f"Ingestion error: {exc}",
) from exc
return IngestionResponse(
doc_id=result.doc_id,
chunks_indexed=result.chunks_indexed,
embedding_dimension=result.embedding_dimension,
)
@router.delete("/{doc_id}", status_code=status.HTTP_204_NO_CONTENT)
async def delete_ingested(doc_id: str, request: Request) -> None:
"""Remove all indexed chunks for a document from OpenSearch."""
svc = _get_ingestion_service(request)
try:
from services.ingestion_service import IngestionError
await svc.delete(doc_id)
except IngestionError as exc:
logger.warning("Delete failed for document %s: %s", doc_id, exc)
raise HTTPException(
status_code=status.HTTP_422_UNPROCESSABLE_ENTITY, detail=str(exc)
) from exc
except Exception as exc:
logger.exception("Unexpected error deleting chunks for document %s", doc_id)
raise HTTPException(
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
detail=f"Delete error: {exc}",
) from exc

36
document-parser/main.py
View file

@ -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
@ -108,12 +109,46 @@ def _build_document_service(
# ---------------------------------------------------------------------------
def _build_ingestion_service(
document_repo: SqliteDocumentRepository,
analysis_repo: SqliteAnalysisRepository,
):
"""Build ingestion service if OpenSearch + embedding are configured."""
if not settings.opensearch_url or not settings.embedding_url:
logger.info(
"Ingestion pipeline disabled (OPENSEARCH_URL=%r, EMBEDDING_URL=%r)",
settings.opensearch_url,
settings.embedding_url,
)
return None
from infra.embedding_client import EmbeddingClient
from infra.opensearch_store import OpenSearchStore
from services.ingestion_service import IngestionService
embedding_svc = EmbeddingClient(settings.embedding_url)
vector_store = OpenSearchStore(hosts=[settings.opensearch_url])
logger.info(
"Ingestion pipeline enabled (opensearch=%s, embedding=%s)",
settings.opensearch_url,
settings.embedding_url,
)
return IngestionService(
analysis_repo=analysis_repo,
document_repo=document_repo,
embedding_svc=embedding_svc,
vector_store=vector_store,
embedding_dimension=settings.embedding_dimension,
)
@asynccontextmanager
async def lifespan(app: FastAPI) -> AsyncIterator[None]:
await init_db()
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(document_repo, analysis_repo)
logger.info("Docling Studio backend ready (engine=%s)", settings.conversion_engine)
yield
@ -140,6 +175,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)

223
document-parser/services/ingestion_service.py Normal file
View file

@ -0,0 +1,223 @@
"""Ingestion service — orchestrates the full Docling → embedding → OpenSearch pipeline.
Pipeline stages:
1. Load analysis job + document metadata
2. Parse chunks from the completed job
3. Generate embeddings for each chunk
4. Idempotently index chunks in OpenSearch
Idempotency: existing chunks for the document are deleted before re-indexing,
so re-ingesting a document always produces a clean, up-to-date index.
"""
from __future__ import annotations
import json
import logging
from dataclasses import dataclass
from typing import TYPE_CHECKING
from domain.models import AnalysisStatus
from domain.vector_schema import (
DEFAULT_INDEX_NAME,
ChunkBboxEntry,
ChunkOrigin,
IndexedChunk,
build_index_mapping,
)
if TYPE_CHECKING:
from domain.ports import (
AnalysisRepository,
DocumentRepository,
EmbeddingService,
VectorStore,
)
logger = logging.getLogger(__name__)
@dataclass(frozen=True)
class IngestionResult:
"""Result returned by a successful ingestion run."""
doc_id: str
chunks_indexed: int
embedding_dimension: int
class IngestionError(Exception):
"""Raised when the ingestion pipeline cannot complete."""
class IngestionService:
"""Orchestrates the full ingestion pipeline for a single analysis job."""
def __init__(
self,
analysis_repo: AnalysisRepository,
document_repo: DocumentRepository,
embedding_svc: EmbeddingService,
vector_store: VectorStore,
*,
index_name: str = DEFAULT_INDEX_NAME,
embedding_dimension: int = 384,
) -> None:
self._analysis_repo = analysis_repo
self._document_repo = document_repo
self._embedding_svc = embedding_svc
self._vector_store = vector_store
self._index_name = index_name
self._embedding_dimension = embedding_dimension
# ------------------------------------------------------------------
# Public API
# ------------------------------------------------------------------
async def ingest(self, job_id: str) -> IngestionResult:
"""Run the full ingestion pipeline for the given analysis job.
Args:
job_id: ID of a COMPLETED analysis job that has chunks.
Returns:
IngestionResult with doc_id and number of indexed chunks.
Raises:
IngestionError: if any pipeline stage fails.
"""
logger.info("Ingestion started for job %s", job_id)
# --- Stage 1: load job ---
job = await self._analysis_repo.find_by_id(job_id)
if job is None:
raise IngestionError(f"Analysis job not found: {job_id}")
if job.status != AnalysisStatus.COMPLETED:
raise IngestionError(
f"Job {job_id} is not COMPLETED (status={job.status}). Run analysis first."
)
if not job.chunks_json:
raise IngestionError(f"Job {job_id} has no chunks. Run chunking first.")
doc = await self._document_repo.find_by_id(job.document_id)
if doc is None:
raise IngestionError(f"Document not found: {job.document_id}")
logger.info("Loaded job %s — document: %s", job_id, doc.filename)
# --- Stage 2: parse chunks ---
try:
raw_chunks: list[dict] = json.loads(job.chunks_json)
except json.JSONDecodeError as exc:
raise IngestionError(f"Invalid chunks_json in job {job_id}") from exc
if not raw_chunks:
raise IngestionError(f"Job {job_id} has empty chunk list. Chunk first.")
texts = [c.get("text", "") for c in raw_chunks]
logger.info("Parsed %d chunks from job %s", len(texts), job_id)
# --- Stage 3: generate embeddings ---
try:
embeddings = await self._embedding_svc.embed(texts)
except Exception as exc:
raise IngestionError(f"Embedding generation failed: {exc}") from exc
if len(embeddings) != len(texts):
raise IngestionError(
f"Embedding dimension mismatch: got {len(embeddings)} vectors for {len(texts)} texts"
)
# Detect embedding dimension from the first non-empty vector
detected_dim = self._embedding_dimension
for vec in embeddings:
if vec:
detected_dim = len(vec)
break
logger.info("Generated %d embeddings (dim=%d)", len(embeddings), detected_dim)
# --- Stage 4: ensure index exists ---
mapping = build_index_mapping(detected_dim)
try:
await self._vector_store.ensure_index(self._index_name, mapping)
except Exception as exc:
raise IngestionError(f"Failed to ensure index: {exc}") from exc
# --- Stage 4b: delete existing chunks (idempotency) ---
try:
deleted = await self._vector_store.delete_document(self._index_name, doc.id)
if deleted:
logger.info("Deleted %d existing chunks for document %s", deleted, doc.id)
except Exception as exc:
raise IngestionError(f"Failed to delete existing chunks: {exc}") from exc
# --- Stage 5: build IndexedChunk list ---
origin = ChunkOrigin(
binary_hash=doc.id, # use doc id as stable identifier
filename=doc.filename,
)
indexed_chunks: list[IndexedChunk] = []
for i, (raw, emb) in enumerate(zip(raw_chunks, embeddings, strict=True)):
bboxes = [
ChunkBboxEntry(
page=b.get("page", 0),
x=b["bbox"][0] if b.get("bbox") else 0.0,
y=b["bbox"][1] if b.get("bbox") else 0.0,
w=(b["bbox"][2] - b["bbox"][0])
if b.get("bbox") and len(b["bbox"]) >= 4
else 0.0,
h=(b["bbox"][3] - b["bbox"][1])
if b.get("bbox") and len(b["bbox"]) >= 4
else 0.0,
)
for b in raw.get("bboxes", [])
]
chunk = IndexedChunk(
doc_id=doc.id,
filename=doc.filename,
content=raw.get("text", ""),
embedding=emb,
chunk_index=i,
chunk_type="text",
page_number=raw.get("sourcePage") or 0,
bboxes=bboxes,
headings=raw.get("headings", []),
doc_items=[],
origin=origin,
)
indexed_chunks.append(chunk)
# --- Stage 6: bulk index ---
try:
indexed_count = await self._vector_store.index_chunks(self._index_name, indexed_chunks)
except Exception as exc:
raise IngestionError(f"Bulk indexing failed: {exc}") from exc
logger.info(
"Ingestion complete — %d/%d chunks indexed for document %s",
indexed_count,
len(indexed_chunks),
doc.id,
)
return IngestionResult(
doc_id=doc.id,
chunks_indexed=indexed_count,
embedding_dimension=detected_dim,
)
async def delete(self, doc_id: str) -> int:
"""Remove all indexed chunks for a document.
Returns:
Number of chunks deleted.
"""
logger.info("Deleting indexed chunks for document %s", doc_id)
try:
deleted = await self._vector_store.delete_document(self._index_name, doc_id)
except Exception as exc:
raise IngestionError(f"Failed to delete from index: {exc}") from exc
logger.info("Deleted %d chunks for document %s", deleted, doc_id)
return deleted

112
document-parser/tests/test_ingestion_api.py Normal file
View file

@ -0,0 +1,112 @@
"""Tests for the ingestion REST API endpoints."""
from __future__ import annotations
from unittest.mock import AsyncMock, MagicMock
import pytest
from fastapi.testclient import TestClient
from main import app
from services.ingestion_service import IngestionError, IngestionResult
@pytest.fixture()
def client():
return TestClient(app, raise_server_exceptions=False)
@pytest.fixture(autouse=True)
def _clear_ingestion_service():
"""Ensure ingestion_service is reset after each test."""
original = getattr(app.state, "ingestion_service", None)
yield
app.state.ingestion_service = original
# ---------------------------------------------------------------------------
# GET /api/ingestion/status
# ---------------------------------------------------------------------------
def test_status_available(client):
app.state.ingestion_service = MagicMock()
resp = client.get("/api/ingestion/status")
assert resp.status_code == 200
assert resp.json()["available"] is True
def test_status_unavailable(client):
app.state.ingestion_service = None
resp = client.get("/api/ingestion/status")
assert resp.status_code == 200
assert resp.json()["available"] is False
# ---------------------------------------------------------------------------
# POST /api/ingestion/{job_id}
# ---------------------------------------------------------------------------
def test_ingest_success(client):
svc = MagicMock()
svc.ingest = AsyncMock(
return_value=IngestionResult(doc_id="doc-1", chunks_indexed=5, embedding_dimension=384)
)
app.state.ingestion_service = svc
resp = client.post("/api/ingestion/job-1")
assert resp.status_code == 200
data = resp.json()
assert data["doc_id"] == "doc-1"
assert data["chunks_indexed"] == 5
assert data["embedding_dimension"] == 384
def test_ingest_no_service(client):
app.state.ingestion_service = None
resp = client.post("/api/ingestion/job-1")
assert resp.status_code == 503
def test_ingest_ingestion_error(client):
svc = MagicMock()
svc.ingest = AsyncMock(side_effect=IngestionError("job not found"))
app.state.ingestion_service = svc
resp = client.post("/api/ingestion/job-1")
assert resp.status_code == 422
assert "job not found" in resp.json()["detail"]
def test_ingest_unexpected_error(client):
svc = MagicMock()
svc.ingest = AsyncMock(side_effect=RuntimeError("boom"))
app.state.ingestion_service = svc
resp = client.post("/api/ingestion/job-1")
assert resp.status_code == 500
# ---------------------------------------------------------------------------
# DELETE /api/ingestion/{doc_id}
# ---------------------------------------------------------------------------
def test_delete_success(client):
svc = MagicMock()
svc.delete = AsyncMock(return_value=3)
app.state.ingestion_service = svc
resp = client.delete("/api/ingestion/doc-1")
assert resp.status_code == 204
def test_delete_no_service(client):
app.state.ingestion_service = None
resp = client.delete("/api/ingestion/doc-1")
assert resp.status_code == 503
def test_delete_ingestion_error(client):
svc = MagicMock()
svc.delete = AsyncMock(side_effect=IngestionError("opensearch error"))
app.state.ingestion_service = svc
resp = client.delete("/api/ingestion/doc-1")
assert resp.status_code == 422

268
document-parser/tests/test_ingestion_service.py Normal file
View file

@ -0,0 +1,268 @@
"""Tests for IngestionService — orchestrated pipeline."""
from __future__ import annotations
import json
from unittest.mock import AsyncMock, MagicMock
import pytest
from domain.models import AnalysisJob, AnalysisStatus, Document
from domain.vector_schema import IndexedChunk
from services.ingestion_service import IngestionError, IngestionService
# ---------------------------------------------------------------------------
# Fixtures
# ---------------------------------------------------------------------------
def _make_doc(doc_id: str = "doc-1", filename: str = "test.pdf") -> Document:
return Document(id=doc_id, filename=filename)
def _make_job(
job_id: str = "job-1",
doc_id: str = "doc-1",
status: AnalysisStatus = AnalysisStatus.COMPLETED,
chunks_json: str | None = None,
) -> AnalysisJob:
if chunks_json is None:
chunks_json = json.dumps(
[
{
"text": "Hello world",
"headings": ["Introduction"],
"sourcePage": 1,
"bboxes": [{"page": 1, "bbox": [10.0, 20.0, 100.0, 50.0]}],
},
{
"text": "Second chunk",
"headings": [],
"sourcePage": 2,
"bboxes": [],
},
]
)
job = AnalysisJob(id=job_id, document_id=doc_id, status=status)
job.chunks_json = chunks_json
return job
def _make_service(
*,
job: AnalysisJob | None = None,
doc: Document | None = None,
embeddings: list[list[float]] | None = None,
indexed_count: int = 2,
) -> IngestionService:
if job is None:
job = _make_job()
if doc is None:
doc = _make_doc()
if embeddings is None:
embeddings = [[0.1] * 384, [0.2] * 384]
analysis_repo = MagicMock()
analysis_repo.find_by_id = AsyncMock(return_value=job)
document_repo = MagicMock()
document_repo.find_by_id = AsyncMock(return_value=doc)
embedding_svc = MagicMock()
embedding_svc.embed = AsyncMock(return_value=embeddings)
vector_store = MagicMock()
vector_store.ensure_index = AsyncMock()
vector_store.delete_document = AsyncMock(return_value=0)
vector_store.index_chunks = AsyncMock(return_value=indexed_count)
return (
IngestionService(
analysis_repo=analysis_repo,
document_repo=document_repo,
embedding_svc=embedding_svc,
vector_store=vector_store,
),
analysis_repo,
document_repo,
embedding_svc,
vector_store,
)
# ---------------------------------------------------------------------------
# Happy path
# ---------------------------------------------------------------------------
async def test_ingest_success():
svc, _, _, embedding_svc, vector_store = _make_service()
result = await svc.ingest("job-1")
assert result.doc_id == "doc-1"
assert result.chunks_indexed == 2
assert result.embedding_dimension == 384
# embedding called with 2 texts
embedding_svc.embed.assert_awaited_once()
texts = embedding_svc.embed.call_args[0][0]
assert texts == ["Hello world", "Second chunk"]
# ensure_index called
vector_store.ensure_index.assert_awaited_once()
# delete_document called for idempotency
vector_store.delete_document.assert_awaited_once_with("docling-studio-chunks", "doc-1")
# index_chunks called with IndexedChunk list
vector_store.index_chunks.assert_awaited_once()
chunks_arg = vector_store.index_chunks.call_args[0][1]
assert len(chunks_arg) == 2
assert all(isinstance(c, IndexedChunk) for c in chunks_arg)
assert chunks_arg[0].content == "Hello world"
assert chunks_arg[0].chunk_index == 0
assert chunks_arg[1].chunk_index == 1
async def test_ingest_preserves_bboxes():
# Rebuild to inspect chunks
svc2, _, _, _, vector_store = _make_service()
await svc2.ingest("job-1")
chunks_arg = vector_store.index_chunks.call_args[0][1]
first_chunk = chunks_arg[0]
assert len(first_chunk.bboxes) == 1
bbox = first_chunk.bboxes[0]
assert bbox.page == 1
assert bbox.x == 10.0
assert bbox.y == 20.0
assert bbox.w == 90.0 # right - left = 100 - 10
assert bbox.h == 30.0 # bottom - top = 50 - 20
async def test_ingest_preserves_headings():
svc, _, _, _, vector_store = _make_service()
await svc.ingest("job-1")
chunks = vector_store.index_chunks.call_args[0][1]
assert chunks[0].headings == ["Introduction"]
assert chunks[1].headings == []
async def test_ingest_idempotent_deletes_first():
svc, _, _, _, vector_store = _make_service()
vector_store.delete_document = AsyncMock(return_value=5)
result = await svc.ingest("job-1")
vector_store.delete_document.assert_awaited_once()
assert result.chunks_indexed == 2
async def test_ingest_detects_embedding_dimension():
embeddings = [[0.1] * 768, [0.2] * 768]
svc, *_ = _make_service(embeddings=embeddings)
result = await svc.ingest("job-1")
assert result.embedding_dimension == 768
# ---------------------------------------------------------------------------
# Error cases — job not found
# ---------------------------------------------------------------------------
async def test_ingest_job_not_found():
svc, analysis_repo, *_ = _make_service()
analysis_repo.find_by_id = AsyncMock(return_value=None)
with pytest.raises(IngestionError, match="not found"):
await svc.ingest("missing-job")
async def test_ingest_job_not_completed():
job = _make_job(status=AnalysisStatus.RUNNING)
svc, *_ = _make_service(job=job)
with pytest.raises(IngestionError, match="not COMPLETED"):
await svc.ingest("job-1")
async def test_ingest_job_no_chunks():
job = _make_job(chunks_json=None)
job.chunks_json = None
svc, *_ = _make_service(job=job)
with pytest.raises(IngestionError, match="no chunks"):
await svc.ingest("job-1")
async def test_ingest_job_empty_chunks():
job = _make_job(chunks_json="[]")
svc, *_ = _make_service(job=job)
with pytest.raises(IngestionError, match="empty"):
await svc.ingest("job-1")
async def test_ingest_doc_not_found():
svc, _, document_repo, *_ = _make_service()
document_repo.find_by_id = AsyncMock(return_value=None)
with pytest.raises(IngestionError, match="Document not found"):
await svc.ingest("job-1")
# ---------------------------------------------------------------------------
# Error cases — pipeline failures
# ---------------------------------------------------------------------------
async def test_ingest_embedding_failure():
svc, _, _, embedding_svc, _ = _make_service()
embedding_svc.embed = AsyncMock(side_effect=ConnectionError("embedding service down"))
with pytest.raises(IngestionError, match="Embedding generation failed"):
await svc.ingest("job-1")
async def test_ingest_index_failure():
svc, _, _, _, vector_store = _make_service()
vector_store.ensure_index = AsyncMock(side_effect=RuntimeError("opensearch down"))
with pytest.raises(IngestionError, match="Failed to ensure index"):
await svc.ingest("job-1")
async def test_ingest_bulk_index_failure():
svc, _, _, _, vector_store = _make_service()
vector_store.index_chunks = AsyncMock(side_effect=RuntimeError("bulk failed"))
with pytest.raises(IngestionError, match="Bulk indexing failed"):
await svc.ingest("job-1")
async def test_ingest_embedding_count_mismatch():
svc, _, _, embedding_svc, _ = _make_service()
embedding_svc.embed = AsyncMock(return_value=[[0.1] * 384]) # only 1 instead of 2
with pytest.raises(IngestionError, match="mismatch"):
await svc.ingest("job-1")
# ---------------------------------------------------------------------------
# Delete
# ---------------------------------------------------------------------------
async def test_delete_calls_vector_store():
svc, _, _, _, vector_store = _make_service()
vector_store.delete_document = AsyncMock(return_value=7)
count = await svc.delete("doc-1")
assert count == 7
vector_store.delete_document.assert_awaited_once_with("docling-studio-chunks", "doc-1")
async def test_delete_raises_ingestion_error_on_failure():
svc, _, _, _, vector_store = _make_service()
vector_store.delete_document = AsyncMock(side_effect=RuntimeError("opensearch down"))
with pytest.raises(IngestionError, match="Failed to delete"):
await svc.delete("doc-1")