tools/pdf_reader.py

from __future__ import annotations

import csv
from pathlib import Path
from typing import Any

import fitz
import pandas as pd
from docx import Document


def read_supported_file(file_path: str | None, question: str, max_chars: int = 22000) -> dict[str, Any]:
    if not file_path:
        raise ValueError("This question has no downloaded file.")
    path = Path(file_path)
    if not path.exists():
        raise FileNotFoundError(path)

    suffix = path.suffix.lower()
    if suffix == ".pdf":
        content = _read_pdf(path, question)
    elif suffix in {".csv", ".xlsx", ".xls"}:
        content = _read_spreadsheet(path)
    elif suffix == ".docx":
        content = _read_docx(path)
    elif suffix in {".txt", ".md", ".json", ".xml", ".html"}:
        content = path.read_text(encoding="utf-8", errors="replace")
    else:
        content = path.read_bytes()[:2000].hex()

    return {
        "ok": True,
        "source": str(path),
        "content": content[:max_chars],
        "metadata": {"suffix": suffix, "size": path.stat().st_size},
    }


def _read_pdf(path: Path, question: str) -> str:
    terms = {word.lower().strip(".,?!:;()[]") for word in question.split() if len(word) > 4}
    pages: list[tuple[int, str, int]] = []
    with fitz.open(path) as document:
        for number, page in enumerate(document, start=1):
            text = page.get_text("text").strip()
            score = sum(text.lower().count(term) for term in terms)
            pages.append((number, text, score))
    pages.sort(key=lambda item: item[2], reverse=True)
    selected = pages[: min(12, len(pages))]
    selected.sort(key=lambda item: item[0])
    return "\n\n".join(f"--- Page {number} ---\n{text}" for number, text, _ in selected)


def _read_spreadsheet(path: Path) -> str:
    sheets = pd.read_excel(path, sheet_name=None) if path.suffix.lower() != ".csv" else {"csv": pd.read_csv(path)}
    chunks: list[str] = []
    for name, frame in sheets.items():
        chunks.append(f"SHEET: {name}\nSHAPE: {frame.shape}\nCOLUMNS: {list(frame.columns)}\n{frame.head(80).to_csv(index=False)}")
    return "\n\n".join(chunks)


def _read_docx(path: Path) -> str:
    doc = Document(path)
    paragraphs = [paragraph.text for paragraph in doc.paragraphs if paragraph.text.strip()]
    for table in doc.tables:
        for row in table.rows:
            paragraphs.append(" | ".join(cell.text for cell in row.cells))
    return "\n".join(paragraphs)