patristic-be / src /api /routers /addbook.py
Mario33333's picture
deploy: reader-access security hardening (feature/audiobook@06a5ed8)
7c2d250 verified
Raw
History Blame Contribute Delete
10.5 kB
"""``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,
)