Coverage for src/qdrant_loader/core/pipeline/document_pipeline.py: 73%
166 statements
« prev ^ index » next coverage.py v7.15.0, created at 2026-07-20 10:15 +0000
1"""Document processing pipeline that coordinates chunking, embedding, and upserting."""
3import asyncio
4import time
5from collections import defaultdict
6from collections.abc import Awaitable, Callable
7from dataclasses import dataclass
8from typing import Any
10import qdrant_loader_core.graph.registry as _registry # noqa: F401
11from qdrant_loader_core.graph import get_graph_store
12from qdrant_loader_core.graph.extractor.base_extractor import EntityExtractor
14from qdrant_loader.config import get_settings
15from qdrant_loader.core.document import Document
16from qdrant_loader.utils.logging import LoggingConfig
18from .workers import ChunkingWorker, EmbeddingWorker, UpsertWorker
19from .workers.upsert_worker import PipelineResult
21logger = LoggingConfig.get_logger(__name__)
24@dataclass
25class BatchResult:
26 """Result of processing a bounded batch of documents."""
28 success_count: int = 0
29 failure_count: int = 0
30 skipped_count: int = 0
31 successfully_processed_documents: set[str] | None = None
32 failed_document_ids: set[str] | None = None
33 errors: list[str] | None = None
35 def __post_init__(self) -> None:
36 if self.successfully_processed_documents is None:
37 self.successfully_processed_documents = set()
38 if self.failed_document_ids is None:
39 self.failed_document_ids = set()
40 if self.errors is None:
41 self.errors = []
44class DocumentPipeline:
45 """Handles the chunking -> embedding -> upsert pipeline."""
47 def __init__(
48 self,
49 chunking_worker: ChunkingWorker,
50 embedding_worker: EmbeddingWorker,
51 upsert_worker: UpsertWorker,
52 ):
53 self.chunking_worker = chunking_worker
54 self.embedding_worker = embedding_worker
55 self.upsert_worker = upsert_worker
57 async def _process_graph(
58 self,
59 documents: list[Document],
60 current_project_id: str | None = None,
61 ) -> None:
62 """
63 Graph write hook (optional):
64 - Extract entities + relations from documents
65 - Build nodes/edges
66 - Deduplicate
67 - Batch upsert into GraphStore
68 - Must NOT affect ingestion if fails
69 """
71 # ---------------------------
72 # 1. Load config safely
73 # ---------------------------
74 try:
75 settings = get_settings()
76 graph_cfg = getattr(settings.global_config, "graph", None)
77 graph_enabled = (
78 bool(getattr(graph_cfg, "enabled", False)) if graph_cfg else False
79 )
80 except Exception:
81 logger.warning("Graph config not available → graph disabled")
82 return
84 if not graph_enabled:
85 return
87 logger.info("🔄 Graph extraction started (batch size=%s)", len(documents))
89 # ---------------------------
90 # 2. Prepare dedup storage
91 # ---------------------------
92 nodes_dict: dict[str, Any] = {}
93 edges_dict: dict[tuple[str, str, str, str | None], Any] = {}
95 # Optional: group by source_type for efficiency
96 grouped_docs: dict[str, list[Document]] = defaultdict(list)
97 for doc in documents:
98 grouped_docs[doc.source_type].append(doc)
100 # ---------------------------
101 # 3. Extract graph per source_type
102 # ---------------------------
103 for source_type, docs in grouped_docs.items():
104 try:
105 extractor = EntityExtractor.for_source(source_type)
106 if not extractor:
107 logger.warning("No extractor found for source_type=%s", source_type)
108 continue
110 for doc in docs:
111 try:
112 subgraph = await extractor.extract(doc)
114 if getattr(subgraph, "nodes", None):
115 for node in subgraph.nodes:
116 nodes_dict[node.id] = node
118 if getattr(subgraph, "edges", None):
119 for edge in subgraph.edges:
120 kind = (edge.properties or {}).get("kind")
121 edge_key = (
122 edge.source,
123 edge.target,
124 edge.edge_type,
125 kind,
126 )
127 edges_dict[edge_key] = edge
129 except Exception as e:
130 logger.error(
131 "⚠️ Graph extract failed doc_id=%s error=%s",
132 getattr(doc, "id", "<unknown>"),
133 e,
134 exc_info=True,
135 )
137 except Exception as e:
138 logger.error(
139 "⚠️ Extractor failed for source_type=%s error=%s",
140 source_type,
141 e,
142 exc_info=True,
143 )
145 # ---------------------------
146 # 4. Build batches
147 # ---------------------------
148 nodes_batch = list(nodes_dict.values())
149 edges_batch = list(edges_dict.values())
151 # ---------------------------
152 # 5. Attach project context
153 # ---------------------------
154 if current_project_id:
155 for node in nodes_batch:
156 if not hasattr(node, "properties") or node.properties is None:
157 node.properties = {}
158 node.properties["project"] = current_project_id
160 for edge in edges_batch:
161 if not hasattr(edge, "properties") or edge.properties is None:
162 edge.properties = {}
163 edge.properties["project"] = current_project_id
165 # ---------------------------
166 # 6. Batch upsert
167 # ---------------------------
168 if not nodes_batch and not edges_batch:
169 logger.info("Graph extraction result is empty → skip upsert")
170 return
172 try:
173 graph_store = await get_graph_store(
174 **(graph_cfg.store_kwargs() if graph_cfg else {})
175 )
177 if nodes_batch:
178 await graph_store.upsert_nodes_batch(nodes_batch)
180 if edges_batch:
181 await graph_store.upsert_edges_batch(edges_batch)
183 logger.info(
184 "✅ Graph upsert success | nodes=%s edges=%s",
185 len(nodes_batch),
186 len(edges_batch),
187 )
188 logger.debug("=== NODES ===")
189 for i, node in enumerate(nodes_batch, start=1):
190 logger.debug("Node %s: %s", i, node)
192 logger.debug("=== EDGES ===")
193 for i, edge in enumerate(edges_batch, start=1):
194 logger.debug("Edge %s: %s", i, edge)
195 except Exception as e:
196 logger.error(
197 "⚠️ Graph upsert failed (non-fatal): %s",
198 e,
199 exc_info=True,
200 )
202 async def process_batch(
203 self,
204 batch: list[Document],
205 current_project_id: str | None = None,
206 on_document_complete: Callable[[Document, bool], Awaitable[None]] | None = None,
207 ) -> BatchResult:
208 """Process a bounded batch of documents through the pipeline.
210 Args:
211 batch: List of documents to process (bounded size, typically 256)
212 current_project_id: Optional project id context for graph extraction/upsert.
213 on_document_complete: Optional async callback invoked as soon as a
214 document's chunks are all accounted for (success or failure),
215 well before this whole batch finishes. Passed straight through
216 to ``UpsertWorker.process_embedded_chunks``.
218 Returns:
219 BatchResult with processing statistics.
220 """
221 logger.info(f"⚙️ Processing batch of {len(batch)} documents through pipeline")
222 start_time = time.time()
224 try:
225 logger.debug("🔄 Starting chunking phase for batch...")
226 chunking_start = time.time()
227 chunks_iter = self.chunking_worker.process_documents(batch)
229 logger.debug("🔄 Chunking completed, transitioning to embedding phase...")
230 chunking_duration = time.time() - chunking_start
231 logger.debug(f"⏱️ Chunking phase took {chunking_duration:.2f} seconds")
233 embedding_start = time.time()
234 embedded_chunks_iter = self.embedding_worker.process_chunks(chunks_iter)
236 logger.debug("🔄 Embedding phase ready, starting upsert phase...")
238 try:
239 pipeline_result = await asyncio.wait_for(
240 self.upsert_worker.process_embedded_chunks(
241 embedded_chunks_iter, on_document_complete=on_document_complete
242 ),
243 timeout=600.0, # 10 minute timeout per batch
244 )
245 except TimeoutError:
246 logger.error("❌ Batch processing timed out after 10 minutes")
247 return BatchResult(
248 failure_count=len(batch),
249 errors=["Batch processing timed out after 10 minutes"],
250 )
252 total_duration = time.time() - start_time
253 embedding_duration = time.time() - embedding_start
255 logger.debug(
256 f"⏱️ Embedding + Upsert phase took {embedding_duration:.2f} seconds"
257 )
258 logger.info(
259 f"✅ Batch processing completed: {pipeline_result.success_count} chunks, "
260 f"{pipeline_result.error_count} errors in {total_duration:.2f}s"
261 )
263 # add node and edge
264 if (
265 pipeline_result.successfully_processed_documents
266 ): # get list of strings of successful Qdrant upsert docs
267 ok_docs = [
268 doc
269 for doc in batch
270 if doc.id in pipeline_result.successfully_processed_documents
271 ]
272 if ok_docs:
273 await self._process_graph(ok_docs, current_project_id)
275 return BatchResult(
276 success_count=pipeline_result.success_count,
277 failure_count=pipeline_result.error_count,
278 skipped_count=0,
279 successfully_processed_documents=pipeline_result.successfully_processed_documents,
280 failed_document_ids=pipeline_result.failed_document_ids,
281 errors=pipeline_result.errors,
282 )
284 except Exception as e:
285 total_duration = time.time() - start_time
286 logger.error(
287 f"❌ Batch processing failed after {total_duration:.2f} seconds: {e}",
288 exc_info=True,
289 )
290 return BatchResult(
291 failure_count=len(batch),
292 errors=[f"Batch processing failed: {e}"],
293 )
295 async def process_documents(self, documents: list[Document]) -> PipelineResult:
296 """Process documents through the pipeline.
298 Args:
299 documents: List of documents to process
301 Returns:
302 PipelineResult with processing statistics
303 """
304 logger.info(f"⚙️ Processing {len(documents)} documents through pipeline")
305 start_time = time.time()
307 try:
308 logger.info("🔄 Starting chunking phase...")
309 chunking_start = time.time()
310 chunks_iter = self.chunking_worker.process_documents(documents)
312 logger.info("🔄 Chunking completed, transitioning to embedding phase...")
313 chunking_duration = time.time() - chunking_start
314 logger.info(f"⏱️ Chunking phase took {chunking_duration:.2f} seconds")
316 embedding_start = time.time()
317 embedded_chunks_iter = self.embedding_worker.process_chunks(chunks_iter)
319 logger.info("🔄 Embedding phase ready, starting upsert phase...")
321 try:
322 result = await asyncio.wait_for(
323 self.upsert_worker.process_embedded_chunks(embedded_chunks_iter),
324 timeout=3600.0, # 1 hour timeout for the entire pipeline
325 )
326 except TimeoutError:
327 logger.error("❌ Pipeline timed out after 1 hour")
328 result = PipelineResult()
329 result.error_count = len(documents)
330 result.errors = ["Pipeline timed out after 1 hour"]
331 return result
333 total_duration = time.time() - start_time
334 embedding_duration = time.time() - embedding_start
336 logger.info(
337 f"⏱️ Embedding + Upsert phase took {embedding_duration:.2f} seconds"
338 )
339 logger.info(f"⏱️ Total pipeline duration: {total_duration:.2f} seconds")
340 logger.info(
341 f"✅ Pipeline completed: {result.success_count} chunks processed, "
342 f"{result.error_count} errors"
343 )
345 return result
347 except Exception as e:
348 total_duration = time.time() - start_time
349 logger.error(
350 f"❌ Document pipeline failed after {total_duration:.2f} seconds: {e}",
351 exc_info=True,
352 )
353 result = PipelineResult()
354 result.error_count = len(documents)
355 result.errors = [f"Pipeline failed: {e}"]
356 return result