Coverage for /home/admin/Documents/AI/applications/lexigram-dev/lexigram/experimental/ai/lexigram-ai-workers/src/lexigram/ai/workers/document_ingestion/parser.py: 24%

139 statements  

« prev     ^ index     » next       coverage.py v7.15.4, created at 2026-08-25 07:19 +0800

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())