Coverage for src/qdrant_loader/core/async_ingestion_pipeline.py: 90%
126 statements
« prev ^ index » next coverage.py v7.15.0, created at 2026-07-20 10:15 +0000
1"""Refactored async ingestion pipeline using the new modular architecture."""
3import asyncio
4from pathlib import Path
6from qdrant_loader.config import Settings, SourcesConfig
7from qdrant_loader.core.monitoring import prometheus_metrics
8from qdrant_loader.core.monitoring.ingestion_metrics import IngestionMonitor
9from qdrant_loader.core.project_manager import ProjectManager
10from qdrant_loader.core.qdrant_manager import QdrantManager
11from qdrant_loader.core.state.state_manager import StateManager
12from qdrant_loader.utils.logging import LoggingConfig
13from qdrant_loader.utils.sensitive import sanitize_exception_message
15from .pipeline import (
16 PipelineComponentsFactory,
17 PipelineConfig,
18 PipelineOrchestrator,
19 ResourceManager,
20)
22logger = LoggingConfig.get_logger(__name__)
25class AsyncIngestionPipeline:
26 """Async ingestion pipeline using modular architecture.
28 This class provides a streamlined interface for the modular pipeline
29 architecture, handling document ingestion and processing workflows.
30 """
32 def __init__(
33 self,
34 settings: Settings,
35 qdrant_manager: QdrantManager,
36 state_manager: StateManager | None = None,
37 max_chunk_workers: int | None = None,
38 max_embed_workers: int | None = None,
39 max_upsert_workers: int | None = None,
40 queue_size: int | None = None,
41 upsert_batch_size: int | None = None,
42 enable_metrics: bool = False,
43 metrics_dir: Path | None = None, # New parameter for workspace support
44 ):
45 """Initialize the async ingestion pipeline.
47 Args:
48 settings: Application settings
49 qdrant_manager: QdrantManager instance
50 state_manager: Optional state manager
52 max_chunk_workers: Maximum number of chunking workers. Defaults to
53 ``settings.global_config.concurrency.max_chunk_workers``.
54 max_embed_workers: Maximum number of embedding workers. Defaults to
55 ``settings.global_config.concurrency.max_embed_workers``.
56 max_upsert_workers: Maximum number of upsert workers. Defaults to
57 ``settings.global_config.concurrency.max_upsert_workers``.
58 queue_size: Queue size for workers. Defaults to
59 ``settings.global_config.concurrency.queue_size``.
60 upsert_batch_size: Batch size for upserts. Defaults to
61 ``settings.global_config.concurrency.upsert_batch_size``.
62 enable_metrics: Whether to enable metrics server
63 metrics_dir: Custom metrics directory (for workspace support)
64 """
65 self.settings = settings
66 self.qdrant_manager = qdrant_manager
68 # Validate that global configuration is available for pipeline operation.
69 if not settings.global_config:
70 raise ValueError(
71 "Global configuration not available. Please check your configuration file."
72 )
74 # Explicit constructor args win; otherwise fall back to the user's
75 # settings.yaml (global.concurrency), so these knobs are actually
76 # reachable from configuration instead of being stuck at a literal.
77 concurrency = settings.global_config.concurrency
78 self.pipeline_config = PipelineConfig(
79 max_chunk_workers=(
80 max_chunk_workers
81 if max_chunk_workers is not None
82 else concurrency.max_chunk_workers
83 ),
84 max_embed_workers=(
85 max_embed_workers
86 if max_embed_workers is not None
87 else concurrency.max_embed_workers
88 ),
89 max_upsert_workers=(
90 max_upsert_workers
91 if max_upsert_workers is not None
92 else concurrency.max_upsert_workers
93 ),
94 queue_size=queue_size if queue_size is not None else concurrency.queue_size,
95 upsert_batch_size=(
96 upsert_batch_size
97 if upsert_batch_size is not None
98 else concurrency.upsert_batch_size
99 ),
100 enable_metrics=enable_metrics,
101 )
103 # Create resource manager to handle cleanup and signal handling.
104 self.resource_manager = ResourceManager()
105 self.resource_manager.register_signal_handlers()
107 # Create state manager instance if not provided by caller.
108 self.state_manager = state_manager or StateManager(
109 settings.global_config.state_management
110 )
112 # Initialize project manager to support multi-project configurations.
113 if not settings.global_config.qdrant:
114 raise ValueError(
115 "Qdrant configuration is required for project manager initialization"
116 )
118 self.project_manager = ProjectManager(
119 projects_config=settings.projects_config,
120 global_collection_name=settings.global_config.qdrant.collection_name,
121 )
123 # Create pipeline components using factory
124 factory = PipelineComponentsFactory()
125 self.components = factory.create_components(
126 settings=settings,
127 config=self.pipeline_config,
128 qdrant_manager=qdrant_manager,
129 state_manager=self.state_manager,
130 resource_manager=self.resource_manager,
131 )
133 # Create orchestrator with project manager support
134 self.orchestrator = PipelineOrchestrator(
135 settings, self.components, self.project_manager
136 )
138 # Initialize performance monitor with custom or default metrics directory
139 if metrics_dir:
140 # Use provided metrics directory (workspace mode)
141 # Accept both Path and str inputs
142 final_metrics_dir = (
143 metrics_dir if isinstance(metrics_dir, Path) else Path(metrics_dir)
144 )
145 else:
146 # Use default metrics directory
147 final_metrics_dir = Path.cwd() / "metrics"
149 final_metrics_dir.mkdir(parents=True, exist_ok=True)
150 logger.info(f"Initializing metrics directory at {final_metrics_dir}")
151 self.monitor = IngestionMonitor(str(final_metrics_dir.absolute()))
153 # Start metrics server if enabled
154 if enable_metrics:
155 prometheus_metrics.start_metrics_server()
157 logger.info("AsyncIngestionPipeline initialized with new modular architecture")
159 # Track cleanup state to prevent duplicate cleanup
160 self._cleanup_performed = False
162 async def initialize(self):
163 """Initialize the pipeline."""
164 logger.debug("Starting pipeline initialization")
166 try:
167 # Ensure the Qdrant collection and its payload indexes exist.
168 # create_collection() is idempotent: it returns early when the
169 # collection already exists, but now also ensures indexes are
170 # present on existing collections (required for filter-based deletes).
171 await asyncio.to_thread(self.qdrant_manager.create_collection)
173 # Initialize state manager first
174 if not self.state_manager.is_initialized:
175 logger.debug("Initializing state manager")
176 await self.state_manager.initialize()
178 # Initialize project manager
179 if not self.project_manager._initialized:
180 logger.debug("Initializing project manager")
181 # Prefer direct use of session factory to match existing tests/mocks
182 session_factory = getattr(self.state_manager, "_session_factory", None)
183 if session_factory is None:
184 logger.error(
185 "State manager session factory is not available during initialization",
186 suggestion="Check database configuration and ensure proper state manager setup",
187 )
188 raise RuntimeError("State manager session factory is not available")
190 try:
191 async with session_factory() as session: # type: ignore
192 await self.project_manager.initialize(session)
193 logger.debug("Project manager initialization completed")
195 except Exception as e:
196 logger.error(
197 "Failed to initialize project manager during pipeline startup",
198 error=sanitize_exception_message(e),
199 error_type=type(e).__name__,
200 suggestion="Check database connectivity and project configuration",
201 )
202 raise
203 except Exception as e:
204 logger.error(
205 "Pipeline initialization failed during startup sequence",
206 error=sanitize_exception_message(e),
207 error_type=type(e).__name__,
208 suggestion="Check configuration, database connectivity, and system resources",
209 )
210 raise
212 async def process_documents(
213 self,
214 sources_config: SourcesConfig | None = None,
215 source_type: str | None = None,
216 source: str | None = None,
217 project_id: str | None = None,
218 force: bool = False,
219 resume: bool = True,
220 ) -> int:
221 """Process documents from all configured sources.
223 Args:
224 sources_config: Sources configuration to use (deprecated, use project_id instead)
225 source_type: Filter by source type
226 source: Filter by specific source name
227 project_id: Process documents for a specific project
228 force: Force processing of all documents, bypassing change detection
230 Returns:
231 Number of documents successfully processed.
232 """
233 # Ensure the pipeline is initialized
234 await self.initialize()
236 # Reset metrics for new run
237 self.monitor.clear_metrics()
238 self.monitor.start_operation(
239 "ingestion_process",
240 metadata={
241 "source_type": source_type,
242 "source": source,
243 "project_id": project_id,
244 "force": force,
245 },
246 )
248 processed_count = (
249 0 # Initialize to avoid UnboundLocalError in exception handler
250 )
251 try:
252 logger.debug("Starting document processing with new pipeline architecture")
254 # Use the orchestrator to process documents with project support
255 processed_count = await self.orchestrator.process_documents(
256 sources_config=sources_config,
257 source_type=source_type,
258 source=source,
259 project_id=project_id,
260 force=force,
261 resume=resume,
262 )
264 # Update metrics
265 if processed_count:
266 pipeline_result = getattr(
267 self.orchestrator, "last_pipeline_result", None
268 )
269 total_chunks = getattr(pipeline_result, "success_count", 0)
270 total_size_bytes = getattr(pipeline_result, "total_size_bytes", 0)
272 self.monitor.start_batch(
273 "document_batch",
274 batch_size=processed_count,
275 metadata={
276 "source_type": source_type,
277 "source": source,
278 "project_id": project_id,
279 "force": force,
280 },
281 )
282 # Note: Success/error counts are handled internally by the new architecture
283 self.monitor.end_batch(
284 "document_batch",
285 processed_count,
286 0,
287 [],
288 total_chunks=total_chunks,
289 total_size_bytes=total_size_bytes,
290 )
292 self.monitor.end_operation("ingestion_process")
294 logger.debug(
295 f"Document processing completed. Processed {processed_count} documents"
296 )
297 return processed_count
299 except Exception as e:
300 safe_error = sanitize_exception_message(e)
301 logger.error(
302 "Document processing pipeline failed during ingestion",
303 error=safe_error,
304 error_type=type(e).__name__,
305 documents_attempted=processed_count,
306 suggestion="Check data source connectivity, document formats, and system resources",
307 )
308 self.monitor.end_operation(
309 "ingestion_process", success=False, error=safe_error
310 )
311 raise
313 async def cleanup(self):
314 """Clean up resources."""
315 if self._cleanup_performed:
316 return
318 logger.info("Cleaning up pipeline resources")
319 self._cleanup_performed = True
321 try:
322 # Save metrics
323 if hasattr(self, "monitor"):
324 self.monitor.save_metrics()
326 # Stop metrics server
327 try:
328 prometheus_metrics.stop_metrics_server()
329 except Exception as e:
330 logger.warning(
331 f"Error stopping metrics server: {sanitize_exception_message(e)}"
332 )
334 # Use resource manager for cleanup
335 if hasattr(self, "resource_manager"):
336 await self.resource_manager.cleanup()
338 logger.info("Pipeline cleanup completed")
339 except Exception as e:
340 logger.error(
341 f"Error during pipeline cleanup: {sanitize_exception_message(e)}"
342 )
344 def __del__(self):
345 """Destructor to ensure cleanup."""
346 try:
347 # Can't await in __del__, so use the sync cleanup method
348 self._sync_cleanup()
349 except Exception as e:
350 logger.error(
351 f"Error in destructor cleanup: {sanitize_exception_message(e)}"
352 )
354 def _sync_cleanup(self):
355 """Synchronous cleanup for destructor and signal handlers."""
356 if self._cleanup_performed:
357 return
359 logger.info("Cleaning up pipeline resources (sync)")
360 self._cleanup_performed = True
362 # Save metrics
363 try:
364 if hasattr(self, "monitor"):
365 self.monitor.save_metrics()
366 except Exception as e:
367 logger.error(f"Error saving metrics: {sanitize_exception_message(e)}")
369 # Stop metrics server
370 try:
371 prometheus_metrics.stop_metrics_server()
372 except Exception as e:
373 logger.error(
374 f"Error stopping metrics server: {sanitize_exception_message(e)}"
375 )
377 # Use resource manager sync cleanup
378 try:
379 if hasattr(self, "resource_manager"):
380 self.resource_manager._cleanup()
381 except Exception as e:
382 logger.error(
383 f"Error in resource manager cleanup: {sanitize_exception_message(e)}"
384 )
386 logger.info("Pipeline cleanup completed (sync)")