« prev ^ index » next coverage.py v7.15.4, created at 2026-08-25 07:19 +0800
1"""
2Core document ingestion logic.
4Contains the main ingestion processing functionality.
5"""
7from __future__ import annotations
9import asyncio
10from dataclasses import dataclass
11from pathlib import Path
12from typing import TYPE_CHECKING, Any
14from lexigram.ai.workers.document_ingestion.parser import (
15 DocumentParser,
16 UniversalDocumentParser,
17)
18from lexigram.ai.workers.document_ingestion.types import (
19 ChunkingConfigDict,
20 Document,
21 IngestionResult,
22 IngestionStatus,
23)
24from lexigram.logging import (
25 get_logger,
26)
28if TYPE_CHECKING:
29 from lexigram.ai.workers.document_ingestion.progress import ProgressTracker
30 from lexigram.contracts import VectorStoreProtocol
31 from lexigram.contracts.ai.rag import ChunkProtocol
32 from lexigram.contracts.ai.vector import ChunkerProtocol
34logger = get_logger(__name__)
37@dataclass(slots=True)
38class _ChunkRecord:
39 """Minimal chunk structure used by the ingestion processor."""
41 text: str
42 metadata: dict[str, Any]
43 score: float | None = None
46class _SimpleChunker:
47 """Fallback chunker used when no chunker implementation is injected."""
49 def __init__(self, chunk_size: int = 1000, overlap: int = 200) -> None:
50 normalized_chunk_size = max(1, chunk_size)
51 normalized_overlap = max(0, overlap)
52 self._chunk_size = normalized_chunk_size
53 self._overlap = min(normalized_overlap, normalized_chunk_size - 1)
55 def chunk(
56 self,
57 text: str,
58 metadata: dict[str, Any] | None = None,
59 ) -> list[_ChunkRecord]:
60 if not text:
61 return []
63 step = max(1, self._chunk_size - self._overlap)
64 chunks: list[_ChunkRecord] = []
65 index = 0
67 for start in range(0, len(text), step):
68 end = min(len(text), start + self._chunk_size)
69 chunk_text = text[start:end]
71 if chunk_text.strip():
72 chunk_metadata = dict(metadata or {})
73 chunk_metadata.update(
74 {
75 "chunk_index": index,
76 "start_index": start,
77 "end_index": end,
78 }
79 )
80 chunks.append(_ChunkRecord(text=chunk_text, metadata=chunk_metadata))
81 index += 1
83 if end >= len(text):
84 break
86 return chunks
89class DocumentProcessor:
90 """Handles the core document processing logic."""
92 def __init__(
93 self,
94 vector_store: VectorStoreProtocol,
95 progress_tracker: ProgressTracker,
96 default_chunking_config: ChunkingConfigDict | None = None,
97 chunker: ChunkerProtocol | None = None,
98 document_parser: DocumentParser | None = None,
99 ):
100 """Initialize document processor."""
101 self.vector_store = vector_store
102 self.progress_tracker = progress_tracker
103 self.default_chunking_config = default_chunking_config or {}
105 chunk_size = self.default_chunking_config.get("chunk_size", 1000)
106 overlap = self.default_chunking_config.get("overlap", 200)
107 self._chunker = chunker or _SimpleChunker(
108 chunk_size=int(chunk_size),
109 overlap=int(overlap),
110 )
111 self.document_parser = document_parser or UniversalDocumentParser()
113 async def process_document(
114 self,
115 document_id: str,
116 file_path: str,
117 collection_name: str,
118 parser_name: str | None,
119 metadata: dict[str, Any],
120 batch_size: int = 50,
121 ) -> IngestionResult:
122 """
123 Process a document through the ingestion pipeline.
125 This is the main processing function that handles parsing, chunking, and storing.
126 """
127 start_time = asyncio.get_event_loop().time()
129 try:
130 # Select document parser
131 if parser_name:
132 logger.info(
133 "Custom parser requested but using default",
134 parser_name=parser_name,
135 document_id=document_id,
136 )
138 # Update progress: parsing
139 await self.progress_tracker.update_progress(
140 document_id=document_id,
141 status=IngestionStatus.PARSING,
142 )
144 # Parse document
145 document = await self._parse_document(Path(file_path), metadata)
146 logger.info(
147 "Parsed document",
148 document_id=document_id,
149 pages=len(document.content.split("\n\n")),
150 )
152 # Update progress: chunking
153 await self.progress_tracker.update_progress(
154 document_id=document_id,
155 status=IngestionStatus.CHUNKING,
156 )
158 # Chunk document
159 chunks = await self._chunk_document(document)
160 logger.info(
161 "Chunked document",
162 document_id=document_id,
163 chunks=len(chunks),
164 )
166 # Update total chunks
167 await self.progress_tracker.update_progress(
168 document_id=document_id,
169 total_chunks=len(chunks),
170 )
172 # Update progress: storing
173 await self.progress_tracker.update_progress(
174 document_id=document_id,
175 status=IngestionStatus.STORING,
176 )
178 # Store chunks in batches
179 chunks_created = 0
180 for i in range(0, len(chunks), batch_size):
181 batch = chunks[i : i + batch_size]
183 await self._store_chunks(batch, collection_name)
184 chunks_created += len(batch)
186 # Update progress
187 await self.progress_tracker.update_progress(
188 document_id=document_id,
189 chunks_processed=chunks_created,
190 )
192 logger.debug(
193 "Stored chunk batch",
194 document_id=document_id,
195 batch_size=len(batch),
196 total_chunks=chunks_created,
197 )
199 # Update progress: completed
200 await self.progress_tracker.update_progress(
201 document_id=document_id,
202 status=IngestionStatus.COMPLETED,
203 )
205 duration = asyncio.get_event_loop().time() - start_time
207 logger.info(
208 "Document ingestion completed",
209 document_id=document_id,
210 chunks_created=chunks_created,
211 duration=f"{duration:.2f}s",
212 )
214 return IngestionResult.success_result(
215 document_id=document_id,
216 chunks_created=chunks_created,
217 duration=duration,
218 metadata={"collection": collection_name},
219 )
221 except Exception as e:
222 duration = asyncio.get_event_loop().time() - start_time
223 error_msg = str(e)
225 # Update progress: failed
226 await self.progress_tracker.update_progress(
227 document_id=document_id,
228 error=error_msg,
229 )
231 logger.exception(
232 "Document ingestion failed",
233 document_id=document_id,
234 error=error_msg,
235 )
237 return IngestionResult.failure_result(
238 document_id=document_id,
239 error=error_msg,
240 duration=duration,
241 )
243 async def _parse_document(
244 self,
245 file_path: Path,
246 metadata: dict[str, Any],
247 ) -> Document:
248 """Parse document from file using configured parser."""
249 document = await self.document_parser.parse(file_path)
250 document.metadata.update(metadata)
251 return document
253 async def _chunk_document(self, document: Document) -> list[ChunkProtocol]:
254 """Chunk document into smaller pieces."""
255 raw_chunks = self._chunker.chunk(document.content, metadata=document.metadata)
256 chunks: list[ChunkProtocol] = []
258 for index, chunk in enumerate(raw_chunks):
259 chunk_text = str(getattr(chunk, "text", ""))
260 chunk_metadata = getattr(chunk, "metadata", None)
262 merged_metadata = dict(document.metadata)
263 if isinstance(chunk_metadata, dict):
264 merged_metadata.update(chunk_metadata)
265 merged_metadata.setdefault("chunk_index", index)
267 chunks.append(
268 _ChunkRecord(
269 text=chunk_text,
270 metadata=merged_metadata,
271 score=getattr(chunk, "score", None),
272 )
273 )
275 return chunks
277 async def _store_chunks(
278 self,
279 chunks: list[ChunkProtocol],
280 collection_name: str,
281 ) -> None:
282 """Store chunks in vector store."""
283 texts = [chunk.text for chunk in chunks]
284 metadatas = [chunk.metadata for chunk in chunks]
286 await self.vector_store.add_texts(
287 texts=texts,
288 metadatas=[m or {} for m in metadatas],
289 collection_name=collection_name,
290 )