batch.py

"""Folder-level extraction helpers for locally downloaded PDFs."""

from __future__ import annotations

from pathlib import Path
from typing import Dict, List, Optional

from extractor import extract_pdf
from perovskite import extract_perovskite_fields


def _pdf_files(folder: Path, *, recursive: bool) -> List[Path]:
    pattern = "**/*.pdf" if recursive else "*.pdf"
    return sorted(p for p in folder.glob(pattern) if p.is_file())


def _lighten_extraction(ext: Dict, *, include_full_text: bool) -> Dict:
    item = dict(ext)
    if not include_full_text:
        item.pop("full_text", None)
    return item


def extract_folder_pdfs(
    folder_path: str,
    query: str = "",
    *,
    recursive: bool = True,
    max_files: Optional[int] = None,
    use_llm_filter: Optional[bool] = None,
    include_full_text: bool = False,
    analyze_images: Optional[bool] = None,
    image_query: str = "",
) -> Dict:
    """Extract text/images/perovskite fields from PDFs under ``folder_path``.

    The response is intentionally lightweight by default: per-paper full text is
    omitted unless ``include_full_text`` is true, so MCP responses stay small.
    """
    folder = Path(folder_path).expanduser()
    if not folder.exists() or not folder.is_dir():
        return {
            "ok": False,
            "error": "folder_not_found",
            "folder_path": str(folder),
            "results": [],
            "summary": {
                "total_pdf_files": 0,
                "processed": 0,
                "succeeded": 0,
                "failed": 0,
            },
        }

    pdfs = _pdf_files(folder, recursive=recursive)
    total_pdf_files = len(pdfs)
    if max_files is not None:
        limit = max(0, int(max_files))
        pdfs = pdfs[:limit]

    results: List[Dict] = []
    succeeded = 0
    failed = 0

    for pdf_path in pdfs:
        try:
            ext = extract_pdf(
                str(pdf_path),
                query=query,
                analyze_images=analyze_images,
                image_query=image_query,
            )
            if not ext.get("ok"):
                failed += 1
                item = {
                    "ok": False,
                    "pdf_path": str(pdf_path),
                    "filename": pdf_path.name,
                    "error": ext.get("error") or "extract_failed",
                }
            else:
                fields = extract_perovskite_fields(
                    ext.get("full_text") or "",
                    use_llm_filter=use_llm_filter,
                )
                item = _lighten_extraction(ext, include_full_text=include_full_text)
                item["filename"] = pdf_path.name
                item["perovskite_fields"] = fields
                succeeded += 1
        except Exception as exc:  # noqa: BLE001 - keep batch processing going
            failed += 1
            item = {
                "ok": False,
                "pdf_path": str(pdf_path),
                "filename": pdf_path.name,
                "error": f"extract_failed: {exc}",
            }
        results.append(item)

    return {
        "ok": failed == 0,
        "folder_path": str(folder),
        "recursive": recursive,
        "max_files": max_files,
        "results": results,
        "summary": {
            "total_pdf_files": total_pdf_files,
            "processed": len(pdfs),
            "succeeded": succeeded,
            "failed": failed,
        },
    }