Spaces:
Running
Running
| """``POST /books/probe`` and the Add Book preview helpers. | |
| Wraps :mod:`src.lib.metadata_probe` so the FE wizard's Step 2 can show | |
| prefilled metadata + suggested labels + sample pages + a hard ingest | |
| cost estimate before the user commits. | |
| This router is intentionally separate from ``/books`` (CRUD) so that the | |
| two stay small and the probe's expensive Gemini path is easy to mock in | |
| tests. Per CONTRACT.md §1, every wire field is camelCase via the | |
| :class:`ApiModel` base — see :mod:`src.api.dto.addbook`. | |
| Cost rule (BACKEND_BUILD.md §11): the probe IS a paid call (~$0.005 | |
| for the ~5-page metadata sample). All unit tests mock | |
| ``probe_metadata`` so CI never pays for it; the live probe path is | |
| exercised only by the staged-but-skipped ``tests/test_addbook_live.py`` | |
| that's gated behind ``pytest -m live`` (user opt-in). | |
| """ | |
| from __future__ import annotations | |
| from pathlib import Path | |
| import fitz # PyMuPDF | |
| from fastapi import APIRouter, Depends | |
| from src.api.deps import require_admin | |
| from src.api.dto.addbook import ( | |
| GuessedMetadata, | |
| ProbeRequest, | |
| ProbeResponse, | |
| SamplePage, | |
| SuggestedLabel, | |
| ) | |
| from src.api.dto.books import ( | |
| EraEnum, | |
| ExtractionModeEnum, | |
| LanguageEnum, | |
| ReligionEnum, | |
| TraditionEnum, | |
| ) | |
| from src.api.errors import BadRequest | |
| from src.api.limits import enforce_paid_request | |
| from src.api.validation import require_safe_id | |
| from src.config import load_config | |
| from src.lib.auth.models import User | |
| from src.lib.costs.repo import compute_cost_usd | |
| from src.lib.metadata_probe import probe_metadata | |
| from src.lib.tools.repo import get_tool | |
| router = APIRouter(prefix="/books", tags=["addbook"]) | |
| # How many sample pages to surface in the wizard. The probe itself | |
| # samples up to ``metadata_probe.pages_to_sample`` (default 5) for the | |
| # Gemini call; we mirror that for the on-screen preview so what the | |
| # user sees matches what the model saw. | |
| _SAMPLE_PAGES = 5 | |
| def _load_upload_bytes(upload_id: str) -> bytes: | |
| """Read bytes for a previously-uploaded file. Raises 400 on miss. | |
| Doesn't open the file via :class:`fitz` to keep this helper cheap; | |
| the probe path opens it once via PyMuPDF for the sample text. | |
| """ | |
| require_safe_id(upload_id, field="uploadId") | |
| cfg = load_config() | |
| path: Path = cfg.paths.data_dir / "uploads" / f"{upload_id}.pdf" | |
| if not path.exists(): | |
| raise BadRequest( | |
| f"No upload with id {upload_id!r}. Re-upload the file and try again." | |
| ) | |
| return path.read_bytes() | |
| def _fetch_source_url(url: str) -> bytes: | |
| """Download PDF bytes from a public URL. Raises 400 on transport errors. | |
| Reuses the Stage-2 acquisition downloader so the probe sees exactly | |
| the bytes the real ingest will. The download is bounded by the | |
| same upload-size limit we apply to multipart uploads. | |
| """ | |
| # Local import: keeps requests + the net helper off the import-time cost | |
| # of every router module under ``uvicorn --reload``. | |
| from src.lib.net_safety import UnsafeUrlError, safe_get_bytes | |
| try: | |
| # SSRF guard: scheme allow-list + every resolved hop must be a public | |
| # address, redirects re-validated. Blocks fetches to cloud metadata / | |
| # loopback / private-LAN hosts via a user-supplied URL. | |
| return safe_get_bytes(url, timeout=30) | |
| except UnsafeUrlError as e: | |
| raise BadRequest(f"That URL isn't allowed: {e}") from e | |
| except Exception as e: # noqa: BLE001 | |
| raise BadRequest(f"Could not fetch URL {url!r}: {type(e).__name__}: {e}") from e | |
| def _sample_pages_from_pdf(pdf_bytes: bytes, n: int = _SAMPLE_PAGES) -> list[SamplePage]: | |
| """Extract per-page **native** text for the wizard preview. | |
| Returns at most ``n`` pages, fewer if the PDF is shorter. Despite the | |
| field name, ``SamplePage.ocrText`` here carries the *PyMuPDF native* | |
| extraction, NOT a Gemini OCR pass — even for scanned PDFs (where it may | |
| be empty). The FE renders the empty string explicitly so the user sees | |
| "this page has no native text" rather than a hidden zero-state. | |
| NOTE (capability gap): the richer "run OCR + cleanup on a sample and | |
| compare against the page image" flow is implemented in | |
| :func:`src.lib.probe_ocr.run_ocr_sample`, but it is currently only | |
| consumed by the legacy Streamlit UI | |
| (``src/stage9_ui/pages_v2/library_add.py``). This FastAPI probe path | |
| does NOT run OCR — it only surfaces native text + a Gemini *metadata* | |
| guess. Porting ``run_ocr_sample`` to an endpoint is what gives the new | |
| SPA a true text-only-vs-OCR sample comparison before ingest. | |
| """ | |
| out: list[SamplePage] = [] | |
| try: | |
| with fitz.open(stream=pdf_bytes, filetype="pdf") as doc: | |
| count = min(n, doc.page_count) | |
| for i in range(count): | |
| txt = doc.load_page(i).get_text("text") or "" | |
| out.append(SamplePage(pdf_page=i + 1, ocr_text=txt)) | |
| except Exception: # noqa: BLE001 — corrupted PDFs surface as empty list | |
| return [] | |
| return out | |
| def _estimate_ingest_cost_usd( | |
| *, | |
| pages_total: int, | |
| extraction_mode: str, | |
| ) -> float: | |
| """Cheap counterfactual estimate for ``estimated_ingest_cost_usd``. | |
| For ``extraction_mode == "native_text"`` the answer is $0 — every | |
| downstream stage runs locally (PyMuPDF text, BGE-M3 on CUDA, | |
| Qdrant upsert). For ``ocr`` mode we estimate two Gemini Flash | |
| passes (OCR + cleanup) at the active model's per-image pricing | |
| from ``tools_registry.yaml``. Returns 0.0 if pricing is missing | |
| so the wizard doesn't render a misleading "$nan" cell. | |
| The estimate is intentionally conservative — it under-counts | |
| nothing the user is about to be charged for (no images per page | |
| discount for short pages, no cache hits) so the "Confirm" button | |
| reads at least as much as the real bill. | |
| """ | |
| if extraction_mode == "native_text" or pages_total <= 0: | |
| return 0.0 | |
| # Active OCR + cleanup model. Defaults align with ``config.yaml`` | |
| # so a fresh checkout works without pulling apart the YAML. | |
| cfg = load_config() | |
| ocr_model = cfg.section("ocr").get("model", "gemini-2.5-flash") | |
| cleanup_model = cfg.section("cleanup").get("model", ocr_model) | |
| def _per_page_cost(model: str) -> float: | |
| tool = get_tool(model) | |
| if tool is None or tool.pricing is None: | |
| return 0.0 | |
| # Average tokens per page: ~1500 output (extracted text) + the | |
| # built-in per-image input. The pricing model has | |
| # ``image_input_tokens_per_image`` for image-grounded calls; | |
| # ``compute_cost_usd`` handles the math. | |
| return compute_cost_usd( | |
| tool.pricing, | |
| input_tokens=500, | |
| output_tokens=1500, | |
| image_count=1, | |
| n_calls=1, | |
| ) | |
| per_page = _per_page_cost(ocr_model) + _per_page_cost(cleanup_model) | |
| return round(per_page * pages_total, 4) | |
| def _coerce_enum(value: str | None, EnumCls): | |
| """Tolerate the lowercase strings :mod:`metadata_probe` emits. | |
| The probe returns plain lowercase strings filtered against | |
| ``RELIGIONS`` / ``ERAS`` / etc; the DTO uses ``Enum`` so the FE | |
| gets a TS union. Returns ``None`` for unknown / missing values so | |
| the wizard renders an empty selector rather than a 422. | |
| """ | |
| if not value: | |
| return None | |
| try: | |
| return EnumCls(value) | |
| except ValueError: | |
| return None | |
| def post_probe( | |
| req: ProbeRequest, | |
| user: User = Depends(require_admin), | |
| ) -> ProbeResponse: | |
| """Pre-ingest probe: count pages, detect text layer, ask Gemini for | |
| metadata, render sample pages, estimate ingest cost. | |
| Body must include either ``uploadId`` (preferred) or ``sourceUrl``. | |
| Admin-only because the probe is a paid call (~$0.005 of Gemini | |
| metadata per invocation). | |
| The endpoint never raises if the underlying Gemini call fails — | |
| instead the ``probeError`` field carries the model-side error and | |
| the wizard surfaces it inline so the user can retry without | |
| losing their already-typed metadata. | |
| """ | |
| enforce_paid_request(user) # daily spend ceiling + per-user rate limit | |
| if not req.upload_id and not req.source_url: | |
| raise BadRequest("Provide either uploadId or sourceUrl.") | |
| if req.upload_id: | |
| pdf_bytes = _load_upload_bytes(req.upload_id) | |
| else: | |
| # mypy/pyright narrowing: the `if` above guarantees source_url. | |
| assert req.source_url is not None | |
| pdf_bytes = _fetch_source_url(req.source_url) | |
| # Run the probe. On Gemini failure ``probe.error`` is populated but | |
| # the probe still returns the cheap structural fields (pages_total, | |
| # has_text_layer, extraction_mode_suggested) so the wizard can keep | |
| # rendering. Test paths mock ``probe_metadata`` so this is 0 paid | |
| # calls in CI. | |
| probe = probe_metadata(pdf_bytes) | |
| guessed = GuessedMetadata( | |
| title_ar=probe.title_ar, | |
| title_en=probe.title_en, | |
| author=probe.author, | |
| author_id=probe.author_id, | |
| era=_coerce_enum(probe.era, EraEnum), | |
| tradition=_coerce_enum(probe.tradition, TraditionEnum), | |
| language=_coerce_enum(probe.language, LanguageEnum), | |
| book_religion=_coerce_enum(probe.book_religion, ReligionEnum), | |
| author_religion=_coerce_enum(probe.author_religion, ReligionEnum), | |
| confidence=probe.confidence, | |
| rationale=probe.rationale, | |
| ) | |
| suggested_labels = [ | |
| SuggestedLabel(id=name, reason=None) for name in probe.suggested_labels | |
| ] | |
| sample_pages = _sample_pages_from_pdf(pdf_bytes) | |
| estimated = _estimate_ingest_cost_usd( | |
| pages_total=probe.pages_total, | |
| extraction_mode=probe.extraction_mode_suggested, | |
| ) | |
| return ProbeResponse( | |
| pages_total=probe.pages_total, | |
| has_text_layer=probe.extraction_mode_suggested == "native_text", | |
| suggested_extraction_mode=ExtractionModeEnum(probe.extraction_mode_suggested), | |
| guessed_metadata=guessed, | |
| suggested_labels=suggested_labels, | |
| sample_pages=sample_pages, | |
| estimated_ingest_cost_usd=estimated, | |
| probe_error=probe.error, | |
| ) | |