Spaces:
Running on Zero
Running on Zero
File size: 5,274 Bytes
087643a | 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 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 | # FastAPI app: /health, /detect, /ask. Kept thin on purpose - upload
# validation, error mapping, logging, then hand off to Detector or the
# reasoning pipeline, both already tested elsewhere.
from __future__ import annotations
import os
from contextlib import asynccontextmanager
# loads .env if one exists (local dev) - does nothing if it doesn't (Docker,
# HF Spaces, Kaggle all set real env vars/secrets instead)
from dotenv import load_dotenv
load_dotenv()
from fastapi import Depends, FastAPI, File, Form, HTTPException, Request, UploadFile
from fastapi.responses import JSONResponse
from PIL import Image, UnidentifiedImageError
from app.constants import MODEL_VERSION
from app.detector import Detector
from app.logging_config import configure_logging, log_request
from app.reasoning.pipeline import answer_question
from app.schemas import DetectResponse
logger = configure_logging()
# Read once at import, not per-request. Tests override this directly
# rather than faking an env var.
MAX_UPLOAD_BYTES = int(os.environ.get("MAX_UPLOAD_MB", "10")) * 1024 * 1024
_detector_instance: Detector | None = None
def get_detector() -> Detector:
"""FastAPI dependency - same Detector instance for every route.
Not a bare global so tests can override it with a stub."""
global _detector_instance
if _detector_instance is None:
_detector_instance = Detector()
return _detector_instance
@asynccontextmanager
async def lifespan(app: FastAPI):
"""Tries to load weights at startup but doesn't crash if that fails -
a service that's up and reporting "model not loaded" is way easier to
debug than one that won't start at all."""
try:
get_detector().load()
logger.info("startup: model loaded", extra={"extra_fields": {}})
except Exception as error:
logger.error(
"startup: model failed to load, API will report 503 on inference routes",
extra={"extra_fields": {"error": str(error)}},
)
yield
app = FastAPI(title="Constrained Document Layout Detection API", lifespan=lifespan)
@app.exception_handler(Exception)
async def handle_unexpected_error(request: Request, exc: Exception):
"""Catches anything a route didn't handle and returns a safe generic
message instead of leaking a stack trace. Full error still goes to logs."""
logger.error(
"unhandled exception",
extra={"extra_fields": {"path": str(request.url), "error": str(exc)}},
)
return JSONResponse(
status_code=500,
content={"detail": "Internal server error. See server logs for details."},
)
def _read_and_validate_image(contents: bytes, content_type: str | None) -> Image.Image:
"""Size check, then content-type check, then an actual decode attempt -
a file can claim image/png and still be garbage."""
if len(contents) > MAX_UPLOAD_BYTES:
raise HTTPException(
status_code=413,
detail=f"Upload exceeds the {MAX_UPLOAD_BYTES // (1024 * 1024)} MB limit.",
)
if content_type is None or not content_type.startswith("image/"):
raise HTTPException(
status_code=415,
detail=f"Expected an image upload, got content-type '{content_type}'.",
)
try:
import io
return Image.open(io.BytesIO(contents)).convert("RGB")
except UnidentifiedImageError:
raise HTTPException(status_code=415, detail="Could not decode the upload as an image.")
@app.get("/health")
def health(detector: Detector = Depends(get_detector)):
return {"status": "ok", "model_loaded": detector.is_loaded, "model_version": MODEL_VERSION}
@app.post("/detect", response_model=DetectResponse)
async def detect(file: UploadFile = File(...), detector: Detector = Depends(get_detector)):
with log_request(logger, "/detect") as ctx:
if not detector.is_loaded:
raise HTTPException(
status_code=503,
detail="Model is not loaded. Check MODEL_PATH and server startup logs.",
)
contents = await file.read()
image = _read_and_validate_image(contents, file.content_type)
detections, inference_ms = detector.predict(image)
ctx["detection_count"] = len(detections)
return DetectResponse(
detections=detections,
image_size={"width": image.width, "height": image.height},
inference_ms=round(inference_ms, 2),
model_version=MODEL_VERSION,
)
@app.post("/ask")
async def ask(
file: UploadFile = File(...),
question: str = Form(..., min_length=1, max_length=500),
detector: Detector = Depends(get_detector),
):
"""Part B. All the actual routing/evidence/guardrail/synthesis logic
lives in app.reasoning.pipeline - this is just upload validation + wiring."""
with log_request(logger, "/ask", question_length=len(question)) as ctx:
contents = await file.read()
image = _read_and_validate_image(contents, file.content_type)
response = answer_question(image=image, question=question, detector=detector)
ctx["used_detection"] = response.used_detection
ctx["insufficient_information"] = response.insufficient_information
return response
|