Fix startup and Gemini provider handling

services/extractor.py

@@ -1,4 +1,5 @@
 import base64
+import importlib.util
 import json
 import re
 from pathlib import Path
@@ -61,7 +62,7 @@ Rules:
 
 def _get_ai_verification_provider(preferred_provider: str = "") -> str:
     """Return a provider capable of vision reasoning for verification."""
-    providers = [p["id"] for p in get_available_providers() if p["id"] in {"claude", "gemini"}]
+    providers = [p["id"] for p in get_available_providers() if provider_supports_ai_verification(p["id"])]
     if preferred_provider in providers:
         return preferred_provider
     if DEFAULT_PROVIDER in providers:
@@ -71,6 +72,10 @@ def _get_ai_verification_provider(preferred_provider: str = "") -> str:
     raise ValueError("No AI verification provider configured")
 
 
+def provider_supports_ai_verification(provider_id: str) -> bool:
+    return provider_id == "claude" or provider_id == "gemini" or provider_id.startswith("gemini-")
+
+
 def _build_correlation_user_prompt(current_context: dict, candidate_context: dict) -> str:
     return (
         "Determine whether PAGE_A and PAGE_B belong to the same multi-page request/document for the same person. "
@@ -249,15 +254,29 @@ def _parse_json_lenient(text: str) -> dict:
     raise ValueError(f"Could not parse model response as JSON: {text[:200]}")
 
 
+# Individual Gemini models exposed to the UI
+_GEMINI_MODEL_OPTIONS = [
+    ("gemini-2.5-pro",        "Gemini 2.5 Pro (أفضل دقة)"),
+    ("gemini-2.5-flash",      "Gemini 2.5 Flash (سريع)"),
+    ("gemini-2.0-flash",      "Gemini 2.0 Flash"),
+    ("gemini-2.0-flash-lite", "Gemini 2.0 Flash Lite (احتياطي)"),
+]
+
+
+def _easyocr_is_available() -> bool:
+    return importlib.util.find_spec("easyocr") is not None
+
+
 def get_available_providers() -> list[dict]:
-    """Return list of all available providers (EasyOCR is always available)."""
-    providers = [
-        {"id": "easyocr", "name": "EasyOCR (مجاني)", "model": "local"},
-    ]
+    """Return only providers that are actually usable in the current runtime."""
+    providers = []
+    if _easyocr_is_available():
+        providers.append({"id": "easyocr", "name": "EasyOCR (مجاني)", "model": "local"})
     if ANTHROPIC_API_KEY:
         providers.append({"id": "claude", "name": "Claude (Anthropic)", "model": CLAUDE_MODEL})
     if GEMINI_API_KEY:
-        providers.append({"id": "gemini", "name": "Gemini (Google)", "model": GEMINI_MODEL})
+        for model_id, label in _GEMINI_MODEL_OPTIONS:
+            providers.append({"id": model_id, "name": label, "model": model_id})
     return providers
 
 
@@ -265,8 +284,12 @@ def get_default_provider() -> str:
     """Return the default provider, falling back to whichever is available."""
     providers = get_available_providers()
     if not providers:
-        return "easyocr"
+        return GEMINI_MODEL if DEFAULT_PROVIDER == "gemini" else DEFAULT_PROVIDER
     ids = [p["id"] for p in providers]
+    # Map legacy "gemini" default to the configured GEMINI_MODEL
+    resolved = GEMINI_MODEL if DEFAULT_PROVIDER == "gemini" else DEFAULT_PROVIDER
+    if resolved in ids:
+        return resolved
     if DEFAULT_PROVIDER in ids:
         return DEFAULT_PROVIDER
     return ids[0]
@@ -329,7 +352,7 @@ async def _extract_with_claude(image_path: str) -> dict:
 
 # ─── Gemini extraction ───────────────────────────────────────────
 
-async def _extract_with_gemini(image_path: str) -> dict:
+async def _extract_with_gemini(image_path: str, model: str = "") -> dict:
     import asyncio
     from google import genai
     from google.genai import types
@@ -345,8 +368,17 @@ async def _extract_with_gemini(image_path: str) -> dict:
                 ".png": "image/png", ".webp": "image/webp"}
     mime_type = mime_map.get(suffix, "image/jpeg")
 
