4c3870bf3e
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.
223 lines
7.6 KiB
Python
223 lines
7.6 KiB
Python
"""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
|