Spaces:
Running
Running
File size: 5,232 Bytes
5ea3240 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 | from __future__ import annotations
import json
import shutil
from pathlib import Path
from typing import Iterable, TYPE_CHECKING
import pandas as pd
from bs4 import BeautifulSoup
from docx import Document as DocxDocument
from pptx import Presentation
from pypdf import PdfReader
from .schemas import Document
from .security import safe_extract_zip, validate_upload, sanitize_filename
if TYPE_CHECKING:
from .llm import GeminiGateway
class DocumentLoader:
def __init__(self, ocr_gateway: GeminiGateway | None = None):
self.ocr_gateway = ocr_gateway
def expand_inputs(self, paths: Iterable[Path], workspace_dir: Path) -> list[Path]:
expanded: list[Path] = []
for path in paths:
validate_upload(path)
if path.suffix.lower() == ".zip":
expanded.extend(safe_extract_zip(path, workspace_dir / "unzipped"))
else:
target = workspace_dir / "uploads" / sanitize_filename(path.name)
target.parent.mkdir(parents=True, exist_ok=True)
if path.resolve() != target.resolve():
shutil.copy2(path, target)
expanded.append(target)
return expanded
def load(self, path: Path) -> tuple[list[Document], list[tuple[str, pd.DataFrame]]]:
ext = path.suffix.lower()
if ext == ".pdf":
return self._pdf(path), []
if ext == ".docx":
return self._docx(path), []
if ext == ".pptx":
return self._pptx(path), []
if ext == ".csv":
df = pd.read_csv(path)
return self._dataframe_docs(path.name, df), [(path.stem, df)]
if ext in {".xlsx", ".xls"}:
sheets = pd.read_excel(path, sheet_name=None)
docs: list[Document] = []
tables: list[tuple[str, pd.DataFrame]] = []
for sheet, df in sheets.items():
docs.extend(self._dataframe_docs(f"{path.name}:{sheet}", df))
tables.append((f"{path.stem}_{sheet}", df))
return docs, tables
if ext == ".json":
data = json.loads(path.read_text(encoding="utf-8", errors="ignore"))
return [Document(json.dumps(data, indent=2, ensure_ascii=False), path.name)], []
if ext in {".html", ".htm"}:
soup = BeautifulSoup(path.read_text(encoding="utf-8", errors="ignore"), "lxml")
title = soup.title.string.strip() if soup.title and soup.title.string else None
return [Document(soup.get_text("\n", strip=True), path.name, section=title)], []
if ext in {".png", ".jpg", ".jpeg", ".webp"}:
if not self.ocr_gateway:
return [Document("[Image file indexed without OCR. Enable Gemini OCR to extract its text.]", path.name)], []
return [Document(self.ocr_gateway.extract_file_text(path), path.name, metadata={"ocr": "gemini"})], []
text = path.read_text(encoding="utf-8", errors="ignore")
return [Document(text, path.name)], []
def _pdf(self, path: Path) -> list[Document]:
reader = PdfReader(str(path))
docs: list[Document] = []
total_chars = 0
for i, page in enumerate(reader.pages, start=1):
text = (page.extract_text() or "").strip()
total_chars += len(text)
if text:
docs.append(Document(text, path.name, page=i))
if total_chars < 80 and self.ocr_gateway:
extracted = self.ocr_gateway.extract_file_text(path)
return [Document(extracted, path.name, metadata={"ocr": "gemini"})]
return docs
def _docx(self, path: Path) -> list[Document]:
doc = DocxDocument(str(path))
blocks: list[str] = []
for p in doc.paragraphs:
if p.text.strip():
blocks.append(p.text.strip())
for table in doc.tables:
rows = []
for row in table.rows:
rows.append(" | ".join(cell.text.strip() for cell in row.cells))
if rows:
blocks.append("\n".join(rows))
return [Document("\n\n".join(blocks), path.name)]
def _pptx(self, path: Path) -> list[Document]:
prs = Presentation(str(path))
docs: list[Document] = []
for i, slide in enumerate(prs.slides, start=1):
texts = []
for shape in slide.shapes:
if hasattr(shape, "text") and shape.text.strip():
texts.append(shape.text.strip())
if texts:
docs.append(Document("\n".join(texts), path.name, page=i, section=f"Slide {i}"))
return docs
def _dataframe_docs(self, source: str, df: pd.DataFrame) -> list[Document]:
docs: list[Document] = []
clean = df.fillna("")
for start in range(0, len(clean), 50):
block = clean.iloc[start:start + 50]
text = block.to_csv(index=False)
docs.append(Document(text, source, section=f"Rows {start + 1}-{start + len(block)}", metadata={"structured": True}))
if not docs:
docs.append(Document("Columns: " + ", ".join(map(str, df.columns)), source, metadata={"structured": True}))
return docs
|