1"""
2Document parsing functionality.
3
4Handles parsing of various document formats into unified Document objects.
5"""
6
7from __future__ import annotations
8
9import asyncio
10from dataclasses import dataclass
11from datetime import UTC, datetime
12from pathlib import Path
13import re
14from typing import Any, Protocol, cast
15
16from lexigram.ai.workers.document_ingestion.types import Document
17from lexigram.contracts.ai.exceptions import RAGError
18
19
20@dataclass(slots=True)
21class _LoadedChunk:
22 """Internal chunk representation produced by parser-side loaders."""
23
24 text: str
25 metadata: dict[str, Any]
26 score: float | None = None
27
28
29class DocumentLoader(Protocol):
30 """Protocol for document loaders used by ``UniversalDocumentParser``."""
31
32 async def load(self, source: str | Path) -> list[_LoadedChunk]:
33 """Load a source into chunk-like records."""
34 ...
35
36
37class _TextLoader:
38 async def load(self, source: str | Path) -> list[_LoadedChunk]:
39 path = Path(source)
40 try:
41 content = await asyncio.to_thread(path.read_text, encoding="utf-8")
42 except (OSError, UnicodeDecodeError) as exc:
43 msg = f"Failed to read text file {path}: {exc}"
44 raise RAGError(msg) from exc
45
46 return [
47 _LoadedChunk(
48 text=content,
49 metadata={
50 "source": str(path),
51 "type": "text",
52 },
53 )
54 ]
55
56
57class _MarkdownLoader:
58 async def load(self, source: str | Path) -> list[_LoadedChunk]:
59 path = Path(source)
60 try:
61 content = await asyncio.to_thread(path.read_text, encoding="utf-8")
62 except (OSError, UnicodeDecodeError) as exc:
63 msg = f"Failed to read markdown file {path}: {exc}"
64 raise RAGError(msg) from exc
65
66 chunks: list[_LoadedChunk] = []
67 current_chunk: list[str] = []
68 current_header: str | None = None
69
70 for line in content.split("\n"):
71 if line.startswith("#") and current_chunk:
72 chunks.append(
73 _LoadedChunk(
74 text="\n".join(current_chunk),
75 metadata={
76 "source": str(path),
77 "header": current_header,
78 "type": "markdown",
79 },
80 )
81 )
82 current_chunk = [line]
83 current_header = line.strip("# ")
84 continue
85
86 if line.startswith("#"):
87 current_header = line.strip("# ")
88 current_chunk.append(line)
89
90 if current_chunk:
91 chunks.append(
92 _LoadedChunk(
93 text="\n".join(current_chunk),
94 metadata={
95 "source": str(path),
96 "header": current_header,
97 "type": "markdown",
98 },
99 )
100 )
101
102 return chunks
103
104
105class _HTMLLoader:
106 async def load(self, source: str | Path) -> list[_LoadedChunk]:
107 path = Path(source)
108 try:
109 html = await asyncio.to_thread(path.read_text, encoding="utf-8")
110 except (OSError, UnicodeDecodeError) as exc:
111 msg = f"Failed to read HTML file {path}: {exc}"
112 raise RAGError(msg) from exc
113
114 without_scripts = re.sub(
115 r"<script[\\s\\S]*?</script>",
116 " ",
117 html,
118 flags=re.IGNORECASE,
119 )
120 without_styles = re.sub(
121 r"<style[\\s\\S]*?</style>",
122 " ",
123 without_scripts,
124 flags=re.IGNORECASE,
125 )
126 text = re.sub(r"<[^>]+>", " ", without_styles)
127 text = re.sub(r"\\s+", " ", text).strip()
128
129 return [
130 _LoadedChunk(
131 text=text,
132 metadata={
133 "source": str(path),
134 "type": "html",
135 },
136 )
137 ]
138
139
140class _PDFLoader:
141 async def load(self, source: str | Path) -> list[_LoadedChunk]:
142 path = Path(source)
143
144 def _read_pdf() -> list[_LoadedChunk]:
145 import importlib
146
147 try:
148 pypdf_module = importlib.import_module("pypdf")
149 except ImportError as exc:
150 msg = "PDF loading requires the 'pypdf' package"
151 raise RAGError(msg) from exc
152
153 chunks: list[_LoadedChunk] = []
154 with path.open("rb") as file_obj:
155 reader = pypdf_module.PdfReader(file_obj)
156 for page_index, page in enumerate(reader.pages):
157 page_text = page.extract_text() or ""
158 if page_text.strip():
159 chunks.append(
160 _LoadedChunk(
161 text=page_text,
162 metadata={
163 "source": str(path),
164 "type": "pdf",
165 "page": page_index + 1,
166 },
167 )
168 )
169 return chunks
170
171 try:
172 return cast("list[_LoadedChunk]", await asyncio.to_thread(_read_pdf))
173 except (OSError, ValueError, TypeError, RAGError) as exc:
174 msg = f"Failed to read PDF file {path}: {exc}"
175 raise RAGError(msg) from exc
176
177
178class DocumentParser(Protocol):
179 """Protocol for document parsers."""
180
181 async def parse(self, file_path: Path) -> Document:
182 """Parse document from file."""
183 ...
184
185 async def extract_metadata(self, file_path: Path) -> dict[str, Any]:
186 """Extract metadata from document."""
187 ...
188
189
190class UniversalDocumentParser:
191 """Universal document parser using RAG loaders.
192
193 Supports multiple document formats by delegating to appropriate loaders.
194 """
195
196 def __init__(self, allowed_root: Path | None = None) -> None:
197 """Initialize parser with all supported loaders.
198
199 Args:
200 allowed_root: Optional directory that every parsed source must
201 resolve inside of. When ``None`` (default), no containment
202 check is performed and behavior is unchanged from earlier
203 versions. When set, ``parse`` and ``extract_metadata`` raise
204 ``RAGError`` for any source whose resolved path (symlinks
205 followed) is not inside ``allowed_root``.
206 """
207 self.allowed_root = allowed_root.resolve() if allowed_root is not None else None
208 self._loaders: dict[str, DocumentLoader] = {
209 # Text formats
210 ".txt": _TextLoader(),
211 ".md": _MarkdownLoader(),
212 ".markdown": _MarkdownLoader(),
213 # Binary formats
214 ".pdf": _PDFLoader(),
215 # Web formats
216 ".html": _HTMLLoader(),
217 ".htm": _HTMLLoader(),
218 }
219
220 def _validate_source_within_root(self, file_path: Path) -> None:
221 """Reject sources that resolve outside the configured allowed root.
222
223 Args:
224 file_path: Source path the caller wants to parse or inspect.
225
226 Raises:
227 RAGError: If ``allowed_root`` is set and the resolved path of
228 ``file_path`` is not inside it.
229 """
230 if self.allowed_root is None:
231 return
232 resolved = file_path.resolve()
233 if not resolved.is_relative_to(self.allowed_root):
234 msg = (
235 f"Access denied: {file_path} is outside the allowed root "
236 f"{self.allowed_root}"
237 )
238 raise RAGError(msg)
239
240 async def parse(self, file_path: Path) -> Document:
241 """Parse document from file using appropriate loader.
242
243 Args:
244 file_path: Path to document file
245
246 Returns:
247 Parsed document
248
249 Raises:
250 RAGError: If ``allowed_root`` is set and the resolved
251 ``file_path`` is outside it.
252 ValueError: If file type is not supported
253 """
254 self._validate_source_within_root(file_path)
255
256 if not file_path.exists():
257 msg = f"Document not found: {file_path}"
258 raise FileNotFoundError(msg)
259
260 # Get file extension (lowercase)
261 suffix = file_path.suffix.lower()
262
263 # Get appropriate loader
264 loader = self._loaders.get(suffix)
265 if loader is None:
266 supported = ", ".join(sorted(self._loaders.keys()))
267 msg = f"Unsupported file type '{suffix}'. Supported types: {supported}"
268 raise ValueError(msg)
269
270 # Load document using the appropriate loader
271 chunks = await loader.load(file_path)
272
273 # Combine all chunks into a single document
274 # Preserve chunk metadata in the document metadata
275 combined_text = "\n\n".join(chunk.text for chunk in chunks)
276
277 # Collect metadata from all chunks
278 doc_metadata = {
279 "file_path": str(file_path),
280 "file_name": file_path.name,
281 "file_size": file_path.stat().st_size,
282 "file_type": suffix,
283 "total_chunks": len(chunks),
284 }
285
286 # Add chunk-specific metadata if available
287 if chunks:
288 # Use metadata from first chunk as base
289 first_chunk = chunks[0]
290 if first_chunk.metadata:
291 doc_metadata.update(first_chunk.metadata)
292
293 # Add page info if available (for PDFs)
294 if "page" in first_chunk.metadata:
295 doc_metadata["total_pages"] = max(
296 chunk.metadata.get("page", 0) for chunk in chunks if chunk.metadata
297 )
298
299 return Document(
300 content=combined_text,
301 metadata=doc_metadata,
302 )
303
304 async def extract_metadata(self, file_path: Path) -> dict[str, Any]:
305 """Extract metadata from document without full parsing.
306
307 Args:
308 file_path: Path to document file
309
310 Returns:
311 Document metadata
312
313 Raises:
314 RAGError: If ``allowed_root`` is set and the resolved
315 ``file_path`` is outside it.
316 """
317 self._validate_source_within_root(file_path)
318
319 if not file_path.exists():
320 msg = f"Document not found: {file_path}"
321 raise FileNotFoundError(msg)
322
323 # Basic file metadata
324 stat = file_path.stat()
325 metadata = {
326 "file_path": str(file_path),
327 "file_name": file_path.name,
328 "file_size": stat.st_size,
329 "file_modified": datetime.fromtimestamp(stat.st_mtime, UTC).isoformat(),
330 "file_type": file_path.suffix.lower(),
331 }
332
333 # Try to get additional metadata from loader if available
334 suffix = file_path.suffix.lower()
335 loader = self._loaders.get(suffix)
336
337 if loader is not None:
338 try:
339 # For now, just load and extract from first chunk
340 # In future, loaders could have dedicated metadata methods
341 chunks = await loader.load(file_path)
342 if chunks and chunks[0].metadata:
343 metadata.update(chunks[0].metadata)
344 except (OSError, ValueError, TypeError, AttributeError):
345 # If metadata extraction fails, just return basic metadata
346 pass
347
348 return metadata
349
350 @property
351 def supported_extensions(self) -> list[str]:
352 """Get list of supported file extensions."""
353 return sorted(self._loaders.keys())