""" DocuLens — Batch processing endpoints. Accepts multiple files in a single request, processes them sequentially, and returns per-file results with a batch summary. """ import os import json import time import tempfile import logging from typing import Optional from uuid import uuid4 from fastapi import APIRouter, UploadFile, File, Form, HTTPException, Depends from PIL import Image from extraction import ( extract_invoice, extract_from_pdf, ExtractionResult, ) from db.supabase import save_document, create_run, complete_run, save_stage_result from middleware import check_rate_limit, validate_file_upload, sanitize_filename from api_webhooks import deliver_webhook logger = logging.getLogger(__name__) batch_router = APIRouter() HF_TOKEN = os.environ.get("HF_TOKEN", "") # Maximum files per batch (free tier) MAX_BATCH_SIZE = int(os.environ.get("MAX_BATCH_SIZE", "20")) # --------------------------------------------------------------------------- # Batch processing endpoint # --------------------------------------------------------------------------- @batch_router.post("/api/v1/batch/extract") async def batch_extract( files: list[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 multiple uploaded files. Returns a batch result with per-file status, extraction data, and a summary of passed/failed counts. """ if not files: raise HTTPException(status_code=400, detail="No files provided") # ── Tier-based batch limits ── effective_max = MAX_BATCH_SIZE if user_id: try: from api_usage import check_feature, check_quota, TIER_FEATURES, get_user_tier if not check_feature(user_id, "batch_upload"): raise HTTPException( status_code=403, detail="Batch upload is not available on the free plan. Upgrade to Starter or above.", ) tier = get_user_tier(user_id) tier_max = TIER_FEATURES.get(tier, {}).get("max_batch_files", MAX_BATCH_SIZE) effective_max = min(tier_max, MAX_BATCH_SIZE) if tier_max else MAX_BATCH_SIZE # Check quota for all files in the batch quota = check_quota(user_id, pages_requested=len(files)) if not quota["allowed"]: raise HTTPException(status_code=429, detail=quota["message"]) except HTTPException: raise except Exception as e: logger.warning("Batch tier check failed (non-blocking): %s", e) if len(files) > effective_max: raise HTTPException( status_code=400, detail=f"Batch size exceeds limit. Maximum {effective_max} files per batch.", ) batch_id = str(uuid4()) batch_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") token = "" if use_ocr_only else HF_TOKEN results = [] passed = 0 failed = 0 for i, file in enumerate(files): file_start = time.time() safe_filename = sanitize_filename(file.filename) file_result = { "index": i, "filename": safe_filename, "status": "pending", "run_id": None, "data": None, "error": None, "processing_time_ms": 0, } try: # Validate each file file_size = file.size or 0 validate_file_upload(file.filename, file_size, file.content_type) # Persistence doc_id = None run_id = None if should_persist: 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, ) file_result["run_id"] = run_id # Read file contents contents = await file.read() ext = os.path.splitext(safe_filename)[1].lower() # Extract if ext == ".pdf": page_results = extract_from_pdf( contents, hf_token=token, force_ocr=use_ocr_only, use_case=use_case, model_id=model, ) if not page_results: raise ValueError("Could not extract data from PDF") result = page_results[0] # Merge multi-page line items base_li = len(result.line_items) for pi, pr in enumerate(page_results[1:], start=1): for bb in pr.bounding_boxes: if bb.field.startswith("line_item_"): li_idx = int(bb.field.split("_")[-1]) bb.field = f"line_item_{base_li + li_idx}" bb.page = pi result.bounding_boxes.append(bb) result.line_items.extend(pr.line_items) base_li += len(pr.line_items) 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) result = extract_invoice( image, hf_token=token, force_ocr=use_ocr_only, use_case=use_case, model_id=model, ) finally: os.unlink(tmp_path) file_ms = int((time.time() - file_start) * 1000) result.processing_time_ms = file_ms is_failure = result.extraction_method.startswith("failed") if is_failure: file_result["status"] = "failed" file_result["error"] = _describe_failure(result.extraction_method) failed += 1 else: file_result["status"] = "passed" file_result["data"] = result.to_dict() # Remove large base64 images from batch response to keep payload small file_result["data"].pop("page_image", None) passed += 1 # Increment usage for successful extraction if user_id: try: from api_usage import increment_usage increment_usage(user_id, pages=1) except Exception: pass file_result["processing_time_ms"] = file_ms # Persist run result if run_id: save_stage_result( run_id=run_id, stage_type="extraction", stage_name="Data Extraction", status="failed" if is_failure else "passed", sort_order=1, output=result.to_dict() if not is_failure else None, error_message=file_result.get("error"), duration_ms=file_ms, ) complete_run( run_id=run_id, status="completed", overall_result="failed" if is_failure else "passed", extraction_data=result.to_dict() if not is_failure else {}, processing_time_ms=file_ms, error_message=file_result.get("error"), ) except HTTPException as he: file_result["status"] = "failed" file_result["error"] = he.detail file_result["processing_time_ms"] = int((time.time() - file_start) * 1000) failed += 1 except Exception as e: logger.warning(f"Batch file {i} ({safe_filename}) failed: {e}") file_result["status"] = "failed" file_result["error"] = str(e) file_result["processing_time_ms"] = int((time.time() - file_start) * 1000) failed += 1 results.append(file_result) batch_ms = int((time.time() - batch_start) * 1000) batch_response = { "batch_id": batch_id, "total": len(files), "passed": passed, "failed": failed, "processing_time_ms": batch_ms, "results": results, } # Fire webhook try: deliver_webhook("batch.completed", batch_response) except Exception: pass return batch_response # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- def _describe_failure(method: str) -> str: if "no_token" in method: return "HF_TOKEN not set — AI Vision unavailable." if "vlm" in method: return "AI Vision models failed. Inference API may be loading or rate-limited." return "Could not extract readable text. Image may be blurry or handwritten."