api.py

import os
import io
import re
import json
import uuid
from difflib import SequenceMatcher
from typing import Optional, Dict, Any
from fastapi.responses import RedirectResponse
from fastapi import (
    FastAPI,
    File,
    UploadFile,
    HTTPException
)

from PIL import Image
from dotenv import load_dotenv
from pydantic import BaseModel

# OCR pipeline
import app

# =========================================================
# LOAD ENVIRONMENT
# =========================================================
load_dotenv(override=True)

# =========================================================
# FALLBACK MODELS
# =========================================================
# Tried in order when a model-level failure occurs (high demand / overload /
# model not found / unsupported). Imported from app.py so there is a single
# source of truth; if app.py doesn't define it for some reason, fall back to
# a locally defined copy so this module still works standalone.
try:
    FALLBACK_MODELS = app.FALLBACK_MODELS
except AttributeError:
    FALLBACK_MODELS = [
        "gemini-3.1-pro-preview",          # Primary choice: latest and most powerful
        "gemini-3.0-pro",                  # Fallback 1: previous generation Pro, more stable
        "gemini-2.5-flash-image-preview",  # Fallback 2: Flash version, fast speed
        "gemini-2.5-flash",                # Ultimate fallback: most stable Flash
    ]

# =========================================================
# GEMINI API KEY + MODEL ROTATION STATE
# =========================================================
# These remember the last combination that worked, so the next request
# starts from there instead of always retrying everything from scratch.
CURRENT_KEY_INDEX = 0
CURRENT_MODEL_INDEX = 0


def get_gemini_api_keys():
    """
    Load all Gemini keys from environment variables / HF Secrets.

    Example:
        GEMINI_API_KEY_1
        GEMINI_API_KEY_2
        GEMINI_API_KEY_3
    """
    keys = []

    for var, val in os.environ.items():
        if var.startswith("GEMINI_API_KEY") and val:
            masked = "*" * max(len(val.strip()) - 4, 0) + val.strip()[-4:]
            print(f"using this API key: {var} = {masked}")
            keys.append(val.strip())

    keys.sort()

    return keys


def is_quota_error(error_message: str) -> bool:
    """
    Detect Gemini quota / rate-limit errors. These are tied to a specific
    API key, so the correct response is to rotate to the NEXT KEY while
    keeping the same model.
    """
    error_message = error_message.lower()

    quota_patterns = [
        "quota",
        "rate limit",
        "resource exhausted",
        "429",
        "too many requests",
        "limit exceeded",
        "exceeded your current quota",
        "billing",
        "401", 
        "unauthenticated", 
        "invalid authentication", 
        "api key not valid"
    ]

    return any(pattern in error_message for pattern in quota_patterns)


def is_high_demand_error(error_message: str) -> bool:
    """
    Detect Gemini overload / high-demand / model-unavailable errors. These
    are tied to the MODEL itself (the model is overloaded or doesn't exist
    on this account/region), so the correct response is to move to the
    NEXT FALLBACK MODEL rather than trying more keys against the same
    broken/overloaded model.
    """
    error_message = error_message.lower()

    high_demand_patterns = [
        "high demand",
        "overloaded",
        "service unavailable",
        "503",
        "model is overloaded",
        "currently unavailable",
        "try again later",
        "unavailable. please retry",
        "is unavailable",          # raised by app.run_gemini_correction for unknown models
        "404",
        "not found",
        "is not supported",
        "not supported for",
        "invalid model",
        "unknown model",
        "repeated/duplicate coordinates",  # raised by app.run_pipeline when coordinates_suspect is True
    ]

    return any(pattern in error_message for pattern in high_demand_patterns)


