"""``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 @router.post("/probe", response_model=ProbeResponse) 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, )