Coverage for /home/admin/Documents/AI/applications/lexigram-dev/lexigram/experimental/ai/lexigram-ai-rag/src/lexigram/ai/rag/loaders/core.py: 20%

1"""Document loaders for RAG. 

2 

3Load documents from various sources (PDF, HTML, Markdown, JSON, CSV, etc.). 

4""" 

5 

6from __future__ import annotations 

7 

8import asyncio 

9import csv 

10import io 

11from pathlib import Path 

12from typing import Any 

13 

14from lexigram.ai.rag.chunking.types import Chunk 

15from lexigram.ai.rag.types import RAGError 

16from lexigram.serialization import dumps_str 

17from lexigram.serialization import loads as _loads 

18 

19try: 

20 import aiofiles 

21 

22 HAS_AIOFILES = True 

23except ImportError: 

24 HAS_AIOFILES = False 

25 

26 

27class AbstractDocumentLoader: 

28 """Base class for document loaders.""" 

29 

30 async def load(self, source: str | Path) -> list[Chunk]: 

31 """Load documents from source. 

32 

33 Args: 

34 source: Document source (file path, URL, etc.) 

35 

36 Returns: 

37 List of document chunks 

38 

39 Raises: 

40 RAGError: If loading fails 

41 """ 

42 raise NotImplementedError 

43 

44 

45class TextLoader(AbstractDocumentLoader): 

46 """Load plain text files. 

47 

48 Example: 

49 >>> loader = TextLoader() 

50 >>> chunks = await loader.load("document.txt") 

51 """ 

52 

53 async def load(self, source: str | Path) -> list[Chunk]: 

54 """Load text file. 

55 

56 Args: 

57 source: Path to text file 

58 

59 Returns: 

60 Single chunk with file content 

61 

62 Raises: 

63 RAGError: If file cannot be read 

64 """ 

65 try: 

66 path = Path(source) 

67 

68 # Read file asynchronously 

69 if HAS_AIOFILES: 

70 async with aiofiles.open(path, encoding="utf-8") as f: 

71 content = await f.read() 

72 else: 

73 content = await asyncio.to_thread(path.read_text, encoding="utf-8") 

74 

75 return [ 

76 Chunk( 

77 text=content, 

78 source=str(path), 

79 chunk_index=0, 

80 metadata={ 

81 "source": str(path), 

82 "type": "text", 

83 }, 

84 ), 

85 ] 

86 

87 except (OSError, UnicodeDecodeError) as e: 

88 msg = f"Failed to read text file {source}: {e}" 

89 raise RAGError(msg) from e 

90 except Exception as e: 

91 msg = f"Unexpected error loading text file {source}: {e}" 

92 raise RAGError(msg) from e 

93 

94 

95class PDFLoader(AbstractDocumentLoader): 

96 """Load PDF documents. 

97 

98 Requires: pypdf or pdfplumber 

99 """ 

100 

101 def __init__(self, extract_images: bool = False): 

102 """Initialize PDF loader. 

103 

104 Args: 

105 extract_images: Whether to extract images from PDFs 

106 """ 

107 self.extract_images = extract_images 

108 

109 async def load(self, source: str | Path) -> list[Chunk]: 

110 """Load PDF file. 

111 

112 Args: 

113 source: Path to PDF file 

114 

115 Returns: 

116 List of chunks (one per page) 

117 

118 Raises: 

119 RAGError: If PDF cannot be read 

120 """ 

121 try: 

122 try: 

123 import pypdf # type: ignore[import-not-found] 

124 except ImportError as e: 

125 msg = ( 

126 "PDF loading requires 'pypdf' package. " 

127 "Install with: pip install lexigram-ai-rag" 

128 ) 

129 raise ImportError(msg) from e 

130 

131 path = Path(source) 

132 exists = await asyncio.to_thread(path.exists) 

133 if not exists: 

134 raise FileNotFoundError(f"PDF file not found: {source}") 

135 

136 # Load PDF asynchronously 

137 def _load_pdf() -> Any: 

138 chunks: list[Chunk] = [] 

139 try: 

140 with open(path, "rb") as f: 

141 reader = pypdf.PdfReader(f) 

142 for page_num, page in enumerate(reader.pages): 

143 text = page.extract_text() 

