Coverage for src/qdrant_loader/core/pipeline/orchestrator.py: 80%
346 statements
« prev ^ index » next coverage.py v7.15.0, created at 2026-07-20 10:15 +0000
1"""Main orchestrator for the ingestion pipeline."""
3import traceback
4from collections.abc import AsyncIterator
5from datetime import datetime
7from qdrant_loader.config import Settings, SourcesConfig
8from qdrant_loader.connectors.base import ConnectorConfigurationError
9from qdrant_loader.connectors.factory import get_connector_instance
10from qdrant_loader.core.document import Document
11from qdrant_loader.core.project_manager import ProjectManager
12from qdrant_loader.core.qdrant_manager import QdrantManager
13from qdrant_loader.core.state.state_change_detector import StateChangeDetector
14from qdrant_loader.core.state.state_manager import StateManager
15from qdrant_loader.core.worker.handlers import PermanentJobError
16from qdrant_loader.utils.logging import LoggingConfig
17from qdrant_loader.utils.sensitive import sanitize_exception_message
19from .document_pipeline import DocumentPipeline
20from .source_filter import SourceFilter
21from .source_processor import SourceProcessor
22from .workers.upsert_worker import PipelineResult
24logger = LoggingConfig.get_logger(__name__)
27def _safe_document_size(doc: Document) -> int:
28 """Best-effort byte size of a document for metrics purposes."""
29 try:
30 return int(doc.metadata.get("size", 0))
31 except (TypeError, ValueError, AttributeError):
32 return 0
35class PipelineComponents:
36 """Container for pipeline components."""
38 def __init__(
39 self,
40 document_pipeline: DocumentPipeline,
41 source_processor: SourceProcessor,
42 source_filter: SourceFilter,
43 state_manager: StateManager,
44 qdrant_manager: QdrantManager,
45 ):
46 self.document_pipeline = document_pipeline
47 self.source_processor = source_processor
48 self.source_filter = source_filter
49 self.state_manager = state_manager
50 self.qdrant_manager = qdrant_manager
53class PipelineOrchestrator:
54 """Main orchestrator for the ingestion pipeline."""
56 def __init__(
57 self,
58 settings: Settings,
59 components: PipelineComponents,
60 project_manager: ProjectManager | None = None,
61 ):
62 self.settings = settings
63 self.components = components
64 self.project_manager = project_manager
65 self.last_pipeline_result = None
67 async def _stream_batches_from_sources(
68 self,
69 filtered_config: SourcesConfig,
70 batch_size: int = 256,
71 since: datetime | None = None,
72 project_id: str | None = None,
73 seen_uris: set[str] | None = None,
74 resume: bool = True,
75 force: bool = False,
76 ) -> AsyncIterator[list[Document]]:
77 """Stream source documents in bounded micro-batches.
79 This helper collects documents from each source type and yields
80 batches of a fixed size, keeping memory usage bounded.
81 """
82 batch: list[Document] = []
84 # Note: previous implementation contained vestigial async inner
85 # helpers `_flush_batch` and `_append_document` which attempted to
86 # yield from inside non-generator contexts. These were dead code
87 # and confusing. Batching is handled inline in `_process_source_type`.
89 async def _process_source_type(source_type_name: str, source_configs):
90 if not source_configs:
91 return
93 # Tracks the checkpoint cursor of the last document with checkpoint
94 # info, so a batch is flushed before crossing a page boundary
95 # (see _process_source_type body below for why this matters).
96 last_cursor_value = None
97 # True once a size-based mid-page flush has fired for the current page.
98 # Used to emit the warning only once per overflowing page.
99 page_has_overflowed = False
101 async def connector_factory_with_checkpoint(src_config):
102 # Determine if we should attempt to resume from a checkpoint
103 checkpoint_cursor = None
104 try:
105 if resume and not force and project_id is not None:
106 # Lazy import to avoid cycles
107 from qdrant_loader.core.state.checkpoint_manager import (
108 CheckpointManager,
109 )
111 async with (
112 await self.components.state_manager.get_session() as session
113 ):
114 cp_mgr = CheckpointManager(session)
115 cp = await cp_mgr.get_checkpoint(
116 project_id, source_type_name, src_config.source
117 )
118 if cp:
119 checkpoint_cursor = cp.cursor_value
120 except Exception:
121 # On any failure retrieving checkpoint, log and continue without it
122 logger.debug(
123 "Checkpoint lookup failed, proceeding without checkpoint",
124 source_type=source_type_name,
125 source=src_config.source,
126 )
128 return get_connector_instance(
129 src_config, checkpoint_cursor=checkpoint_cursor
130 )
132 async for (
133 document
134 ) in self.components.source_processor.stream_source_documents(
135 source_configs,
136 connector_factory_with_checkpoint,
137 source_type_name,
138 since=since,
139 ):
140 # Inject project metadata when running with project context
141 if project_id and self.project_manager:
142 try:
143 document.metadata = (
144 self.project_manager.inject_project_metadata(
145 project_id, document.metadata
146 )
147 )
148 except Exception:
149 # Don't let metadata injection break streaming; log and continue
150 logger.debug(
151 "Project metadata injection failed for document",
152 document_id=document.id,
153 project_id=project_id,
154 )
156 # Track seen URIs for potential post-stream reconciliation
157 if seen_uris is not None:
158 try:
159 uri = f"{document.source_type}:{document.source}:{document.url.rstrip('/') }"
160 seen_uris.add(uri)
161 except Exception:
162 pass
164 # Flush the batch before crossing a checkpoint page boundary so
165 # that a saved checkpoint never covers a partially-upserted page
166 # (a batch never spans two different page cursors).
167 doc_metadata = getattr(document, "metadata", None) or {}
168 cp_info = (
169 doc_metadata.get("__ingestion_checkpoint")
170 if isinstance(doc_metadata, dict)
171 else None
172 )
173 if isinstance(cp_info, dict) and cp_info:
174 cursor_value = cp_info.get("cursor_value")
175 if (
176 last_cursor_value is not None
177 and cursor_value != last_cursor_value
178 and batch
179 ):
180 # Page boundary: all docs for the previous cursor are
181 # accumulated; safe to save the checkpoint now.
182 yield batch.copy()
183 batch.clear()
184 page_has_overflowed = False # reset for the new page
185 last_cursor_value = cursor_value
187 batch.append(document)
188 if len(batch) >= batch_size:
189 # A size-based flush that fires while we are still inside a
190 # page (same cursor_value across docs) must NOT carry
191 # __ingestion_checkpoint. Saving the page token here would
192 # cause resume to skip the page tail on crash (Jira WS-2
193 # regression: 100 issues × avg attachments > batch_size=256).
194 if last_cursor_value is not None:
195 if not page_has_overflowed:
196 logger.warning(
197 "Source page exceeds batch_size; stripping "
198 "__ingestion_checkpoint from mid-page flush to "
199 "prevent resume from skipping the page tail",
200 source_type=source_type_name,
201 page_cursor=last_cursor_value,
202 batch_size=batch_size,
203 )
204 page_has_overflowed = True
205 for doc in batch:
206 doc_meta = getattr(doc, "metadata", None)
207 if isinstance(doc_meta, dict):
208 doc_meta.pop("__ingestion_checkpoint", None)
209 yield batch.copy()
210 batch.clear()
212 if filtered_config.confluence:
213 async for yielded_batch in _process_source_type(
214 "Confluence", filtered_config.confluence
215 ):
216 yield yielded_batch
218 if filtered_config.git:
219 async for yielded_batch in _process_source_type("Git", filtered_config.git):
220 yield yielded_batch
222 if filtered_config.jira:
223 async for yielded_batch in _process_source_type(
224 "Jira", filtered_config.jira
225 ):
226 yield yielded_batch
228 if filtered_config.publicdocs:
229 async for yielded_batch in _process_source_type(
230 "PublicDocs", filtered_config.publicdocs
231 ):
232 yield yielded_batch
234 if filtered_config.localfile:
235 async for yielded_batch in _process_source_type(
236 "LocalFile", filtered_config.localfile
237 ):
238 yield yielded_batch
240 if batch:
241 yield batch
243 async def process_documents(
244 self,
245 sources_config: SourcesConfig | None = None,
246 source_type: str | None = None,
247 source: str | None = None,
248 project_id: str | None = None,
249 force: bool = False,
250 since: datetime | None = None,
251 resume: bool = True,
252 ) -> int:
253 """Main entry point for document processing.
255 Args:
256 sources_config: Sources configuration to use (for backward compatibility)
257 source_type: Filter by source type
258 source: Filter by specific source name
259 project_id: Process documents for a specific project
260 force: Force processing of all documents, bypassing change detection
261 since: Only stream documents updated after this timestamp (connector-level
262 filtering). Connectors that do not yet support time-based filtering will
263 fall back to full fetch with hash-based change detection.
264 resume: Whether to resume from the last checkpoint when available.
266 Returns:
267 Number of documents successfully processed.
268 """
269 logger.info("🚀 Starting document ingestion")
270 self.last_pipeline_result = None
272 try:
273 if sources_config:
274 # Use provided sources config (backward compatibility)
275 logger.debug("Using provided sources configuration")
276 filtered_config = self.components.source_filter.filter_sources(
277 sources_config, source_type, source
278 )
279 current_project_id = None
280 elif project_id:
281 # Use project-specific sources configuration
282 if not self.project_manager:
283 raise ValueError(
284 "Project manager not available for project-specific processing"
285 )
287 project_context = self.project_manager.get_project_context(project_id)
288 if (
289 not project_context
290 or not project_context.config
291 or not project_context.config.sources
292 ):
293 raise ValueError(
294 f"Project '{project_id}' not found or has no configuration"
295 )
297 logger.debug(f"Using project configuration for project: {project_id}")
298 project_sources_config = project_context.config.sources
299 filtered_config = self.components.source_filter.filter_sources(
300 project_sources_config, source_type, source
301 )
302 current_project_id = project_id
303 else:
304 # Process all projects
305 if not self.project_manager:
306 raise ValueError(
307 "Project manager not available and no sources configuration provided"
308 )
310 logger.debug("Processing all projects")
311 return await self._process_all_projects(
312 source_type, source, force, since
313 )
315 # Check if filtered config is empty
316 if source_type and not any(
317 [
318 filtered_config.git,
319 filtered_config.confluence,
320 filtered_config.jira,
321 filtered_config.publicdocs,
322 filtered_config.localfile,
323 ]
324 ):
325 raise ValueError(f"No sources found for type '{source_type}'")
327 # Fail fast when destination collection is unavailable. Without this
328 # preflight, a job with no changed documents would incorrectly return
329 # success even while Qdrant is misconfigured/unreachable.
330 try:
331 await self.components.qdrant_manager.assert_collection_accessible()
332 except Exception as e:
333 raise PermanentJobError(
334 f"Qdrant collection is unavailable: {sanitize_exception_message(e)}"
335 ) from e
337 # Stream documents in bounded micro-batches and process each batch
338 total_documents = 0
339 processed_count = 0
340 aggregated_result = PipelineResult()
341 batch_count = 0
342 counted_success_doc_ids: set[str] = set()
343 counted_failed_doc_ids: set[str] = set()
344 checkpoint_sources_to_clear: set[tuple[str, str]] = set()
345 streamed_checkpoint_sources: set[tuple[str, str]] = set()
347 if not force and not self.components.state_manager._initialized:
348 logger.debug("Initializing state manager for change detection")
349 await self.components.state_manager.initialize()
351 change_detector = None
352 if not force:
353 change_detector = await StateChangeDetector(
354 self.components.state_manager
355 ).__aenter__()
357 try:
358 stream_iter = self._stream_batches_from_sources(
359 filtered_config,
360 256,
361 since,
362 project_id=current_project_id,
363 resume=resume,
364 force=force,
365 )
367 async for batch in stream_iter:
368 total_documents += len(batch)
369 batch_count += 1
371 for doc in batch:
372 metadata = getattr(doc, "metadata", None) or {}
373 cp_info = (
374 metadata.get("__ingestion_checkpoint")
375 if isinstance(metadata, dict)
376 else None
377 )
378 if isinstance(cp_info, dict) and cp_info:
379 streamed_checkpoint_sources.add(
380 (doc.source_type, doc.source)
381 )
383 if not force and change_detector is not None:
384 batch = await change_detector.classify_batch(
385 batch, filtered_config, current_project_id
386 )
388 if not batch:
389 continue
391 async def _persist_as_completed(
392 doc: Document, success: bool
393 ) -> None:
394 if not success:
395 return
396 await self._persist_single_document_state(
397 doc, current_project_id
398 )
400 batch_result = (
401 await self.components.document_pipeline.process_batch(
402 batch,
403 current_project_id,
404 on_document_complete=_persist_as_completed,
405 )
406 )
407 aggregated_result.success_count += batch_result.success_count
408 aggregated_result.error_count += batch_result.failure_count
409 aggregated_result.errors.extend(batch_result.errors)
410 new_success_doc_ids = (
411 batch_result.successfully_processed_documents
412 - counted_success_doc_ids
413 )
414 counted_success_doc_ids.update(new_success_doc_ids)
415 aggregated_result.processed_document_count = len(
416 counted_success_doc_ids
417 )
419 new_failed_doc_ids = (
420 batch_result.failed_document_ids - counted_failed_doc_ids
421 )
422 counted_failed_doc_ids.update(new_failed_doc_ids)
423 aggregated_result.failed_document_count = len(
424 counted_failed_doc_ids
425 )
427 if batch_result.successfully_processed_documents:
428 await self._update_document_states(
429 batch,
430 batch_result.successfully_processed_documents,
431 current_project_id,
432 )
433 batch_counted_doc_ids: set[str] = set()
434 for doc in batch:
435 if (
436 doc.id in new_success_doc_ids
437 and doc.id not in batch_counted_doc_ids
438 ):
439 batch_counted_doc_ids.add(doc.id)
440 processed_count += 1
441 aggregated_result.total_size_bytes += (
442 _safe_document_size(doc)
443 )
444 # Persist checkpoints found on documents (WS-2).
445 # Save once per source with the furthest-advanced cursor
446 # in this batch, not once per document.
447 if resume and current_project_id is not None and not force:
448 try:
449 from qdrant_loader.core.state.checkpoint_manager import (
450 Checkpoint,
451 CheckpointManager,
452 )
454 # Iteration order is the streaming order, so the
455 # last cp_info seen per source is the furthest
456 # along (later cursor overwrites earlier ones).
457 checkpoints_to_save: dict[
458 tuple[str, str], Checkpoint
459 ] = {}
460 for doc in batch:
461 if (
462 doc.id
463 not in batch_result.successfully_processed_documents
464 ):
465 continue
466 metadata = getattr(doc, "metadata", None) or {}
467 cp_info = (
468 metadata.get("__ingestion_checkpoint")
469 if isinstance(metadata, dict)
470 else None
471 )
472 if not isinstance(cp_info, dict) or not cp_info:
473 continue
474 key = (doc.source_type, doc.source)
475 checkpoints_to_save[key] = Checkpoint(
476 project_id=current_project_id,
477 source_type=doc.source_type,
478 source=doc.source,
479 cursor_kind=cp_info.get("cursor_kind"),
480 cursor_value=cp_info.get("cursor_value"),
481 batch_index=cp_info.get("batch_index", 0),
482 )
483 checkpoint_sources_to_clear.add(key)
485 if checkpoints_to_save:
486 async with (
487 await self.components.state_manager.get_session() as session
488 ):
489 cp_mgr = CheckpointManager(session)
490 for checkpoint in checkpoints_to_save.values():
491 await cp_mgr.save_checkpoint(checkpoint)
492 except Exception as e:
493 logger.error(
494 "Failed to persist checkpoint after batch",
495 error=str(e),
496 error_type=type(e).__name__,
497 )
499 if total_documents == 0 and not force:
500 logger.warning(
501 "⚠️ EMPTY SNAPSHOT in non-force mode. About to enter change detection "
502 "which may classify existing corpus as deleted if source API returned partial/null results. "
503 "This is a known risk (WS-3: add explicit snapshot_is_complete signal or per-source enable_deletion_detection). "
504 "Proceeding carefully."
505 )
507 if total_documents == 0 and force:
508 logger.info("✅ No documents found from sources")
509 return 0
511 sources_to_clear = (
512 checkpoint_sources_to_clear | streamed_checkpoint_sources
513 )
515 # On a clean successful run, clear any saved checkpoints for
516 # the project/sources processed (prevents stale resume state).
517 if (
518 resume
519 and current_project_id is not None
520 and not force
521 and aggregated_result.error_count == 0
522 and sources_to_clear
523 ):
524 try:
525 from qdrant_loader.core.state.checkpoint_manager import (
526 CheckpointManager,
527 )
529 async with (
530 await self.components.state_manager.get_session() as session
531 ):
532 cp_mgr = CheckpointManager(session)
533 for stype, src in sources_to_clear:
534 try:
535 await cp_mgr.clear_checkpoint(
536 current_project_id, stype, src
537 )
538 except Exception as e:
539 logger.warning(
540 "Failed to clear checkpoint for source",
541 source_type=stype,
542 source=src,
543 error=str(e),
544 )
545 except Exception as e:
546 logger.warning(
547 "Failed to clear checkpoints after successful run",
548 error=str(e),
549 )
551 if not force and processed_count == 0:
552 self.last_pipeline_result = aggregated_result
553 if aggregated_result.error_count > 0:
554 logger.error(
555 "No documents were successfully processed",
556 error_count=aggregated_result.error_count,
557 )
558 raise PermanentJobError(
559 self._format_indexing_failure_message(aggregated_result)
560 )
561 else:
562 logger.info("No new or updated documents to process")
563 return 0
565 self.last_pipeline_result = aggregated_result
566 if aggregated_result.error_count > 0:
567 logger.error(
568 f"Ingestion completed with failures: "
569 f"{aggregated_result.success_count} succeeded, "
570 f"{aggregated_result.error_count} failed",
571 error_count=aggregated_result.error_count,
572 )
573 raise PermanentJobError(
574 self._format_indexing_failure_message(aggregated_result)
575 )
576 logger.info(
577 f"✅ Ingestion completed: {aggregated_result.success_count} chunks processed successfully"
578 )
579 return processed_count
580 finally:
581 if change_detector is not None:
582 await change_detector.__aexit__(None, None, None)
584 except Exception as e:
585 logger.error(
586 f"❌ Pipeline orchestration failed: {sanitize_exception_message(e)}",
587 error_type=type(e).__name__,
588 sanitized_traceback=sanitize_exception_message(traceback.format_exc()),
589 )
590 raise
592 @staticmethod
593 def _format_indexing_failure_message(result: PipelineResult) -> str:
594 """Build an error message summarizing chunk-level indexing failures."""
595 error_summary = "; ".join(
596 sanitize_exception_message(error) for error in result.errors[:5]
597 )
598 return (
599 f"{result.error_count} chunk(s) failed to index into Qdrant "
600 f"(out of {result.success_count + result.error_count}): {error_summary}"
601 )
603 async def _process_all_projects(
604 self,
605 source_type: str | None = None,
606 source: str | None = None,
607 force: bool = False,
608 since: datetime | None = None,
609 ) -> int:
610 """Process documents from all configured projects."""
611 if not self.project_manager:
612 raise ValueError("Project manager not available")
614 total_processed_count = 0
615 aggregated_result = PipelineResult()
616 failed_projects: list[str] = []
617 project_ids = self.project_manager.list_project_ids()
619 logger.info(f"Processing {len(project_ids)} projects")
621 for project_id in project_ids:
622 try:
623 logger.debug(f"Processing project: {project_id}")
624 project_documents = await self.process_documents(
625 project_id=project_id,
626 source_type=source_type,
627 source=source,
628 force=force,
629 since=since,
630 )
631 project_result = self.last_pipeline_result
632 total_processed_count += project_documents
634 if project_result is not None:
635 aggregated_result.success_count += project_result.success_count
636 aggregated_result.error_count += project_result.error_count
637 aggregated_result.processed_document_count += (
638 project_result.processed_document_count
639 )
640 aggregated_result.failed_document_count += (
641 project_result.failed_document_count
642 )
643 aggregated_result.total_size_bytes += (
644 project_result.total_size_bytes
645 )
646 aggregated_result.errors.extend(project_result.errors)
648 logger.debug(
649 f"Processed {project_documents} documents from project: {project_id}"
650 )
651 except PermanentJobError:
652 # PermanentJobError from process_documents should propagate
653 raise
654 except ConnectorConfigurationError as e:
655 logger.error(
656 f"Configuration error in project {project_id}: "
657 f"{sanitize_exception_message(e)}. "
658 "Skipping this project — check connector settings.",
659 error_type=type(e).__name__,
660 sanitized_traceback=sanitize_exception_message(
661 traceback.format_exc()
662 ),
663 )
664 aggregated_result.errors.append(
665 f"Configuration error in project {project_id}: "
666 f"{sanitize_exception_message(e)}"
667 )
668 failed_projects.append(project_id)
669 continue
670 except Exception as e:
671 safe_error = sanitize_exception_message(e)
672 sanitized_traceback = sanitize_exception_message(traceback.format_exc())
673 aggregated_result.error_count += 1
674 aggregated_result.errors.append(
675 "project_id="
676 f"{project_id}; "
677 "error_type="
678 f"{type(e).__name__}; "
679 "message="
680 f"{safe_error}; "
681 "traceback="
682 f"{sanitized_traceback}"
683 )
684 logger.error(
685 f"Failed to process project {project_id}: {safe_error}",
686 error_type=type(e).__name__,
687 sanitized_traceback=sanitized_traceback,
688 )
689 failed_projects.append(project_id)
690 # Continue processing other projects
691 continue
693 self.last_pipeline_result = aggregated_result
695 total_count = len(project_ids)
696 failed_count = len(failed_projects)
697 success_count = total_count - failed_count
698 if failed_count > 0:
699 logger.warning(
700 f"Completed processing projects: {success_count}/{total_count} succeeded, "
701 f"{failed_count} failed. Check errors above for details.",
702 total_projects=total_count,
703 successful_projects=success_count,
704 failed_projects=failed_count,
705 )
706 else:
707 logger.info(
708 f"Completed processing all projects: {total_processed_count} total documents"
709 )
710 return total_processed_count
712 async def _detect_document_changes(
713 self,
714 documents: list[Document],
715 filtered_config: SourcesConfig,
716 project_id: str | None = None,
717 ) -> list[Document]:
718 """Detect changes in documents and return only new/updated ones."""
720 logger.debug(f"Starting change detection for {len(documents)} documents")
722 try:
723 # Ensure state manager is initialized before use
724 if not self.components.state_manager._initialized:
725 logger.debug("Initializing state manager for change detection")
726 await self.components.state_manager.initialize()
728 async with StateChangeDetector(
729 self.components.state_manager
730 ) as change_detector:
731 changes = await change_detector.detect_changes(
732 documents, filtered_config
733 )
735 new_documents = list(changes.get("new") or [])
736 updated_documents = list(changes.get("updated") or [])
737 deleted_documents = list(changes.get("deleted") or [])
739 logger.info(
740 f"🔍 Change detection: {len(new_documents)} new, "
741 f"{len(updated_documents)} updated, "
742 f"{len(deleted_documents)} deleted"
743 )
745 if deleted_documents:
746 await self._process_deleted_documents(
747 deleted_documents,
748 project_id,
749 )
751 documents_to_process = new_documents + updated_documents
753 if not documents_to_process and deleted_documents:
754 logger.info(
755 "No new or updated documents to process, "
756 "but deleted documents were handled"
757 )
759 return documents_to_process
761 except Exception as e:
762 logger.error(
763 f"Error during change detection: {sanitize_exception_message(e)}",
764 error_type=type(e).__name__,
765 )
766 raise
768 async def _process_deleted_documents(
769 self,
770 deleted_documents: list[Document],
771 project_id: str | None = None,
772 ) -> None:
773 """Process deleted documents by updating state and removing points from Qdrant."""
774 if not deleted_documents:
775 return
777 logger.info(f"Processing {len(deleted_documents)} deleted documents")
779 if not self.components.state_manager._initialized:
780 logger.debug("Initializing state manager for deleted document processing")
781 await self.components.state_manager.initialize()
783 # Use an atomic operation that marks state and deletes points together.
784 try:
785 deleted_ids = (
786 await self.components.state_manager.mark_documents_deleted_atomic(
787 deleted_documents, self.components.qdrant_manager, project_id
788 )
789 )
790 if deleted_ids:
791 logger.info(
792 f"Deleted {len(deleted_ids)} document points from Qdrant and updated state"
793 )
794 except Exception as e:
795 logger.error(
796 f"Failed to process deleted documents atomically: {sanitize_exception_message(e)}",
797 error_type=type(e).__name__,
798 )
799 raise
801 async def _persist_single_document_state(
802 self,
803 document: Document,
804 project_id: str | None = None,
805 ) -> None:
806 """Persist one document's state as soon as it finishes, not at batch end.
808 Streaming batches are bounded at up to 256 documents and state was
809 previously only committed once the *entire* batch finished
810 chunking/embedding/upserting. If the process was interrupted partway
811 through such a batch, documents already durably upserted to Qdrant
812 had no ``DocumentStateRecord`` yet, so a resume would treat them as
813 new and reprocess them. Called from ``UpsertWorker`` the moment a
814 document's last chunk is accounted for, this closes that gap. The
815 end-of-batch ``_update_document_states`` call still runs afterwards
816 as a safety net (idempotent) for anything this callback missed.
817 """
818 try:
819 if not self.components.state_manager._initialized:
820 await self.components.state_manager.initialize()
821 await self.components.state_manager.update_document_states_batch(
822 [document], project_id
823 )
824 except Exception as e:
825 logger.error(
826 f"Failed to persist incremental document state for {document.id}: "
827 f"{sanitize_exception_message(e)}",
828 error_type=type(e).__name__,
829 )
831 async def _update_document_states(
832 self,
833 documents: list[Document],
834 successfully_processed_doc_ids: set,
835 project_id: str | None = None,
836 ):
837 """Update document states for successfully processed documents."""
838 successfully_processed_docs = [
839 doc for doc in documents if doc.id in successfully_processed_doc_ids
840 ]
842 logger.debug(
843 f"Updating document states for {len(successfully_processed_docs)} documents"
844 )
846 # Ensure state manager is initialized before use
847 if not self.components.state_manager._initialized:
848 logger.debug("Initializing state manager for document state updates")
849 await self.components.state_manager.initialize()
851 if not successfully_processed_docs:
852 return
854 # One session/commit for the whole batch instead of one per document
855 # (each document's write is still isolated via a SAVEPOINT, so a
856 # single failure doesn't affect the others' results below).
857 results = await self.components.state_manager.update_document_states_batch(
858 successfully_processed_docs, project_id
859 )
860 for doc, _record, error in results:
861 if error is None:
862 logger.debug(f"Updated document state for {doc.id}")
863 else:
864 logger.error(
865 f"Failed to update document state for {doc.id}: {sanitize_exception_message(error)}",
866 error_type=type(error).__name__,
867 )