« prev ^ index » next coverage.py v7.15.4, created at 2026-08-25 07:19 +0800
1"""
2Document ingestion worker.
4Main entry point for document ingestion functionality.
5"""
7from __future__ import annotations
9import asyncio
10from typing import TYPE_CHECKING, Any, cast
12from lexigram.ai.workers.document_ingestion.parser import (
13 DocumentParser,
14 UniversalDocumentParser,
15)
16from lexigram.ai.workers.document_ingestion.processor import DocumentProcessor
17from lexigram.ai.workers.document_ingestion.progress import ProgressTracker
19if TYPE_CHECKING:
20 from pathlib import Path
22 from lexigram.ai.workers.document_ingestion.types import (
23 ChunkingConfigDict,
24 IngestionProgress,
25 IngestionResult,
26 )
27 from lexigram.contracts import VectorStoreProtocol
28 from lexigram.contracts.infra.tasks import (
29 TaskQueueError,
30 TaskQueueProtocol,
31 TaskWorkerProtocol,
32 )
33 from lexigram.result import Result
35from lexigram.contracts.core.health import HealthCheckResult, HealthStatus
36from lexigram.logging import (
37 get_logger,
38)
40HAS_TASK_WORKER = True
42logger = get_logger(__name__)
45class DocumentIngestionWorker:
46 """Document ingestion worker."""
48 def __init__(
49 self,
50 vector_store: VectorStoreProtocol,
51 queue: Any,
52 worker_id: str = "document-ingestion",
53 concurrency: int = 3,
54 default_chunking_config: ChunkingConfigDict | None = None,
55 document_parser: DocumentParser | None = None,
56 allowed_root: Path | None = None,
57 ):
58 """
59 Initialize document ingestion worker.
61 Args:
62 vector_store: Vector store for storing document chunks
63 queue: Task queue for job management
64 worker_id: Unique worker identifier
65 concurrency: Number of concurrent document processing tasks
66 default_chunking_config: Default chunking configuration
67 document_parser: Document parser to use (defaults to
68 UniversalDocumentParser)
69 allowed_root: Optional directory that every ingested source must
70 resolve inside of; forwarded to the default
71 ``UniversalDocumentParser``. Ignored when ``document_parser``
72 is supplied — custom parsers own their own path policy.
73 """
74 self.vector_store = vector_store
75 self.queue = queue
76 self.worker_id = worker_id
77 self.concurrency = concurrency
78 self.default_chunking_config = default_chunking_config or {}
79 self.document_parser = document_parser or UniversalDocumentParser(
80 allowed_root=allowed_root
81 )
83 # Initialize components
84 self.progress_tracker = ProgressTracker()
85 self.processor = DocumentProcessor(
86 vector_store=vector_store,
87 progress_tracker=self.progress_tracker,
88 default_chunking_config=self.default_chunking_config,
89 document_parser=self.document_parser,
90 )
92 # TaskWorker resolver happens at module level (see imports above).
93 # Use a forward-ref to the TaskWorkerProtocol type so static checkers can
94 # validate calls like `.start()` and `.stop()` when the provider is available.
95 self._workers: list[TaskWorkerProtocol] = []
96 self._running = False
98 async def start(self) -> None:
99 """Start the ingestion worker pool."""
100 if self._running:
101 logger.warning("Worker %s already running", self.worker_id)
102 return
104 self._running = True
106 # Create handler registry
107 handlers = {
108 "ingest_document": self._handle_ingest_document,
109 }
111 # Start worker pool (if TaskWorker backend is available)
112 if not HAS_TASK_WORKER:
113 logger.warning(
114 "TaskWorker backend not available; worker pool disabled",
115 worker_id=self.worker_id,
116 )
117 # Mark as not running since no workers will be started
118 self._running = False
119 return
121 # Dynamically resolve concrete worker implementation via DI or factory
122 # For now, we assume if HAS_TASK_WORKER is true, we can get the implementation
123 # usage: container.resolve(TaskWorkerProtocol) or factory pattern
124 # Since the original code instantiated TaskWorker directly, we need a way
125 # to get the concrete class without importing it here.
127 # Strategy: Use DI container to get the factory if available,
128 # otherwise we might still need a runtime import inside the method to instantiate
129 # IF we want to avoid top-level import.
130 # But wait, we want to REMOVE the dependency on concrete class.
132 # Ideally:
133 # factory = container.resolve("task_worker_factory")
134 # worker = factory(...)
136 # Current workaround to satisfy the constraint while keeping functionality:
137 # Import inside the method (still a violation but scoped) OR use a factory.
139 from lexigram.contracts.exceptions import (
140 DependencyError,
141 UnresolvableDependencyError,
142 )
143 from lexigram.di.resolution.context import get_resolver
145 try:
146 # Try to resolve a factory or class
147 resolver = get_resolver(self)
148 if resolver is None:
149 raise ValueError("No resolver found in context")
150 worker_class: type[TaskWorkerProtocol] = await cast(
151 "Any", resolver
152 ).resolve("TaskWorkerClass")
153 except (
154 DependencyError,
155 UnresolvableDependencyError,
156 KeyError,
157 RuntimeError,
158 ValueError,
159 ):
160 # TaskWorker could not be resolved from the DI container.
161 # Without a concrete worker implementation the pool cannot start.
162 logger.warning(
163 "TaskWorker not resolvable from DI container; "
164 "document-ingestion worker pool will not start. "
165 "Register a TaskWorkerProtocol implementation before booting."
166 )
167 self._running = False
168 return
170 for i in range(self.concurrency):
171 worker = worker_class(
172 worker_id=f"{self.worker_id}-{i}",
173 queue=self.queue,
174 handler_registry=handlers,
175 )
176 await worker.start()
177 self._workers.append(worker)
179 logger.info(
180 "Started document ingestion worker pool",
181 worker_id=self.worker_id,
182 concurrency=self.concurrency,
183 )
185 async def stop(self) -> None:
186 """Stop the ingestion worker pool."""
187 if not self._running:
188 return
190 self._running = False
192 # Stop all workers
193 await asyncio.gather(
194 *[worker.stop() for worker in self._workers],
195 return_exceptions=True,
196 )
198 self._workers.clear()
200 logger.info("Stopped document ingestion worker pool", worker_id=self.worker_id)
202 async def ingest_document(
203 self,
204 document_id: str,
205 file_path: Path,
206 collection_name: str,
207 parser: DocumentParser | None = None,
208 chunking_config: ChunkingConfigDict | None = None,
209 metadata: dict[str, Any] | None = None,
210 priority: int = 0,
211 batch_size: int = 50,
212 ) -> str:
213 """
214 Submit document for ingestion.
216 Args:
217 document_id: Unique document identifier
218 file_path: Path to document file
219 collection_name: Vector store collection name
220 parser: Optional document parser (uses worker default if None)
221 chunking_config: Optional chunking configuration
222 metadata: Optional document metadata
223 priority: JobProtocol priority (higher = sooner)
224 batch_size: Chunks per batch
226 Returns:
227 JobProtocol ID for tracking
228 """
229 # Create the job data and enqueue using the TaskQueueProtocol API (avoid
230 # instantiating `JobProtocol` directly so providers that expect (job_type, data)
231 # usage continue to work and mypy doesn't require concrete JobProtocol types.)
232 job_data = {
233 "document_id": document_id,
234 "file_path": str(file_path),
235 "collection_name": collection_name,
236 "parser_name": parser.__class__.__name__ if parser else None,
237 "metadata": metadata or {},
238 "batch_size": batch_size,
239 }
241 # Initialize progress tracking placeholder will be set after enqueue
242 enqueue_result: Result[str, TaskQueueError] = await cast(
243 "TaskQueueProtocol", self.queue
244 ).enqueue(
245 {
246 "name": "ingest_document",
247 "args": (),
248 "kwargs": job_data,
249 "priority": priority,
250 }
251 )
252 if enqueue_result.is_err():
253 msg = f"Failed to enqueue ingestion job: {enqueue_result.unwrap_err()}"
254 raise RuntimeError(msg)
256 job_id = enqueue_result.unwrap()
257 await self.progress_tracker.initialize_progress(job_id, document_id)
259 logger.info(
260 "Submitted document ingestion job",
261 job_id=job_id,
262 document_id=document_id,
263 file_path=str(file_path),
264 )
266 return job_id
268 async def get_progress(self, job_id: str) -> IngestionProgress | None:
269 """Get ingestion progress for job."""
270 return await self.progress_tracker.get_progress(job_id)
272 async def _handle_ingest_document(
273 self,
274 document_id: str,
275 file_path: str,
276 collection_name: str,
277 parser_name: str | None,
278 metadata: dict[str, Any],
279 batch_size: int = 50,
280 ) -> IngestionResult:
281 """
282 Handle document ingestion job.
284 This is the main worker function that processes documents.
285 """
286 return await self.processor.process_document(
287 document_id=document_id,
288 file_path=file_path,
289 collection_name=collection_name,
290 parser_name=parser_name,
291 metadata=metadata,
292 batch_size=batch_size,
293 )
295 def get_stats(self) -> dict[str, Any]:
296 """Get worker statistics."""
297 progress_stats = self.progress_tracker.get_stats()
298 return {
299 "worker_id": self.worker_id,
300 "running": self._running,
301 "concurrency": self.concurrency,
302 "active_workers": len(self._workers),
303 **progress_stats,
304 }
306 async def health_check(self, timeout: float = 5.0) -> HealthCheckResult:
307 """Report the health of this worker.
309 Args:
310 timeout: Unused; present for protocol conformance.
312 Returns:
313 HEALTHY when the worker is running, UNHEALTHY otherwise.
314 """
315 status = HealthStatus.HEALTHY if self._running else HealthStatus.UNHEALTHY
316 stats = self.get_stats()
317 return HealthCheckResult(
318 component=f"worker.document_ingestion.{self.worker_id}",
319 status=status,
320 details=stats,
321 )