144 if text.strip(): 

145 chunks.append( 

146 Chunk( 

147 text=text, 

148 source=str(path), 

149 chunk_index=page_num, 

150 metadata={ 

151 "source": str(path), 

152 "page": page_num + 1, 

153 "type": "pdf", 

154 }, 

155 ), 

156 ) 

157 except Exception as e: 

158 raise RAGError(f"PDF parsing error: {e}") from e 

159 return chunks 

160 

161 return await asyncio.to_thread(_load_pdf) 

162 

163 except (FileNotFoundError, ImportError, RAGError): 

164 raise 

165 except Exception as e: 

166 msg = f"Failed to load PDF {source}: {e}" 

167 raise RAGError(msg) from e 

168 

169 

170class MarkdownLoader(AbstractDocumentLoader): 

171 """Load Markdown documents. 

172 

173 Preserves structure while extracting text. 

174 """ 

175 

176 async def load(self, source: str | Path) -> list[Chunk]: 

177 """Load Markdown file. 

178 

179 Args: 

180 source: Path to Markdown file 

181 

182 Returns: 

183 Chunks split by headers 

184 

185 Raises: 

186 RAGError: If file cannot be read 

187 """ 

188 try: 

189 path = Path(source) 

190 

191 # Read file asynchronously 

192 if HAS_AIOFILES: 

193 async with aiofiles.open(path, encoding="utf-8") as f: 

194 content = await f.read() 

195 else: 

196 content = await asyncio.to_thread(path.read_text, encoding="utf-8") 

197 

198 # Split by headers (simple approach) 

199 chunks: list[Chunk] = [] 

200 current_chunk: list[str] = [] 

201 current_header = None 

202 

203 for line in content.split("\n"): 

204 if line.startswith("#"): 

205 # Save previous chunk 

206 if current_chunk: 

207 chunks.append( 

208 Chunk( 

209 text="\n".join(current_chunk), 

210 source=str(path), 

211 chunk_index=len(chunks), 

212 metadata={ 

213 "source": str(path), 

214 "header": current_header, 

215 "type": "markdown", 

216 }, 

217 ), 

218 ) 

219 # Start new chunk 

220 current_header = line.strip("# ") 

221 current_chunk = [line] 

222 else: 

223 current_chunk.append(line) 

224 

225 # Save last chunk 

226 if current_chunk: 

227 chunks.append( 

228 Chunk( 

229 text="\n".join(current_chunk), 

230 source=str(path), 

231 chunk_index=len(chunks), 

232 metadata={ 

233 "source": str(path), 

234 "header": current_header, 

235 "type": "markdown", 

236 }, 

237 ), 

238 ) 

239 except (OSError, UnicodeDecodeError) as e: 

240 msg = f"Failed to read Markdown {source}: {e}" 

241 raise RAGError(msg) from e 

242 except Exception as e: 

243 msg = f"Unexpected error loading Markdown {source}: {e}" 

244 raise RAGError(msg) from e 

245 else: 

246 return chunks 

247 

248 

249class HTMLLoader(AbstractDocumentLoader): 

250 """Load HTML documents. 

251 

252 Requires: beautifulsoup4 

253 """ 

254 

255 async def load(self, source: str | Path) -> list[Chunk]: 

256 """Load HTML file or URL. 

257 

258 Args: 

259 source: Path or URL to HTML 

260 

261 Returns: 

262 Single chunk with extracted text 

263 

264 Raises: 

265 RAGError: If HTML cannot be parsed 

266 """ 

267 try: 

268 try: 

269 from bs4 import BeautifulSoup # type: ignore[import-not-found] 

270 except ImportError as e: 

271 msg = ( 

272 "HTML loading requires 'beautifulsoup4' package. " 

273 "Install with: pip install lexigram-ai-rag" 

274 ) 

275 raise ImportError(msg) from e 

276 # Handle URL vs file path 

277 if str(source).startswith(("http://", "https://")): 

278 # Load from URL 

279 try: 

280 import aiohttp 

281 except ImportError as _e: 

282 msg = ( 

283 "HTML loading from URLs requires 'aiohttp' package. " 

284 "Install with: pip install lexigram-ai-rag" 

285 ) 

286 raise ImportError(msg) from _e 

