Spaces:
Sleeping
Sleeping
File size: 3,458 Bytes
4e8d71a | 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 | """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,
},
}
|