Coverage for src/qdrant_loader/core/pipeline/workers/upsert_worker.py: 99%
175 statements
« prev ^ index » next coverage.py v7.15.0, created at 2026-07-20 10:15 +0000
1"""Upsert worker for upserting embedded chunks to Qdrant."""
3import asyncio
4from collections import Counter
5from collections.abc import AsyncIterator, Awaitable, Callable
6from typing import Any
8from qdrant_client.http import models
10from qdrant_loader.core.monitoring import prometheus_metrics
11from qdrant_loader.core.qdrant_manager import QdrantManager
12from qdrant_loader.utils.logging import LoggingConfig
14from .base_worker import BaseWorker
16logger = LoggingConfig.get_logger(__name__)
19class PipelineResult:
20 """Result of pipeline processing."""
22 def __init__(self):
23 self.success_count: int = 0
24 self.error_count: int = 0
25 self.successfully_processed_documents: set[str] = set()
26 self.failed_document_ids: set[str] = set()
27 self.errors: list[str] = []
28 self.processed_document_count: int = 0
29 self.failed_document_count: int = 0
30 self.total_size_bytes: int = 0
33class UpsertWorker(BaseWorker):
34 """Handles upserting embedded chunks to Qdrant."""
36 def __init__(
37 self,
38 qdrant_manager: QdrantManager,
39 batch_size: int,
40 max_workers: int = 4,
41 queue_size: int = 1000,
42 shutdown_event: asyncio.Event | None = None,
43 ):
44 super().__init__(max_workers, queue_size)
45 self.qdrant_manager = qdrant_manager
46 self.batch_size = batch_size
47 self.shutdown_event = shutdown_event or asyncio.Event()
49 def _handle_duplicate_chunk_ids(
50 self,
51 batch: list[tuple[Any, list[float]]],
52 batch_chunk_id_counts: Counter,
53 duplicate_chunk_ids: set[str],
54 same_batch_duplicates: set[str],
55 cross_batch_duplicates: set[str],
56 new_chunk_ids: set[str],
57 result: PipelineResult,
58 errors: list[str],
59 ) -> None:
60 """Log/record duplicate chunk IDs and their error-count impact.
62 Whether a duplicate-affected document ends up in
63 ``successfully_processed_documents`` or ``failed_document_ids`` is
64 decided by the per-document completion tracking in
65 ``process_embedded_chunks`` (via ``note_chunk_outcome``), not here.
66 """
67 if not duplicate_chunk_ids:
68 return
70 duplicate_doc_ids = set()
71 for chunk, _ in batch:
72 if str(chunk.id) in duplicate_chunk_ids:
73 parent_doc = chunk.metadata.get("parent_document")
74 if parent_doc:
75 duplicate_doc_ids.add(parent_doc.id)
77 same_batch_duplicate_occurrences = sum(
78 count - 1 for count in batch_chunk_id_counts.values() if count > 1
79 )
80 total_duplicate_impact = len(duplicate_doc_ids)
81 duplicate_chunk_attempts = len(batch) - len(new_chunk_ids)
83 logger.warning(
84 "Detected chunk ID collisions during upsert; existing points will be overwritten",
85 duplicate_count=len(duplicate_chunk_ids),
86 same_batch_duplicate_count=len(same_batch_duplicates),
87 same_batch_duplicate_occurrences=same_batch_duplicate_occurrences,
88 cross_batch_duplicate_count=len(cross_batch_duplicates),
89 affected_documents=total_duplicate_impact,
90 )
91 errors.append(
92 "Detected duplicate chunk IDs during upsert: "
93 f"{len(cross_batch_duplicates)} cross-batch IDs and "
94 f"{same_batch_duplicate_occurrences} same-batch duplicate occurrences "
95 f"across {len(same_batch_duplicates)} IDs affecting {total_duplicate_impact} document(s): "
96 f"{sorted(duplicate_doc_ids)}"
97 )
98 result.error_count += duplicate_chunk_attempts
100 async def process(
101 self, batch: list[tuple[Any, list[float]]]
102 ) -> tuple[int, int, set[str], list[str]]:
103 """Process a batch of embedded chunks.
105 Args:
106 batch: List of (chunk, embedding) tuples
108 Returns:
109 Tuple of (success_count, error_count, successful_doc_ids, errors)
110 """
111 if not batch:
112 return 0, 0, set(), []
114 success_count = 0
115 error_count = 0
116 successful_doc_ids = set()
117 errors = []
119 try:
120 with prometheus_metrics.UPSERT_DURATION.time():
121 # QdrantManager.build_point_vector owns the dense / dense+sparse
122 # decision and has its own dense-only fallback on encode failure,
123 # so no defensive wrapper is needed here.
124 points = [
125 models.PointStruct(
126 id=chunk.id,
127 vector=self.qdrant_manager.build_point_vector(
128 embedding, chunk.content
129 ),
130 payload={
131 "content": chunk.content,
132 "contextual_content": chunk.contextual_content,
133 "metadata": {
134 k: v
135 for k, v in chunk.metadata.items()
136 if k != "parent_document"
137 },
138 "source": chunk.source,
139 "source_type": chunk.source_type,
140 "created_at": chunk.created_at.isoformat(),
141 "updated_at": (
142 getattr(
143 chunk, "updated_at", chunk.created_at
144 ).isoformat()
145 if hasattr(chunk, "updated_at")
146 else chunk.created_at.isoformat()
147 ),
148 "title": getattr(
149 chunk, "title", chunk.metadata.get("title", "")
150 ),
151 "url": getattr(chunk, "url", chunk.metadata.get("url", "")),
152 "document_id": chunk.metadata.get(
153 "parent_document_id", chunk.id
154 ),
155 },
156 )
157 for chunk, embedding in batch
158 ]
160 await self.qdrant_manager.upsert_points(points)
161 prometheus_metrics.INGESTED_DOCUMENTS.inc(len(points))
162 success_count = len(points)
164 # Mark parent documents as successfully processed
165 for chunk, _ in batch:
166 parent_doc = chunk.metadata.get("parent_document")
167 if parent_doc:
168 successful_doc_ids.add(parent_doc.id)
170 except Exception as e:
171 for chunk, _ in batch:
172 logger.error(f"Upsert failed for chunk {chunk.id}: {e}")
173 # Mark parent document as failed
174 parent_doc = chunk.metadata.get("parent_document")
175 if parent_doc:
176 successful_doc_ids.discard(parent_doc.id) # Remove if it was added
177 errors.append(f"Upsert failed for chunk {chunk.id}: {e}")
178 error_count = len(batch)
180 return success_count, error_count, successful_doc_ids, errors
182 def _reserve_chunk_ids(
183 self, batch: list[tuple[Any, list[float]]], seen_chunk_ids: set[str]
184 ) -> dict[str, Any]:
185 """Compute duplicate-chunk-id bookkeeping for a batch and reserve its new IDs.
187 This runs synchronously (no ``await``) at batch-formation time, before
188 the batch is handed off for concurrent upserting. That ordering is
189 what keeps duplicate detection correct once batches are processed
190 concurrently: reservations happen strictly in submission order, so two
191 in-flight batches can never both believe the same chunk id is new.
192 """
193 batch_chunk_id_list = [str(chunk.id) for chunk, _ in batch]
194 batch_chunk_ids = set(batch_chunk_id_list)
195 batch_chunk_id_counts = Counter(batch_chunk_id_list)
196 same_batch_duplicates = {
197 chunk_id for chunk_id, count in batch_chunk_id_counts.items() if count > 1
198 }
199 cross_batch_duplicates = batch_chunk_ids & seen_chunk_ids
200 duplicate_chunk_ids = cross_batch_duplicates | same_batch_duplicates
201 new_chunk_ids = batch_chunk_ids - seen_chunk_ids - same_batch_duplicates
203 # Reserve now so the next batch's cross-batch check sees it, regardless
204 # of how long this batch's upsert takes to actually complete.
205 seen_chunk_ids.update(new_chunk_ids)
207 return {
208 "batch_chunk_id_counts": batch_chunk_id_counts,
209 "same_batch_duplicates": same_batch_duplicates,
210 "cross_batch_duplicates": cross_batch_duplicates,
211 "duplicate_chunk_ids": duplicate_chunk_ids,
212 "new_chunk_ids": new_chunk_ids,
213 }
215 @staticmethod
216 def _note_chunk_outcome(
217 chunk: Any,
218 failed: bool,
219 result: PipelineResult,
220 doc_totals: dict[str, int],
221 doc_seen: dict[str, int],
222 doc_failed: dict[str, bool],
223 ) -> bool | None:
224 """Record one chunk's fate and finalize its document once complete.
226 A document is only added to ``successfully_processed_documents`` once
227 *every* chunk chunking produced for it has been accounted for (via
228 ``parent_document_total_chunks``) with none failed — a document isn't
229 "done" just because the first batch containing one of its chunks
230 happened to succeed. If any chunk failed (embedding failure, upsert
231 failure, or duplicate-id collision), the document lands in
232 ``failed_document_ids`` instead so a later incremental run retries it.
234 Returns ``None`` while the document still has outstanding chunks, or
235 the document's final success flag (``True``/``False``) the moment it
236 completes — callers use this to persist document state immediately
237 instead of waiting for the whole streaming batch to finish.
238 """
239 parent_doc = chunk.metadata.get("parent_document")
240 if not parent_doc:
241 return None
243 doc_id = parent_doc.id
244 total = chunk.metadata.get("parent_document_total_chunks") or 1
245 doc_totals.setdefault(doc_id, total)
247 seen = doc_seen.get(doc_id, 0) + 1
248 doc_seen[doc_id] = seen
249 if failed:
250 doc_failed[doc_id] = True
252 if seen >= doc_totals[doc_id]:
253 if doc_failed.get(doc_id):
254 result.failed_document_ids.add(doc_id)
255 result.successfully_processed_documents.discard(doc_id)
256 return False
257 else:
258 result.successfully_processed_documents.add(doc_id)
259 return True
260 return None
262 def _merge_batch_outcome(
263 self,
264 batch: list[tuple[Any, list[float]]],
265 dedup: dict[str, Any],
266 outcome: tuple[int, int, set[str], list[str]],
267 result: PipelineResult,
268 seen_chunk_ids: set[str],
269 doc_totals: dict[str, int],
270 doc_seen: dict[str, int],
271 doc_failed: dict[str, bool],
272 ) -> list[tuple[Any, bool]]:
273 """Fold one batch's process() outcome into the running PipelineResult.
275 Returns the documents that just completed as a result of this batch
276 (parent document object, success flag), so the caller can persist
277 their state right away instead of waiting for the whole streaming
278 batch to finish.
279 """
280 success_count, error_count, successful_doc_ids, errors = outcome
282 if success_count > 0:
283 if dedup["duplicate_chunk_ids"]:
284 self._handle_duplicate_chunk_ids(
285 batch=batch,
286 batch_chunk_id_counts=dedup["batch_chunk_id_counts"],
287 duplicate_chunk_ids=dedup["duplicate_chunk_ids"],
288 same_batch_duplicates=dedup["same_batch_duplicates"],
289 cross_batch_duplicates=dedup["cross_batch_duplicates"],
290 new_chunk_ids=dedup["new_chunk_ids"],
291 result=result,
292 errors=errors,
293 )
294 result.success_count += len(dedup["new_chunk_ids"])
295 else:
296 # Nothing was actually written; release the reservation so a
297 # later batch with the same id isn't wrongly treated as a dup.
298 seen_chunk_ids.difference_update(dedup["new_chunk_ids"])
300 result.error_count += error_count
301 result.errors.extend(errors)
303 finalized: list[tuple[Any, bool]] = []
304 for chunk, _ in batch:
305 chunk_failed = (
306 success_count == 0 or str(chunk.id) in dedup["duplicate_chunk_ids"]
307 )
308 outcome_flag = self._note_chunk_outcome(
309 chunk, chunk_failed, result, doc_totals, doc_seen, doc_failed
310 )
311 if outcome_flag is not None:
312 parent_doc = chunk.metadata.get("parent_document")
313 if parent_doc:
314 finalized.append((parent_doc, outcome_flag))
315 return finalized
317 async def process_embedded_chunks(
318 self,
319 embedded_chunks: AsyncIterator[tuple[Any, list[float] | None]],
320 on_document_complete: Callable[[Any, bool], Awaitable[None]] | None = None,
321 ) -> PipelineResult:
322 """Upsert embedded chunks to Qdrant.
324 The number of batches dispatched-and-not-yet-merged is capped at
325 ``max_workers``: once that many are in flight, the oldest is awaited
326 before a new one is created. This bounds memory use (not just upsert
327 concurrency) when embeddings arrive far faster than Qdrant upserts
328 complete, while still merging outcomes in submission order. A chunk
329 arriving with ``embedding is None`` (embedding failed upstream) is
330 never sent to Qdrant, but is still accounted for so its parent
331 document doesn't get marked successful with pieces missing.
333 Args:
334 embedded_chunks: AsyncIterator of (chunk, embedding) tuples
335 on_document_complete: Optional async callback invoked with
336 ``(parent_document, success)`` the moment a document's last
337 chunk is accounted for — before the whole iterator finishes.
338 Lets callers persist document state incrementally, so a
339 mid-run interruption doesn't lose the state of documents
340 that were already fully upserted. Callback failures are
341 logged and never abort the pipeline.
343 Returns:
344 PipelineResult with processing statistics
345 """
347 async def _notify(parent_doc: Any, success: bool) -> None:
348 if on_document_complete is None:
349 return
350 try:
351 await on_document_complete(parent_doc, success)
352 except Exception as e:
353 logger.error(
354 "on_document_complete callback failed",
355 document_id=getattr(parent_doc, "id", None),
356 error=str(e),
357 error_type=type(e).__name__,
358 )
360 logger.debug("UpsertWorker started")
361 logger.info("🔄 Starting upsert processing", max_workers=self.max_workers)
362 result = PipelineResult()
363 seen_chunk_ids: set[str] = set()
364 doc_totals: dict[str, int] = {}
365 doc_seen: dict[str, int] = {}
366 doc_failed: dict[str, bool] = {}
367 batch: list[tuple[Any, list[float]]] = []
368 pending: list[
369 tuple[asyncio.Task, list[tuple[Any, list[float]]], dict[str, Any]]
370 ] = []
372 async def drain_oldest() -> None:
373 task, task_batch, dedup = pending.pop(0)
374 outcome = await task
375 finalized = self._merge_batch_outcome(
376 task_batch,
377 dedup,
378 outcome,
379 result,
380 seen_chunk_ids,
381 doc_totals,
382 doc_seen,
383 doc_failed,
384 )
385 for parent_doc, success in finalized:
386 await _notify(parent_doc, success)
388 def dispatch(batch_to_dispatch: list[tuple[Any, list[float]]]) -> None:
389 dedup = self._reserve_chunk_ids(batch_to_dispatch, seen_chunk_ids)
390 task = asyncio.create_task(self.process_with_semaphore(batch_to_dispatch))
391 pending.append((task, batch_to_dispatch, dedup))
393 try:
394 async for chunk, embedding in embedded_chunks:
395 if self.shutdown_event.is_set():
396 logger.debug("UpsertWorker exiting due to shutdown")
397 break
399 if embedding is None:
400 result.error_count += 1
401 result.errors.append(
402 f"Embedding failed for chunk {chunk.id}, skipped upsert"
403 )
404 outcome_flag = self._note_chunk_outcome(
405 chunk, True, result, doc_totals, doc_seen, doc_failed
406 )
407 if outcome_flag is not None:
408 parent_doc = chunk.metadata.get("parent_document")
409 if parent_doc:
410 await _notify(parent_doc, outcome_flag)
411 continue
413 batch.append((chunk, embedding))
415 # Dispatch batch when it reaches the desired size
416 if len(batch) >= self.batch_size:
417 batch_to_submit = batch
418 batch = []
420 # Keep at most max_workers batches in flight so memory
421 # use stays bounded, not just upsert concurrency.
422 if len(pending) >= self.max_workers:
423 await drain_oldest()
425 dispatch(batch_to_submit)
427 # Dispatch any remaining chunks in the final batch
428 if batch and not self.shutdown_event.is_set():
429 dispatch(batch)
431 while pending:
432 await drain_oldest()
434 except asyncio.CancelledError:
435 logger.debug("UpsertWorker cancelled")
436 for task, _, _ in pending:
437 if not task.done():
438 task.cancel()
439 raise
440 finally:
441 logger.debug("UpsertWorker exited")
443 return result