From 54dd52639d0c87f2a18f46b939b941c78a60fe4a Mon Sep 17 00:00:00 2001 From: Krikorios <99836218+Krikorios@users.noreply.github.com> Date: Thu, 7 May 2026 12:45:01 +0300 Subject: [PATCH] Fix startup and Gemini provider handling --- .gitignore | 2 + main.py | 6 +- requirements.txt | 1 + routers/auth.py | 48 +++++------ routers/review.py | 4 +- routers/upload.py | 177 +++++++++++++++++++++------------------- services/extractor.py | 64 +++++++++++---- services/pdf_handler.py | 19 ++++- 8 files changed, 189 insertions(+), 132 deletions(-) diff --git a/.gitignore b/.gitignore index bee1019..ae18f4a 100644 --- a/.gitignore +++ b/.gitignore @@ -7,6 +7,7 @@ backups/ __pycache__/ *.pyc .venv/ +.venv312/ venv/ *.egg-info/ .claude/ @@ -14,3 +15,4 @@ venv/ *.jpeg *.jpg *.png +tmp_samples/ diff --git a/main.py b/main.py index 1ec74b8..6df0810 100644 --- a/main.py +++ b/main.py @@ -12,6 +12,10 @@ from routers import documents, review, search, upload, auth from services.auth_service import get_session, cleanup_expired_sessions +Path(UPLOAD_DIR).mkdir(parents=True, exist_ok=True) +Path("data").mkdir(exist_ok=True) + + class CachedStaticFiles(StaticFiles): """StaticFiles that sets a long Cache-Control header so browsers don't re-download the same image on every refresh.""" @@ -29,8 +33,6 @@ class CachedStaticFiles(StaticFiles): @asynccontextmanager async def lifespan(app: FastAPI): - Path(UPLOAD_DIR).mkdir(parents=True, exist_ok=True) - Path("data").mkdir(exist_ok=True) create_tables() try: cleanup_expired_sessions() diff --git a/requirements.txt b/requirements.txt index a92888f..1fc4f33 100644 --- a/requirements.txt +++ b/requirements.txt @@ -5,6 +5,7 @@ google-genai>=1.70.0 easyocr>=1.7.0 pymupdf>=1.24.0 jinja2>=3.1.4 +python-dotenv>=1.0.1 python-multipart>=0.0.9 aiofiles>=23.0.0 pillow>=10.0.0 diff --git a/routers/auth.py b/routers/auth.py index 3990299..78af952 100644 --- a/routers/auth.py +++ b/routers/auth.py @@ -45,7 +45,7 @@ def require_admin(request: Request): @router.get("/login", response_class=HTMLResponse) async def login_get(request: Request): - return templates.TemplateResponse(request=request, name="login.html", context={"error": None}) + return templates.TemplateResponse("login.html", {"request": request, "error": None}) @router.post("/login", response_class=HTMLResponse) async def login_post(request: Request, username: str = Form(...), password: str = Form(...)): @@ -54,9 +54,8 @@ async def login_post(request: Request, username: str = Form(...), password: str if is_login_blocked(username, ip): return templates.TemplateResponse( - request=request, - name="login.html", - context={"error": "تم حجب محاولات تسجيل الدخول مؤقتاً. حاول بعد 15 دقيقة."}, + "login.html", + {"request": request, "error": "تم حجب محاولات تسجيل الدخول مؤقتاً. حاول بعد 15 دقيقة."}, status_code=429, ) @@ -64,9 +63,8 @@ async def login_post(request: Request, username: str = Form(...), password: str if not user or not verify_password(user["password_hash"], password): record_login_attempt(username, ip, success=False) return templates.TemplateResponse( - request=request, - name="login.html", - context={"error": "Invalid username or password"}, + "login.html", + {"request": request, "error": "Invalid username or password"}, status_code=401, ) @@ -109,9 +107,8 @@ async def logout(request: Request): async def users_list(request: Request, current=Depends(require_admin)): users = get_all_users() return templates.TemplateResponse( - request=request, - name="users.html", - context={"users": users, "current_user": current}, + "users.html", + {"request": request, "users": users, "current_user": current}, ) @router.post("/users/create") @@ -128,16 +125,14 @@ async def add_user( if len(password) < 8: users = get_all_users() return templates.TemplateResponse( - request=request, - name="users.html", - context={"users": users, "current_user": current, "error": "Password must be at least 8 characters."}, + "users.html", + {"request": request, "users": users, "current_user": current, "error": "Password must be at least 8 characters."}, ) if get_user_by_username(username): users = get_all_users() return templates.TemplateResponse( - request=request, - name="users.html", - context={"users": users, "current_user": current, "error": f"User '{username}' already exists."}, + "users.html", + {"request": request, "users": users, "current_user": current, "error": f"User '{username}' already exists."}, ) create_user(username, password, role=role) return RedirectResponse(url="/auth/users", status_code=status.HTTP_303_SEE_OTHER) @@ -147,9 +142,8 @@ async def remove_user(request: Request, user_id: int, current=Depends(require_ad if user_id == current["user_id"]: users = get_all_users() return templates.TemplateResponse( - request=request, - name="users.html", - context={"users": users, "current_user": current, "error": "لا يمكنك حذف حسابك الخاص."}, + "users.html", + {"request": request, "users": users, "current_user": current, "error": "لا يمكنك حذف حسابك الخاص."}, status_code=400, ) # Prevent removing the last admin @@ -157,9 +151,8 @@ async def remove_user(request: Request, user_id: int, current=Depends(require_ad if target and target["role"] == "admin" and count_admins() <= 1: users = get_all_users() return templates.TemplateResponse( - request=request, - name="users.html", - context={"users": users, "current_user": current, "error": "لا يمكن حذف آخر مسؤول في النظام."}, + "users.html", + {"request": request, "users": users, "current_user": current, "error": "لا يمكن حذف آخر مسؤول في النظام."}, status_code=400, ) delete_user(user_id) @@ -180,9 +173,8 @@ async def change_role( if target and target["role"] == "admin" and role != "admin" and count_admins() <= 1: users = get_all_users() return templates.TemplateResponse( - request=request, - name="users.html", - context={"users": users, "current_user": current, "error": "لا يمكن تخفيض رتبة آخر مسؤول."}, + "users.html", + {"request": request, "users": users, "current_user": current, "error": "لا يمكن تخفيض رتبة آخر مسؤول."}, status_code=400, ) set_user_role(user_id, role) @@ -196,9 +188,9 @@ async def backup_db(request: Request, current=Depends(require_admin)): except Exception as e: msg = f"Backup failed: {str(e)}" return templates.TemplateResponse( - request=request, - name="users.html", - context={ + "users.html", + { + "request": request, "users": get_all_users(), "current_user": current, "backup_msg": msg, diff --git a/routers/review.py b/routers/review.py index 9aba7b0..27b2888 100644 --- a/routers/review.py +++ b/routers/review.py @@ -12,6 +12,7 @@ from services.extractor import ( extract_document, get_available_providers, get_default_provider, + provider_supports_ai_verification, verify_page_correlation, ) from services.search_service import normalize_arabic, _normalize_scope @@ -336,7 +337,8 @@ async def review_document(request: Request, doc_id: int, wait: int = 0): "providers": get_available_providers(), "current_provider": doc.get("provider") or get_default_provider(), "ai_verification_available": any( - provider["id"] in {"claude", "gemini"} for provider in get_available_providers() + provider_supports_ai_verification(provider["id"]) + for provider in get_available_providers() ), }, ) diff --git a/routers/upload.py b/routers/upload.py index b53e46d..5115fee 100644 --- a/routers/upload.py +++ b/routers/upload.py @@ -9,7 +9,7 @@ from fastapi import APIRouter, File, Form, HTTPException, Request, UploadFile from fastapi.responses import JSONResponse, RedirectResponse from fastapi.templating import Jinja2Templates -from config import UPLOAD_DIR +from config import UPLOAD_DIR, MAX_CONCURRENT_EXTRACTIONS from database.connection import get_db from services.extractor import extract_document, get_available_providers, get_default_provider from services.pdf_handler import pdf_to_images @@ -17,6 +17,16 @@ from services.pdf_handler import pdf_to_images router = APIRouter() templates = Jinja2Templates(directory="templates") +# Semaphore: limits simultaneous AI API calls to avoid quota exhaustion on bulk uploads +_extraction_semaphore: asyncio.Semaphore | None = None + + +def _get_semaphore() -> asyncio.Semaphore: + global _extraction_semaphore + if _extraction_semaphore is None: + _extraction_semaphore = asyncio.Semaphore(MAX_CONCURRENT_EXTRACTIONS) + return _extraction_semaphore + ALLOWED_IMAGE_EXTS = {".jpg", ".jpeg", ".png", ".webp"} ALLOWED_PDF_EXTS = {".pdf"} ALLOWED_EXTENSIONS = ALLOWED_IMAGE_EXTS | ALLOWED_PDF_EXTS @@ -91,95 +101,96 @@ def _save_image(file_bytes: bytes, original_name: str) -> str: async def _extract_and_save(doc_id: int, image_path: str, provider: str = ""): """Background task: call extractor, parse result, update DB.""" - try: - data = await extract_document(image_path, provider=provider) - raw_json = json.dumps(data, ensure_ascii=False) + async with _get_semaphore(): + try: + data = await extract_document(image_path, provider=provider) + raw_json = json.dumps(data, ensure_ascii=False) - with get_db() as conn: - conn.execute( - """UPDATE documents SET status='extracted', - raw_extraction_json=?, - request_number=?, - request_date=?, - applicant_name_raw=?, - request_purpose=?, - data_valid_until=?, - registry_office=?, - owns_properties=?, - declared_property_count=?, - page_info=?, - search_scope=?, - updated_at=CURRENT_TIMESTAMP - WHERE id=?""", - ( - raw_json, - data.get("request_number"), - data.get("request_date"), - data.get("applicant_name_raw"), - data.get("request_purpose"), - data.get("data_valid_until"), - data.get("registry_office"), - data.get("owns_properties"), - data.get("declared_property_count"), - data.get("page_info"), - data.get("search_scope"), - doc_id, - ), - ) - - for i, prop in enumerate(data.get("properties", [])): + with get_db() as conn: conn.execute( - """INSERT INTO properties - (document_id, row_order, party_name, property_number, - section, block, real_estate_district, qaza, num_shares, ownership_type) - VALUES (?,?,?,?,?,?,?,?,?,?)""", + """UPDATE documents SET status='extracted', + raw_extraction_json=?, + request_number=?, + request_date=?, + applicant_name_raw=?, + request_purpose=?, + data_valid_until=?, + registry_office=?, + owns_properties=?, + declared_property_count=?, + page_info=?, + search_scope=?, + updated_at=CURRENT_TIMESTAMP + WHERE id=?""", ( - doc_id, i, - prop.get("party_name"), - prop.get("property_number"), - prop.get("section"), - prop.get("block"), - prop.get("real_estate_district"), - prop.get("qaza"), - prop.get("num_shares"), - prop.get("ownership_type"), + raw_json, + data.get("request_number"), + data.get("request_date"), + data.get("applicant_name_raw"), + data.get("request_purpose"), + data.get("data_valid_until"), + data.get("registry_office"), + data.get("owns_properties"), + data.get("declared_property_count"), + data.get("page_info"), + data.get("search_scope"), + doc_id, ), ) - # Flag logical duplicates: same request_number + search_scope + page_number - req_num = (data.get("request_number") or "").strip() - scope = (data.get("search_scope") or "").strip() - page_info = (data.get("page_info") or "").strip() - if req_num: - existing = conn.execute( - """SELECT id FROM documents - WHERE id != ? AND request_number=? - AND COALESCE(search_scope,'')=? - AND COALESCE(page_info,'')=? - AND status IN ('extracted','confirmed') - AND COALESCE(duplicate_dismissed, 0) = 0 - ORDER BY id LIMIT 1""", - (doc_id, req_num, scope, page_info), - ).fetchone() - if existing: - # Only auto-flag if THIS document hasn't itself been dismissed. - self_row = conn.execute( - "SELECT COALESCE(duplicate_dismissed, 0) AS d FROM documents WHERE id=?", - (doc_id,), - ).fetchone() - if not (self_row and self_row["d"]): - conn.execute( - "UPDATE documents SET duplicate_of=? WHERE id=?", - (existing["id"], doc_id), - ) + for i, prop in enumerate(data.get("properties", [])): + conn.execute( + """INSERT INTO properties + (document_id, row_order, party_name, property_number, + section, block, real_estate_district, qaza, num_shares, ownership_type) + VALUES (?,?,?,?,?,?,?,?,?,?)""", + ( + doc_id, i, + prop.get("party_name"), + prop.get("property_number"), + prop.get("section"), + prop.get("block"), + prop.get("real_estate_district"), + prop.get("qaza"), + prop.get("num_shares"), + prop.get("ownership_type"), + ), + ) - except Exception as e: - with get_db() as conn: - conn.execute( - """UPDATE documents SET status='error', extraction_error=?, - updated_at=CURRENT_TIMESTAMP WHERE id=?""", - (str(e), doc_id), - ) + # Flag logical duplicates: same request_number + search_scope + page_number + req_num = (data.get("request_number") or "").strip() + scope = (data.get("search_scope") or "").strip() + page_info = (data.get("page_info") or "").strip() + if req_num: + existing = conn.execute( + """SELECT id FROM documents + WHERE id != ? AND request_number=? + AND COALESCE(search_scope,'')=? + AND COALESCE(page_info,'')=? + AND status IN ('extracted','confirmed') + AND COALESCE(duplicate_dismissed, 0) = 0 + ORDER BY id LIMIT 1""", + (doc_id, req_num, scope, page_info), + ).fetchone() + if existing: + # Only auto-flag if THIS document hasn't itself been dismissed. + self_row = conn.execute( + "SELECT COALESCE(duplicate_dismissed, 0) AS d FROM documents WHERE id=?", + (doc_id,), + ).fetchone() + if not (self_row and self_row["d"]): + conn.execute( + "UPDATE documents SET duplicate_of=? WHERE id=?", + (existing["id"], doc_id), + ) + + except Exception as e: + with get_db() as conn: + conn.execute( + """UPDATE documents SET status='error', extraction_error=?, + updated_at=CURRENT_TIMESTAMP WHERE id=?""", + (str(e), doc_id), + ) @router.get("/") diff --git a/services/extractor.py b/services/extractor.py index eda2945..b4324d6 100644 --- a/services/extractor.py +++ b/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}") diff --git a/services/pdf_handler.py b/services/pdf_handler.py index 91de74e..052f8a4 100644 --- a/services/pdf_handler.py +++ b/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()