287 async with aiohttp.ClientSession( 

288 timeout=aiohttp.ClientTimeout(total=30.0), 

289 ) as _session: 

290 async with _session.get(str(source)) as _resp: 

291 html = await _resp.text() 

292 else: 

293 # Load from file 

294 path = Path(source) 

295 if HAS_AIOFILES: 

296 async with aiofiles.open(path, encoding="utf-8") as f: 

297 html = await f.read() 

298 else: 

299 html = await asyncio.to_thread(path.read_text, encoding="utf-8") 

300 

301 # Parse HTML 

302 def _parse_html() -> Any: 

303 soup = BeautifulSoup(html, "html.parser") 

304 # Remove script and style elements 

305 for script in soup(["script", "style"]): 

306 script.decompose() 

307 return soup.get_text(separator="\n", strip=True) 

308 

309 text = await asyncio.to_thread(_parse_html) 

310 

311 return [ 

312 Chunk( 

313 text=text, 

314 source=str(source), 

315 chunk_index=0, 

316 metadata={ 

317 "source": str(source), 

318 "type": "html", 

319 }, 

320 ), 

321 ] 

322 

323 except Exception as e: 

324 msg = f"Failed to load HTML {source}: {e}" 

325 raise RAGError(msg) from e 

326 

327 

328class JSONLoader(AbstractDocumentLoader): 

329 """Load JSON and JSONL documents. 

330 

331 Produces one chunk per top-level array element (JSON) or per line 

332 (JSONL). Falls back to a single chunk for scalar/object JSON files. 

333 

334 Example: 

335 >>> loader = JSONLoader() 

336 >>> chunks = await loader.load("records.json") 

337 >>> chunks = await loader.load("records.jsonl") 

338 """ 

339 

340 def __init__(self, *, text_key: str | None = None) -> None: 

341 """Initialize JSON loader. 

342 

343 Args: 

344 text_key: When set, extract only this key's value as the chunk 

345 text. If not set, the raw JSON string of each record is used. 

346 """ 

347 self.text_key = text_key 

348 

349 async def load(self, source: str | Path) -> list[Chunk]: 

350 """Load JSON or JSONL file. 

351 

352 Args: 

353 source: Path to JSON or JSONL file. 

354 

355 Returns: 

356 List of chunks — one per top-level item or one for the whole file. 

357 

358 Raises: 

359 RAGError: If the file cannot be read or parsed. 

360 """ 

361 try: 

362 path = Path(source) 

363 if HAS_AIOFILES: 

364 async with aiofiles.open(path, encoding="utf-8") as f: 

365 raw = await f.read() 

366 else: 

367 raw = await asyncio.to_thread(path.read_text, encoding="utf-8") 

368 

369 file_type = "jsonl" if path.suffix.lower() == ".jsonl" else "json" 

370 records: list[Any] 

371 

372 if file_type == "jsonl": 

373 records = [] 

374 for line in raw.splitlines(): 

375 line = line.strip() 

376 if line: 

377 records.append(_loads(line)) 

378 else: 

379 data = _loads(raw) 

380 records = data if isinstance(data, list) else [data] 

381 

382 chunks: list[Chunk] = [] 

383 for idx, record in enumerate(records): 

384 if self.text_key and isinstance(record, dict): 

385 text = str(record.get(self.text_key, "")) 

386 meta: dict[str, Any] = { 

387 k: v for k, v in record.items() if k != self.text_key 

388 } 

389 else: 

390 text = record if isinstance(record, str) else dumps_str(record) 

391 meta = {} 

392 

393 chunks.append( 

394 Chunk( 

395 text=text, 

396 source=str(path), 

397 chunk_index=idx, 

398 metadata={ 

399 "source": str(path), 

400 "type": file_type, 

401 "record_index": idx, 

402 **meta, 

403 }, 

404 ) 

405 ) 

406 

407 return chunks 

408 

409 except (OSError, UnicodeDecodeError) as e: 

410 msg = f"Failed to read JSON file {source}: {e}" 

411 raise RAGError(msg) from e 

412 except ValueError as e: 

413 msg = f"Failed to parse JSON file {source}: {e}" 

414 raise RAGError(msg) from e 

415 except Exception as e: 