-    # Try primary model first; only fall back on 503 (overloaded), NOT on 429 (quota)
-    models_to_try = [GEMINI_MODEL, "gemini-2.0-flash-lite"]
+    # Full fallback chain: start from requested model, cascade through cheaper/available ones
+    _ALL_GEMINI_FALLBACK = [
+        "gemini-2.5-pro",
+        "gemini-2.5-flash",
+        "gemini-2.0-flash",
+        "gemini-2.0-flash-lite",
+    ]
+    # If a specific model was requested, start from it; otherwise start from the configured default
+    start_model = model if model else GEMINI_MODEL
+    # Build ordered list: requested model first, then remaining fallbacks in order
+    models_to_try = [start_model] + [m for m in _ALL_GEMINI_FALLBACK if m != start_model]
 
     def _call(model_name: str):
         response = client.models.generate_content(
@@ -731,6 +763,8 @@ async def extract_document(image_path: str, provider: str = "") -> dict:
     """
     Extract structured data from a document image using the specified provider.
     Falls back to the default provider if none specified.
+    provider can be: "claude", "easyocr", "gemini" (uses GEMINI_MODEL),
+    or a specific Gemini model ID like "gemini-2.5-flash".
     """
     if not provider:
         provider = get_default_provider()
@@ -739,11 +773,13 @@ async def extract_document(image_path: str, provider: str = "") -> dict:
         if not ANTHROPIC_API_KEY:
             raise ValueError("ANTHROPIC_API_KEY not set")
         return await _extract_with_claude(image_path)
-    elif provider == "gemini":
-        if not GEMINI_API_KEY:
-            raise ValueError("GEMINI_API_KEY not set")
-        return await _extract_with_gemini(image_path)
     elif provider == "easyocr":
         return await _extract_with_easyocr(image_path)
+    elif provider == "gemini" or provider.startswith("gemini-"):
+        if not GEMINI_API_KEY:
+            raise ValueError("GEMINI_API_KEY not set")
+        # Pass the specific model if the provider ID encodes one
+        model_override = provider if provider.startswith("gemini-") else ""
+        return await _extract_with_gemini(image_path, model=model_override)
     else:
         raise ValueError(f"Unknown provider: {provider}")

services/pdf_handler.py

@@ -6,25 +6,35 @@ from pathlib import Path
 
 from config import UPLOAD_DIR
 
+# Maximum pages we will extract from a single PDF upload.
+# Beyond this limit pages are silently dropped to prevent runaway queue growth.
+MAX_PDF_PAGES = 200
+
 
 def pdf_to_images(pdf_bytes: bytes, original_name: str) -> list[dict]:
     """
     Convert a PDF to individual page images.
-    Returns list of dicts: [{"image_path": "relative/path.png", "page_number": 1}, ...]
+
+    Each page is rendered at 250 DPI — sufficient for Arabic OCR while keeping
+    file sizes manageable.  For scanned/image-only PDFs PyMuPDF simply renders
+    the embedded raster; no additional image extraction step is needed.
+
+    Returns list of dicts: [{"image_path": "relative/path.png", "page_number": 1, ...}, ...]
     """
     import fitz  # PyMuPDF
 
     doc = fitz.open(stream=pdf_bytes, filetype="pdf")
+    total_pages = len(doc)
     group_id = uuid.uuid4().hex
     today = date.today().isoformat()
     dest_dir = Path(UPLOAD_DIR) / today
     dest_dir.mkdir(parents=True, exist_ok=True)
 
     pages = []
-    for page_num in range(len(doc)):
+    for page_num in range(min(total_pages, MAX_PDF_PAGES)):
         page = doc[page_num]
-        # Render at 200 DPI for good OCR quality
-        pix = page.get_pixmap(dpi=200)
+        # 250 DPI gives sharper Arabic text while staying under ~2 MB per PNG
+        pix = page.get_pixmap(dpi=250)
 
         filename = f"{group_id}_p{page_num + 1}.png"
         dest = dest_dir / filename
@@ -34,6 +44,7 @@ def pdf_to_images(pdf_bytes: bytes, original_name: str) -> list[dict]:
             "image_path": str(Path(today) / filename),
             "page_number": page_num + 1,
             "pdf_group_id": group_id,
+            "total_pages": total_pages,
         })
 
     doc.close()