def run_pipeline_with_fallback(image, progress):
    """
    Run app.run_pipeline with automatic rotation across both Gemini API keys
    and fallback models.

    Strategy:
      - Outer loop: iterate over FALLBACK_MODELS, starting from the last
        known-good model index.
      - Inner loop: for each model, iterate over all API keys, starting
        from the last known-good key index.
      - Quota/rate-limit error on a key -> try the next key, same model.
      - High-demand/overload/model-not-found error -> abandon remaining
        keys for this model and move straight to the next fallback model.
      - Any other/unrecognized error -> raised immediately, no further
        rotation (it's not something key or model rotation can fix).
      - On success, remember which key + model worked so the next request
        starts there instead of from the top.
    """
    global CURRENT_KEY_INDEX, CURRENT_MODEL_INDEX

    api_keys = get_gemini_api_keys()

    if not api_keys:
        raise Exception("No Gemini API keys configured.")

    total_keys = len(api_keys)
    total_models = len(FALLBACK_MODELS)

    last_error: Optional[Exception] = None

    for model_offset in range(total_models):

        model_idx = (CURRENT_MODEL_INDEX + model_offset) % total_models
        model_name = FALLBACK_MODELS[model_idx]

        print(f"[Gemini] Trying model {model_idx + 1}/{total_models}: {model_name}")

        model_failed_high_demand = False

        for key_offset in range(total_keys):

            idx = (CURRENT_KEY_INDEX + key_offset) % total_keys
            api_key = api_keys[idx]

            try:
                print(f"[Gemini] Trying key {idx + 1}/{total_keys} with model {model_name}")

                result = app.run_pipeline(
                    image=image,
                    ocr_engine="PaddleOCR",
                    api_key=api_key,
                    model=model_name,
                    progress=progress
                )

                # Success — remember this combination for next time.
                CURRENT_KEY_INDEX = idx
                CURRENT_MODEL_INDEX = model_idx

                print(f"[Gemini] Success with key {idx + 1} on model {model_name}")

                return result

            except Exception as e:

                last_error = e
                error_message = str(e)

                print(f"[Gemini] Key {idx + 1} on model {model_name} failed: {error_message}")

                if is_quota_error(error_message):
                    print("[Gemini] Quota/rate-limit detected. Trying next key...")
                    continue

                if is_high_demand_error(error_message):
                    print(
                        "[Gemini] High demand / overload / unavailable model "
                        "detected. Switching to next fallback model..."
                    )
                    model_failed_high_demand = True
                    break

                # Unrecognized error type — don't keep blindly rotating keys
                # or models for something rotation can't fix (e.g. a bad
                # image, malformed schema, programming error).
                raise

        if model_failed_high_demand:
            continue

        print(f"[Gemini] All keys exhausted (quota errors) for model {model_name}. Trying next fallback model...")

    raise Exception(
        f"All Gemini API keys and fallback models failed. Last error: {last_error}"
    )


# =========================================================
# FASTAPI APP
# =========================================================
app_api = FastAPI(
    title="Route OCR Extraction API",
    version="4.0.0"
)

# =========================================================
# MEMORY STORAGE
# =========================================================
uploaded_files: Dict[str, bytes] = {}

# =========================================================
# RESPONSE MODELS
# =========================================================
class RootResponse(BaseModel):
    message: str


class UploadResponse(BaseModel):
    file_id: str
    filename: str
    message: str


class ExtractRequest(BaseModel):
    file_id: str


class ExtractResponse(BaseModel):
    status: str
    message: str
    extracted_data: Optional[Dict[str, Any]] = None
    debug: Optional[Any] = None
    route_document_detected: bool = False
    confidence_score: int = 0
    error: Optional[str] = None


class EvaluateRequest(BaseModel):
    file_id: str
    ground_truth: str


class EvaluateResponse(BaseModel):
    status: str
    message: str
    accuracy_score: float
    route_document_detected: bool
    confidence_score: int
    extracted_text: str
    ground_truth: str
    debug: Optional[Any] = None


class ErrorResponse(BaseModel):
    detail: str


# =========================================================
# SAFE PROGRESS CALLBACK
# =========================================================
def empty_progress(*args, **kwargs):
    pass


# =========================================================
# SAFELY CAUGHT BY EXCEPTION NOW (Validate logic moved to app.py)
# =========================================================


# =========================================================
# OCR ACCURACY EVALUATION
# =========================================================
def evaluate_ocr_accuracy(extracted_text: str, ground_truth: str):
    """
    Compare OCR extracted text with actual text.
    """

    # Normalize text
    extracted_text = extracted_text.lower().strip()
    ground_truth = ground_truth.lower().strip()

    # Remove extra spaces
    extracted_text = re.sub(r'\s+', ' ', extracted_text)
    ground_truth = re.sub(r'\s+', ' ', ground_truth)

    # Similarity ratio
    similarity = SequenceMatcher(None, extracted_text, ground_truth).ratio()

    accuracy_score = round(similarity * 100, 2)

    return accuracy_score


# =========================================================
# ROOT ENDPOINT
# =========================================================
@app_api.get("/")
def root():
    return RedirectResponse(url="/docs")


@app_api.get("/health")
async def health():
    return {"message": "OK"}


# =========================================================
# UPLOAD ENDPOINT
# =========================================================
@app_api.post(
    "/upload",
    response_model=UploadResponse,
    responses={
        500: {"model": ErrorResponse}
    }
)
async def upload_document(file: UploadFile = File(...)):

    try:
        file_id = str(uuid.uuid4())

        contents = await file.read()

        uploaded_files[file_id] = contents

        return UploadResponse(
            file_id=file_id,
            filename=file.filename or "unknown",
            message="File uploaded successfully."
        )

    except Exception as e:
        raise HTTPException(
            status_code=500,
            detail=f"Upload failed: {str(e)}"
        )


