""" DocuLens - FastAPI REST endpoints. Mounted into Gradio's internal FastAPI app via APIRouter so that ZeroGPU compatibility is preserved (demo.launch handles @spaces.GPU). """ import os import json import time import asyncio import tempfile import logging from fastapi import APIRouter, UploadFile, File, Form, HTTPException, Depends from fastapi.middleware.cors import CORSMiddleware from PIL import Image from extraction import ( extract_invoice, extract_from_pdf, ExtractionResult, validate_with_vlm, run_dynamic_math_checks, image_to_base64, ) # Persistence layer (gracefully no-ops when Supabase is not configured) from db.supabase import save_document, create_run, complete_run, save_stage_result # Security middleware from middleware import check_rate_limit, validate_file_upload, sanitize_filename logger = logging.getLogger(__name__) # --------------------------------------------------------------------------- # Router (mounted by app.py into Gradio's FastAPI app) # --------------------------------------------------------------------------- api_router = APIRouter() HF_TOKEN = os.environ.get("HF_TOKEN", "") # --------------------------------------------------------------------------- # CORS helper — called by app.py after include_router # --------------------------------------------------------------------------- ALLOWED_ORIGINS = [ "http://localhost:5173", # Vite dev server "http://localhost:3000", "https://*.vercel.app", # Vercel preview deploys "https://ai-extractor.in", # Custom domain (apex) "https://www.ai-extractor.in", # Custom domain (www) ] FRONTEND_URL = os.environ.get("FRONTEND_URL", "") if FRONTEND_URL: ALLOWED_ORIGINS.append(FRONTEND_URL) def add_cors_middleware(app): """Add CORS middleware to the given FastAPI/Starlette app.""" app.add_middleware( CORSMiddleware, allow_origins=ALLOWED_ORIGINS, allow_origin_regex=r"https://.*\.vercel\.app", allow_credentials=True, allow_methods=["*"], allow_headers=["*"], ) # --------------------------------------------------------------------------- # Available models # --------------------------------------------------------------------------- EXTRACTION_MODELS = [ {"id": "Qwen/Qwen2.5-VL-72B-Instruct", "name": "Qwen2.5-VL 72B", "tier": "high"}, {"id": "Qwen/Qwen2.5-VL-7B-Instruct", "name": "Qwen2.5-VL 7B", "tier": "mid"}, {"id": "Qwen/Qwen2.5-VL-3B-Instruct", "name": "Qwen2.5-VL 3B", "tier": "fast"}, {"id": "ocr-only", "name": "OCR + Patterns", "tier": "fast"}, ] VALIDATION_MODELS = [ {"id": "Qwen/Qwen2.5-VL-72B-Instruct", "name": "Qwen2.5-VL 72B", "tier": "high"}, {"id": "Qwen/Qwen2.5-VL-7B-Instruct", "name": "Qwen2.5-VL 7B", "tier": "mid"}, ] # --------------------------------------------------------------------------- # Endpoints # --------------------------------------------------------------------------- @api_router.get("/api/v1/models") async def list_models(api_key: str = Depends(check_rate_limit)): """List available extraction and validation models.""" return { "extraction": EXTRACTION_MODELS, "validation": VALIDATION_MODELS, } @api_router.post("/api/v1/extract") async def extract( file: UploadFile = File(...), model: str = Form("Qwen/Qwen2.5-VL-72B-Instruct"), force_ocr: str = Form("false"), persist: str = Form("true"), use_case: str = Form("invoice"), org_id: str = Form(""), user_id: str = Form(""), api_key: str = Depends(check_rate_limit), ): """ Extract structured data from an uploaded invoice image or PDF. Returns all fields, line items, and bounding box coordinates. When persist=true and Supabase is configured, records the run. """ # ── Input validation ── file_size = file.size or 0 validate_file_upload(file.filename, file_size, file.content_type) safe_filename = sanitize_filename(file.filename) # ── Tier quota check (soft — skips if user_id not provided) ── if user_id: try: from api_usage import check_quota, increment_usage quota = check_quota(user_id, pages_requested=1) if not quota["allowed"]: raise HTTPException(status_code=429, detail=quota["message"]) except HTTPException: raise except Exception as e: logger.warning("Quota check failed (non-blocking): %s", e) start = time.time() use_ocr_only = force_ocr.lower() in ("true", "1", "yes") or model == "ocr-only" should_persist = persist.lower() in ("true", "1", "yes") # Determine the HF token to use — skip VLM when OCR-only token = "" if use_ocr_only else HF_TOKEN # ── Persistence: register document & open run ── doc_id = None run_id = None if should_persist: if not file_size and hasattr(file, "size") and file.size: file_size = file.size ext = os.path.splitext(safe_filename)[1].lower() file_type = "pdf" if ext == ".pdf" else "image" doc_id = save_document( filename=safe_filename, file_type=file_type, file_size=file_size, org_id=org_id or None, user_id=user_id or None, ) run_id = create_run( document_id=doc_id, use_case=use_case, org_id=org_id or None, user_id=user_id or None, ) # ── Stage 1: Upload (already complete at this point) ── upload_ms = int((time.time() - start) * 1000) if run_id: save_stage_result( run_id=run_id, stage_type="upload", stage_name="Document Upload", status="passed", sort_order=0, duration_ms=upload_ms, ) try: contents = await file.read() ext = os.path.splitext(safe_filename)[1].lower() # ── Stage 2: Extraction ── extraction_start = time.time() if ext == ".pdf": results = await asyncio.to_thread(extract_from_pdf, contents, hf_token=token, force_ocr=use_ocr_only, use_case=use_case, model_id=model) if not results: if run_id: _record_extraction_failure(run_id, extraction_start, "Could not extract data from PDF") raise HTTPException(status_code=422, detail="Could not extract data from PDF") result = results[0] # Tag page-0 bboxes for bb in result.bounding_boxes: bb.page = 0 # Collect all page images for multi-page preview page_images = [] for r in results: if r.page_image: page_images.append(r.page_image) # Merge line items and their bboxes from subsequent pages. base_li_count = len(result.line_items) for page_idx, r in enumerate(results[1:], start=1): for bb in r.bounding_boxes: if bb.field.startswith("line_item_"): li_idx = int(bb.field.split("_")[-1]) bb.field = f"line_item_{base_li_count + li_idx}" bb.page = page_idx result.bounding_boxes.append(bb) result.line_items.extend(r.line_items) base_li_count += len(r.line_items) result._page_images = page_images else: # Save to temp file and open as image with tempfile.NamedTemporaryFile(suffix=ext or ".jpg", delete=False) as tmp: tmp.write(contents) tmp_path = tmp.name try: image = Image.open(tmp_path) result = await asyncio.to_thread(extract_invoice, image, hf_token=token, force_ocr=use_ocr_only, use_case=use_case, model_id=model) finally: os.unlink(tmp_path) extraction_ms = int((time.time() - extraction_start) * 1000) total_ms = int((time.time() - start) * 1000) result.processing_time_ms = total_ms # Check for failure modes is_failure = result.extraction_method.startswith("failed") if run_id: extraction_status = "failed" if is_failure else "passed" save_stage_result( run_id=run_id, stage_type="extraction", stage_name="Data Extraction", status=extraction_status, sort_order=1, output=result.to_dict() if not is_failure else None, error_message=_describe_failure(result.extraction_method) if is_failure else None, duration_ms=extraction_ms, ) if is_failure: fail_data = result.to_dict() if hasattr(result, '_page_images') and result._page_images: fail_data["page_images"] = result._page_images if run_id: complete_run( run_id=run_id, status="completed", overall_result="failed", extraction_data=fail_data, processing_time_ms=total_ms, error_message=_describe_failure(result.extraction_method), ) return { **fail_data, "run_id": run_id, "error": _describe_failure(result.extraction_method), } # ── Success — mark run as passed (validation is a separate call) ── resp = result.to_dict() # Include all page images for multi-page PDF preview if hasattr(result, '_page_images') and result._page_images: resp["page_images"] = result._page_images if run_id: complete_run( run_id=run_id, status="completed", overall_result="passed", extraction_data=resp, processing_time_ms=total_ms, ) resp["run_id"] = run_id # ── Increment usage counter on successful extraction ── if user_id: try: from api_usage import increment_usage increment_usage(user_id, pages=1) except Exception as e: logger.warning("Usage increment failed (non-blocking): %s", e) # Fire webhook for successful extraction try: from api_webhooks import deliver_webhook deliver_webhook("extraction.completed", resp) except Exception: pass return resp except HTTPException: raise except Exception as e: logger.exception("Extraction failed") if run_id: total_ms = int((time.time() - start) * 1000) complete_run( run_id=run_id, status="completed", overall_result="failed", extraction_data={}, processing_time_ms=total_ms, error_message=str(e), ) raise HTTPException(status_code=500, detail=str(e)) @api_router.post("/api/v1/validate") async def validate( file: UploadFile = File(...), extraction: str = Form(...), model: str = Form("Qwen/Qwen2.5-VL-72B-Instruct"), run_id: str = Form(""), use_case: str = Form("_default"), api_key: str = Depends(check_rate_limit), ): """ Validate extraction results using dynamic math checks and VLM semantic verification. Checks are generated based on the document's use_case. """ # ── Input validation ── validate_file_upload(file.filename, file.size or 0, file.content_type) start = time.time() try: data = json.loads(extraction) except json.JSONDecodeError: raise HTTPException(status_code=400, detail="Invalid extraction JSON") # ── Dynamic math/consistency checks (use_case-aware) ── math_checks = run_dynamic_math_checks(data) # ── VLM semantic validation (cross-check against source image) ── semantic_checks = [] try: contents = await file.read() safe_name = sanitize_filename(file.filename) ext = os.path.splitext(safe_name)[1].lower() if ext == ".pdf": try: import fitz doc = fitz.open(stream=contents, filetype="pdf") page = doc[0] pix = page.get_pixmap(dpi=200) image = Image.frombytes("RGB", [pix.width, pix.height], pix.samples) doc.close() except Exception: image = None else: with tempfile.NamedTemporaryFile(suffix=ext or ".jpg", delete=False) as tmp: tmp.write(contents) tmp_path = tmp.name try: image = Image.open(tmp_path) finally: os.unlink(tmp_path) if image and HF_TOKEN: semantic_checks = await asyncio.to_thread( validate_with_vlm, image=image, extraction_data=data, hf_token=HF_TOKEN, model_id=model, use_case=use_case, ) except Exception as e: logger.warning(f"Semantic validation failed: {e}") semantic_checks = [{ "name": "VLM validation", "status": "skipped", "message": f"Could not run semantic validation: {str(e)[:100]}", }] total_ms = int((time.time() - start) * 1000) # Determine overall status from all check results all_checks = math_checks + semantic_checks statuses = [c["status"] for c in all_checks if c["status"] != "skipped"] if "fail" in statuses: overall = "invalid" elif "warn" in statuses: overall = "needs_review" elif statuses: overall = "valid" else: overall = "needs_review" # ── Persist validation stage ── if run_id: result_map = {"valid": "passed", "invalid": "failed", "needs_review": "needs_review"} save_stage_result( run_id=run_id, stage_type="validation", stage_name="Validation", status=result_map.get(overall, overall), sort_order=2, output={"math_checks": math_checks, "semantic_checks": semantic_checks, "overall": overall}, duration_ms=total_ms, ) complete_run( run_id=run_id, status="completed", overall_result=result_map.get(overall, overall), extraction_data=data, processing_time_ms=total_ms, ) validation_resp = { "overall_status": overall, "math_checks": math_checks, "semantic_checks": semantic_checks, "domain_checks": [], "cross_model_confidence": 0.0, "validation_model": model, "processing_time_ms": total_ms, } # Fire webhook for validation completion try: from api_webhooks import deliver_webhook deliver_webhook("validation.completed", validation_resp) except Exception: pass return validation_resp @api_router.get("/api/v1/health") async def health(api_key: str = Depends(check_rate_limit)): from db.supabase import is_enabled return { "status": "ok", "hf_token_set": bool(HF_TOKEN), "persistence_enabled": is_enabled(), } # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- def _record_extraction_failure(run_id: str, extraction_start: float, message: str): """Record a failed extraction stage and close the run.""" import time as _time extraction_ms = int((_time.time() - extraction_start) * 1000) save_stage_result( run_id=run_id, stage_type="extraction", stage_name="Data Extraction", status="failed", sort_order=1, error_message=message, duration_ms=extraction_ms, ) complete_run( run_id=run_id, status="completed", overall_result="failed", extraction_data={}, processing_time_ms=extraction_ms, error_message=message, ) def _describe_failure(method: str) -> str: if "no_token" in method: return "HF_TOKEN not set — AI Vision unavailable. Set the Space Secret to enable Qwen2.5-VL." if "vlm" in method: return "AI Vision models failed. Inference API may be loading or rate-limited. Try again shortly." return "Could not extract readable text. Image may be blurry or handwritten."