src/api.py

# src/api.py

from fastapi import FastAPI, UploadFile, File, HTTPException, BackgroundTasks
from fastapi.responses import JSONResponse
import shutil
import os
from pathlib import Path
import uuid
import sys

from src.pipeline import process_invoice
from src.schema import InvoiceData 

app = FastAPI(
    title="Invoice Extraction API",
    description="Hybrid ML + Rule-Based Pipeline with LayoutLMv3",
    version="2.0"
)

# Create temp folder if not exists
UPLOAD_DIR = Path("temp_uploads")
UPLOAD_DIR.mkdir(exist_ok=True)

def cleanup_file(path: str):
    """Background task to remove temp file after processing"""
    try:
        if os.path.exists(path):
            os.remove(path)
    except Exception as e:
        print(f"Error cleaning up {path}: {e}")

@app.post("/extract", response_model=InvoiceData) # <--- CONTRACT ENFORCED
async def extract_invoice(
    background_tasks: BackgroundTasks,
    file: UploadFile = File(...)
):
    """
    Upload an invoice (PDF/JPG/PNG) and get structured data.
    """
    # 1. Generate unique filename to prevent collisions
    file_ext = Path(file.filename).suffix
    unique_name = f"{uuid.uuid4()}{file_ext}"
    temp_path = UPLOAD_DIR / unique_name
    
    try:
        # 2. Save Uploaded File
        with open(temp_path, "wb") as buffer:
            shutil.copyfileobj(file.file, buffer)
            
        # 3. Process Logic
        result = process_invoice(str(temp_path), method='ml')
        
        # 4. Cleanup
        # We use background_tasks to delete the file AFTER the response is sent
        background_tasks.add_task(cleanup_file, str(temp_path))
        
        return result

    except Exception as e:
        # Cleanup even on error
        cleanup_file(str(temp_path))
        raise HTTPException(status_code=500, detail=str(e))