416 msg = f"Unexpected error loading JSON file {source}: {e}" 

417 raise RAGError(msg) from e 

418 

419 

420class CSVLoader(AbstractDocumentLoader): 

421 """Load CSV and TSV documents. 

422 

423 Produces one chunk per row by default, or batched rows when 

424 ``batch_size`` is greater than 1. 

425 

426 Example: 

427 >>> loader = CSVLoader() 

428 >>> chunks = await loader.load("data.csv") 

429 >>> loader_tsv = CSVLoader(delimiter="\\t") 

430 >>> chunks = await loader_tsv.load("data.tsv") 

431 """ 

432 

433 def __init__( 

434 self, 

435 *, 

436 delimiter: str = ",", 

437 batch_size: int = 1, 

438 text_columns: list[str] | None = None, 

439 ) -> None: 

440 """Initialize CSV loader. 

441 

442 Args: 

443 delimiter: Field delimiter (default: comma). 

444 batch_size: Number of rows per chunk. Defaults to 1. 

445 text_columns: If set, only these columns are included in chunk 

446 text. All columns are used when not set. 

447 """ 

448 self.delimiter = delimiter 

449 self.batch_size = batch_size 

450 self.text_columns = text_columns 

451 

452 async def load(self, source: str | Path) -> list[Chunk]: 

453 """Load CSV or TSV file. 

454 

455 Args: 

456 source: Path to CSV or TSV file. 

457 

458 Returns: 

459 List of chunks — one per row or per batch of rows. 

460 

461 Raises: 

462 RAGError: If the file cannot be read or parsed. 

463 """ 

464 try: 

465 path = Path(source) 

466 if HAS_AIOFILES: 

467 async with aiofiles.open(path, encoding="utf-8", newline="") as f: 

468 raw = await f.read() 

469 else: 

470 raw = await asyncio.to_thread(path.read_text, encoding="utf-8") 

471 

472 def _parse_csv() -> list[dict[str, str]]: 

473 reader = csv.DictReader( 

474 io.StringIO(raw), 

475 delimiter=self.delimiter, 

476 ) 

477 return list(reader) 

478 

479 rows = await asyncio.to_thread(_parse_csv) 

480 

481 chunks: list[Chunk] = [] 

482 batch: list[dict[str, str]] = [] 

483 

484 def _row_to_text(row: dict[str, str]) -> str: 

485 if self.text_columns: 

486 cols = {k: v for k, v in row.items() if k in self.text_columns} 

487 else: 

488 cols = row 

489 return " ".join(f"{k}: {v}" for k, v in cols.items()) 

490 

491 for row_idx, row in enumerate(rows): 

492 batch.append(row) 

493 if len(batch) >= self.batch_size: 

494 text = "\n".join(_row_to_text(r) for r in batch) 

495 chunk_idx = row_idx // self.batch_size 

496 chunks.append( 

497 Chunk( 

498 text=text, 

499 source=str(path), 

500 chunk_index=chunk_idx, 

501 metadata={ 

502 "source": str(path), 

503 "type": "csv", 

504 "row_start": chunk_idx * self.batch_size, 

505 "row_end": row_idx, 

506 }, 

507 ) 

508 ) 

509 batch = [] 

510 

511 # Flush remaining rows 

512 if batch: 

513 chunk_idx = len(rows) // self.batch_size 

514 text = "\n".join(_row_to_text(r) for r in batch) 

515 chunks.append( 

516 Chunk( 

517 text=text, 

518 source=str(path), 

519 chunk_index=chunk_idx, 

520 metadata={ 

521 "source": str(path), 

522 "type": "csv", 

523 "row_start": chunk_idx * self.batch_size, 

524 "row_end": len(rows) - 1, 

525 }, 

526 ) 

527 ) 

528 

529 return chunks 

530 

531 except (OSError, UnicodeDecodeError) as e: 

532 msg = f"Failed to read CSV file {source}: {e}" 

533 raise RAGError(msg) from e 

534 except csv.Error as e: 

535 msg = f"Failed to parse CSV file {source}: {e}" 

536 raise RAGError(msg) from e 

537 except Exception as e: 

538 msg = f"Unexpected error loading CSV file {source}: {e}" 

539 raise RAGError(msg) from e