Mehramd/multi-file-rag-chat

1from pathlib import Path
2import xml.etree.ElementTree as ET
3import pandas as pd
4
5from pypdf import PdfReader
6from docx import Document as DoxDocument
7
8from langchain_core.documents import Document
9
10# OCR fallback deps
11from pdf2image import convert_from_path
12import pytesseract
13
14DATA_DIR = Path("data")
15
16SUPPORTED_EXTENSIONS = {
17    ".pdf",
18    ".docx",
19    ".xml",
20    ".txt",
21    ".md",
22    ".csv",
23    ".xlsx",
24    ".xls",
25}
26
27def read_pdf(file_path):
28    """
29    Extracts text per page. If a page has no native text layer (scanned /
30    image-only / diagram-heavy page), falls back to OCR via pytesseract.
31    """
32    reader = PdfReader(Path(file_path))
33    documents = []
34    ocr_pages = []
35
36    for page_num, page in enumerate(reader.pages, start=1):
37        text = page.extract_text()
38        used_ocr = False
39
40        if not text or not text.strip():
41            # Native extraction found nothing -> try OCR on this page only
42            try:
43                images = convert_from_path(
44                    str(file_path),
45                    dpi=300,
46                    first_page=page_num,
47                    last_page=page_num,
48                )
49                if images:
50                    text = pytesseract.image_to_string(images[0])
51                    used_ocr = True
52            except Exception as e:
53                print(f"OCR failed on page {page_num} of {file_path.name}: {e}")
54                text = ""
55
56        if text and text.strip():
57            if used_ocr:
58                ocr_pages.append(page_num)
59            documents.append(
60                Document(
61                    page_content=text,
62                    metadata={
63                        "source": file_path.name,
64                        "page": page_num,
65                        "type": "pdf",
66                        "ocr": used_ocr,
67                    }
68                )
69            )
70
71    if ocr_pages:
72        print(f"ℹ️ {file_path.name}: used OCR fallback on pages {ocr_pages} (no native text layer found).")
73
74    return documents
75
76def read_docx(file_path):
77    reader = DoxDocument(Path(file_path))
78    parts = []
79    for paragraph in reader.paragraphs:
80        if paragraph.text.strip():
81            parts.append(paragraph.text.strip())
82
83    for table in reader.tables:
84        for row in table.rows:
85            row_text = "|".join(cell.text.strip() for cell in row.cells)
86            if row_text.strip():
87                parts.append(row_text)
88    return "\n".join(parts)
89
90def read_xml(file_path):
91    try:
92        tree = ET.parse(file_path)
93        root = tree.getroot()
94
95        text_parts = []
96
97        for element in root.iter():
98            if element.text and element.text.strip():
99                text_parts.append(element.text.strip())
100
101        return "\n".join(text_parts)
102
103    except Exception:
104        return Path(file_path).read_text(
105            encoding="utf-8",
106            errors="ignore"
107        )
108
109def read_txt_or_md(file_path):
110    return Path(file_path).read_text(
111        encoding="utf-8",
112        errors="ignore"
113    )
114
115
116def dataframe_to_text(df, source_name):
117    return (
118        f"Source: {source_name}\n\n"
119        + df.fillna("").to_string(index=False)
120    )
121
122def read_csv(file_path):
123    df = pd.read_csv(file_path)
124    return [Document(page_content=dataframe_to_text(df,file_path.name), metadata={"source": file_path.name, "type": "csv", "rows": len(df), "columns":len(df.columns)})]
125
126def read_excel(file_path):
127    docs = []
128    sheets = pd.read_excel(file_path, sheet_name = None)
129    for sheet_name, df in sheets.items():
130        docs.append(Document(page_content=dataframe_to_text(df,file_path.name), metadata={"source": file_path.name, "type": "excel", "rows": len(df), "columns":len(df.columns)}))
131    return docs
132
133def load_single_file(file_path):
134    path = Path(file_path)
135    suffix = path.suffix.lower()
136    if suffix == ".pdf":
137        return read_pdf(path)
138    if suffix == ".docx":
139        return [Document(page_content= read_docx(path), metadata = {"source":file_path.name,"type":"docx"})]
140    if suffix == ".xml":
141        return [Document(page_content= read_xml(path), metadata = {"source":file_path.name,"type":"xml"})]
142    if suffix in {".txt",".md"}:
143        return [Document(page_content= read_txt_or_md(path), metadata = {"source":file_path.name,"type":suffix.replace(".","")})]
144    if suffix == ".csv":
145        return read_csv(path)
146    if suffix == ".xlsx" or suffix == ".xls":
147        return read_excel(path)
148    return []
149
150def load_document_from_directory(directory= DATA_DIR):
151    documents = []
152    for file_path in directory.iterdir():
153        if file_path.suffix.lower() in SUPPORTED_EXTENSIONS:
154            documents.extend(load_single_file(file_path))
155    return documents
156
157def load_documents_from_paths(file_paths):
158    documents = []
159
160    for fp in file_paths:
161        try:
162            documents.extend(
163                load_single_file(Path(fp))
164            )
165        except Exception as e:
166            print(f"Failed to load {fp}: {e}")
167
168    return documents