Files
Mutual_Fund/rag/document_parser.py

52 lines
1.8 KiB
Python

"""Knowledge document parsing for supported upload formats."""
from __future__ import annotations
from pathlib import Path
from xml.etree import ElementTree
from zipfile import BadZipFile, ZipFile
from pypdf import PdfReader
SUPPORTED_EXTENSIONS = {".txt", ".md", ".docx", ".pdf"}
_WORD_NS = "{http://schemas.openxmlformats.org/wordprocessingml/2006/main}"
def parse_document(path: str | Path) -> str:
"""Return document text for a supported knowledge-base upload."""
document_path = Path(path)
suffix = document_path.suffix.lower()
if suffix not in SUPPORTED_EXTENSIONS:
raise ValueError(f"Unsupported document type: {suffix or '<none>'}")
if not document_path.is_file():
raise FileNotFoundError(document_path)
if suffix in {".txt", ".md"}:
return document_path.read_text(encoding="utf-8-sig")
if suffix == ".pdf":
return _parse_pdf(document_path)
return _parse_docx(document_path)
def _parse_docx(path: Path) -> str:
try:
with ZipFile(path) as archive:
xml = archive.read("word/document.xml")
root = ElementTree.fromstring(xml)
except (BadZipFile, KeyError, ElementTree.ParseError) as exc:
raise ValueError(f"Invalid DOCX document: {path}") from exc
paragraphs = []
for paragraph in root.iter(f"{_WORD_NS}p"):
text = "".join(node.text or "" for node in paragraph.iter(f"{_WORD_NS}t"))
if text:
paragraphs.append(text)
return "\n".join(paragraphs)
def _parse_pdf(path: Path) -> str:
try:
reader = PdfReader(str(path))
return "\n".join(page.extract_text() or "" for page in reader.pages).strip()
except Exception as exc: # pypdf raises format-specific exceptions
raise ValueError(f"Invalid PDF document: {path}") from exc