# =========================================================
# EXTRACT ENDPOINT
# =========================================================
@app_api.post(
    "/extract",
    response_model=ExtractResponse,
    responses={
        400: {"model": ErrorResponse},
        404: {"model": ErrorResponse},
        500: {"model": ErrorResponse}
    }
)
async def extract_data(payload: ExtractRequest):

    try:
        file_id = payload.file_id

        # -------------------------------------------------
        # Validate file exists
        # -------------------------------------------------

        if file_id not in uploaded_files:
            raise HTTPException(
                status_code=404,
                detail="File ID not found."
            )

        # -------------------------------------------------
        # Load image
        # -------------------------------------------------

        contents = uploaded_files[file_id]

        image_pil = Image.open(io.BytesIO(contents)).convert("RGB")

        # -------------------------------------------------
        # Run OCR pipeline with key + model fallback
        # -------------------------------------------------

        json_result_str, debug_info = run_pipeline_with_fallback(
            image=image_pil,
            progress=empty_progress
        )

        # -------------------------------------------------
        # Build searchable text
        # -------------------------------------------------

        searchable_text = ""

        if isinstance(debug_info, str):
            searchable_text += debug_info

        searchable_text += str(json_result_str)

        # -------------------------------------------------
        # Parse extracted JSON
        # -------------------------------------------------

        extracted_data = json.loads(json_result_str)

        if "error" in extracted_data:
            raise HTTPException(
                status_code=400,
                detail=extracted_data["error"]
            )

        # -------------------------------------------------
        # Success response
        # -------------------------------------------------

        return ExtractResponse(
            status="success",
            message="Route extraction completed successfully.",
            extracted_data=extracted_data,
            debug=debug_info,
            route_document_detected=True,
            confidence_score=100
        )

    except HTTPException:
        raise

    except Exception as e:
        import traceback
        traceback.print_exc()

        if "Validation failed" in str(e) or "Irrelevant document" in str(e):
            error_detail = str(e).split("Validation failed: ")[-1] if "Validation failed: " in str(e) else str(e)
            raise HTTPException(
                status_code=400,
                detail=error_detail
            )

        raise HTTPException(
            status_code=500,
            detail=str(e)
        )


# =========================================================
# EVALUATE OCR ENDPOINT
# =========================================================
@app_api.post(
    "/evaluate",
    response_model=EvaluateResponse,
    responses={
        400: {"model": ErrorResponse},
        404: {"model": ErrorResponse},
        500: {"model": ErrorResponse}
    }
)
async def evaluate_ocr(payload: EvaluateRequest):

    try:
        file_id = payload.file_id
        ground_truth = payload.ground_truth

        # -------------------------------------------------
        # Validate file exists
        # -------------------------------------------------

        if file_id not in uploaded_files:
            raise HTTPException(
                status_code=404,
                detail="File ID not found."
            )

        # -------------------------------------------------
        # Load image
        # -------------------------------------------------

        contents = uploaded_files[file_id]

        image_pil = Image.open(io.BytesIO(contents)).convert("RGB")

        # -------------------------------------------------
        # Run OCR pipeline with key + model fallback
        # -------------------------------------------------

        json_result_str, debug_info = run_pipeline_with_fallback(
            image=image_pil,
            progress=empty_progress
        )

        # -------------------------------------------------
        # Extract searchable OCR text
        # -------------------------------------------------

        extracted_text = ""

        if isinstance(debug_info, str):
            extracted_text += debug_info

        extracted_text += str(json_result_str)

        extracted_data = json.loads(json_result_str)

        if "error" in extracted_data:
            raise HTTPException(
                status_code=400,
                detail=extracted_data["error"]
            )

        # -------------------------------------------------
        # Calculate OCR accuracy
        # -------------------------------------------------

        accuracy_score = evaluate_ocr_accuracy(
            extracted_text=extracted_text,
            ground_truth=ground_truth
        )

        # -------------------------------------------------
        # Return evaluation response
        # -------------------------------------------------

        return EvaluateResponse(
            status="success",
            message="OCR evaluation completed successfully.",
            accuracy_score=accuracy_score,
            route_document_detected=True,
            confidence_score=100,
            extracted_text=extracted_text,
            ground_truth=ground_truth,
            debug=debug_info
        )

    except HTTPException:
        raise

    except Exception as e:
        import traceback
        traceback.print_exc()

        if "Validation failed" in str(e) or "Irrelevant document" in str(e):
            error_detail = str(e).split("Validation failed: ")[-1] if "Validation failed: " in str(e) else str(e)
            raise HTTPException(
                status_code=400,
                detail=error_detail
            )

        raise HTTPException(
            status_code=500,
            detail=str(e)
        )


# =========================================================
# MAIN
# =========================================================
if __name__ == "__main__":
    import uvicorn

    uvicorn.run(
        app_api,
        host="0.0.0.0",
        port=8000
    )