import os import io import re import base64 import gc import tempfile import uuid import asyncio from typing import List, Dict, Optional, Tuple from collections import Counter, deque from concurrent.futures import ThreadPoolExecutor from threading import Lock, Thread, Event import time import logging from fastapi import FastAPI, File, UploadFile, Form, HTTPException, BackgroundTasks from fastapi.middleware.cors import CORSMiddleware from fastapi.responses import JSONResponse from starlette.requests import Request import fitz # PyMuPDF import google.generativeai as genai from google.api_core import exceptions as google_exceptions from PIL import Image import requests import base64 # Azure Blob Storage try: from azure.storage.blob import ( BlobServiceClient, generate_blob_sas, BlobSasPermissions, ContentSettings ) AZURE_AVAILABLE = True except ImportError: AZURE_AVAILABLE = False print("Warning: azure-storage-blob not installed. Run: pip install azure-storage-blob") # Google Gemini - optional import try: GEMINI_AVAILABLE = True except ImportError: GEMINI_AVAILABLE = False print("Warning: google-generativeai not installed. Image-based PDFs won't be supported.") from datetime import datetime, timedelta # Configure logging logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) app = FastAPI(title="Invoice Splitter API with Azure Blob Storage - Optimized") # Increase request body size limit Request.max_body_size = 200 * 1024 * 1024 # 200MB app.add_middleware( CORSMiddleware, allow_origins=["*"], allow_credentials=True, allow_methods=["*"], allow_headers=["*"], ) # ============================================================================ # ⭐ CONFIGURATION # ============================================================================ # Gemini API Key - REQUIRED for image-based PDFs GEMINI_API_KEY = os.environ.get("GEMINI_API_KEY", "") # Azure Blob Storage Configuration - REQUIRED for blob storage AZURE_STORAGE_CONNECTION_STRING = os.environ.get( "AZURE_STORAGE_CONNECTION_STRING", "") AZURE_STORAGE_ACCOUNT_NAME = os.environ.get("AZURE_STORAGE_ACCOUNT_NAME", "") AZURE_STORAGE_ACCOUNT_KEY = os.environ.get("AZURE_STORAGE_ACCOUNT_KEY", "") # Container name - can be configured or use default AZURE_CONTAINER_NAME = os.environ.get("AZURE_CONTAINER_NAME", "invoice-splits") # ⭐ FOLDER STRUCTURE CONFIGURATION ROOT_FOLDER = os.environ.get("ROOT_FOLDER", "POD") # Root folder name # ⭐ PERFORMANCE CONFIGURATION MAX_PARALLEL_GEMINI_CALLS = int( os.environ.get("MAX_PARALLEL_GEMINI_CALLS", "5")) GEMINI_IMAGE_RESOLUTION = float( os.environ.get("GEMINI_IMAGE_RESOLUTION", "1.2")) USE_SMART_SAMPLING = os.environ.get( "USE_SMART_SAMPLING", "false").lower() == "true" # ⭐ SERVER CONFIGURATION HOST = os.environ.get("HOST", "0.0.0.0") # Hugging Face uses 0.0.0.0 PORT = int(os.environ.get("PORT", "7860")) # Hugging Face default port MAX_WAIT_TIME = 300 # 5 minutes max wait for quota model_lock = Lock() quota_manager_lock = Lock() # ============================================================================ # GLOBAL VARIABLES # ============================================================================ gemini_model = None blob_service_client = None GEMINI_REST_URL = "https://generativelanguage.googleapis.com/v1/models/{model}:generateContent?key={key}" def call_gemini_25(model_name: str, image_bytes: bytes, prompt: str) -> str: global current_model_index while True: model_config = get_current_model_config() url = GEMINI_REST_URL.format( model=model_config["name"], key=GEMINI_API_KEY) encoded = base64.b64encode(image_bytes).decode("utf-8") payload = { "contents": [ { "parts": [ {"inline_data": {"mime_type": "image/png", "data": encoded}}, {"text": prompt} ] } ], "generationConfig": {"temperature": 0} } r = requests.post(url, json=payload, timeout=model_config["timeout"]) # SUCCESS if r.status_code == 200: record_model_request(model_config) data = r.json() return data["candidates"][0]["content"]["parts"][0]["text"] # QUOTA HIT → SWITCH MODEL if r.status_code in (429, 503): print( f"⚠️ RPM exhausted for {model_config['name']} → switching model") model_config["current_rpm"] = model_config["max_requests_per_minute"] next_model = get_next_available_model() if next_model: print(f"🔄 Switched to {next_model['name']}") continue # All models exhausted → wait print("⏳ All models exhausted. Waiting 60s...") time.sleep(60) continue # Other error raise Exception(f"Gemini error {r.status_code}: {r.text}") def get_next_available_model(): global current_model_index for i in range(len(GEMINI_MODELS)): idx = (current_model_index + i) % len(GEMINI_MODELS) if can_use_model(GEMINI_MODELS[idx]): current_model_index = idx return GEMINI_MODELS[idx] return None # Model configuration with quota tracking GEMINI_MODELS = [ { "name": "gemini-2.5-flash-lite", "max_requests_per_minute": 120, "max_requests_per_day": 10000, "max_output_tokens": 16384, "timeout": 60, "description": "Stage 1 - Pre-classification / validation / cheap parsing", "current_rpm": 0, "current_rpd": 0, "last_rpm_reset": None, "last_rpd_reset": None, "quota_reset_time": None, "skip_on_error": True }, { "name": "gemini-2.5-flash-image", "max_requests_per_minute": 50, "max_requests_per_day": 1500, "max_output_tokens": 65536, "timeout": 300, "description": "Stage 2 - Primary invoice OCR extraction", "current_rpm": 0, "current_rpd": 0, "last_rpm_reset": None, "last_rpd_reset": None, "quota_reset_time": None, "skip_on_error": False }, { "name": "gemini-2.5-pro", "max_requests_per_minute": 10, "max_requests_per_day": 1000, "max_output_tokens": 65536, "timeout": 300, "description": "Stage 3 - Complex invoice reasoning fallback", "current_rpm": 0, "current_rpd": 0, "last_rpm_reset": None, "last_rpd_reset": None, "quota_reset_time": None, "skip_on_error": False } ] current_model_index = 0 # ============================================================================ # ⭐ QUOTA MANAGEMENT FUNCTIONS # ============================================================================ def reset_model_quota_counters(model_config): """Reset quota counters based on time windows""" now = datetime.now() with quota_manager_lock: if model_config["last_rpm_reset"] is None: model_config["last_rpm_reset"] = now model_config["current_rpm"] = 0 elif (now - model_config["last_rpm_reset"]).total_seconds() >= 60: model_config["current_rpm"] = 0 model_config["last_rpm_reset"] = now logger.debug(f"🔄 Reset RPM for {model_config['name']}") if model_config["last_rpd_reset"] is None: model_config["last_rpd_reset"] = now model_config["current_rpd"] = 0 elif now.date() > model_config["last_rpd_reset"].date(): model_config["current_rpd"] = 0 model_config["last_rpd_reset"] = now logger.info(f"🔄 Reset daily quota for {model_config['name']}") def can_use_model(model_config): """Check if model has available quota""" reset_model_quota_counters(model_config) with quota_manager_lock: rpm_ok = model_config["current_rpm"] < model_config["max_requests_per_minute"] rpd_ok = model_config["current_rpd"] < model_config["max_requests_per_day"] return rpm_ok and rpd_ok def record_model_request(model_config): """Record a request""" with quota_manager_lock: model_config["current_rpm"] += 1 model_config["current_rpd"] += 1 logger.debug( f"📊 {model_config['name']}: RPM={model_config['current_rpm']}/{model_config['max_requests_per_minute']}") def wait_for_quota_renewal(max_wait=MAX_WAIT_TIME): """Wait for any model to have quota""" start = time.time() logger.info( f"⏳ All models quota exhausted. Waiting for renewal (max {max_wait}s)...") while time.time() - start < max_wait: for i, model in enumerate(GEMINI_MODELS): if can_use_model(model): elapsed = time.time() - start logger.info( f"✅ {model['name']} quota available after {elapsed:.1f}s") return True, i elapsed = time.time() - start remaining = max_wait - elapsed logger.info( f"⏰ Waiting... (elapsed: {elapsed:.0f}s, remaining: {remaining:.0f}s)") time.sleep(10) logger.error(f"❌ Timeout: No quota available after {max_wait}s") return False, -1 def get_current_model_config(): """Get current model config""" return GEMINI_MODELS[current_model_index] # ============================================================================ # STARTUP VALIDATION # ============================================================================ def validate_configuration(): """Validate configuration and warn about missing credentials.""" warnings = [] errors = [] if not GEMINI_API_KEY: warnings.append( "⚠️ GEMINI_API_KEY not set - image-based PDFs will not work") else: print(f"✅ GEMINI_API_KEY configured ({len(GEMINI_API_KEY)} chars)") if not AZURE_STORAGE_CONNECTION_STRING: if not (AZURE_STORAGE_ACCOUNT_NAME and AZURE_STORAGE_ACCOUNT_KEY): errors.append("❌ Azure credentials missing") else: print( f"✅ Azure credentials configured (account: {AZURE_STORAGE_ACCOUNT_NAME})") else: print(f"✅ Azure connection string configured") for warning in warnings: print(warning) for error in errors: print(error) if errors: print("\n⚠️ WARNING: Some required credentials are missing!") return len(errors) == 0 # ============================================================================ # AZURE BLOB STORAGE FUNCTIONS # ============================================================================ def get_blob_service_client(): """Get or create Azure Blob Service Client.""" global blob_service_client if not AZURE_AVAILABLE: print("❌ Azure SDK not available") return None if blob_service_client is None: try: if AZURE_STORAGE_CONNECTION_STRING: blob_service_client = BlobServiceClient.from_connection_string( AZURE_STORAGE_CONNECTION_STRING) print("✅ Azure Blob Storage initialized with connection string") elif AZURE_STORAGE_ACCOUNT_NAME and AZURE_STORAGE_ACCOUNT_KEY: account_url = f"https://{AZURE_STORAGE_ACCOUNT_NAME}.blob.core.windows.net" blob_service_client = BlobServiceClient( account_url=account_url, credential=AZURE_STORAGE_ACCOUNT_KEY) print("✅ Azure Blob Storage initialized with account key") else: print("⚠️ WARNING: No Azure credentials configured") return None except Exception as e: print(f"❌ Failed to initialize Azure Blob Storage: {e}") return None return blob_service_client def ensure_container_exists(container_name: str = None): """Create container if it doesn't exist.""" if container_name is None: container_name = AZURE_CONTAINER_NAME try: client = get_blob_service_client() if client: container_client = client.get_container_client(container_name) if not container_client.exists(): container_client.create_container() print(f"✅ Created container: {container_name}") else: print(f"✅ Container exists: {container_name}") except Exception as e: print(f"⚠️ Container check error: {e}") def upload_raw_pdf_to_blob(pdf_bytes: bytes, filename: str, batch_id: str, container_name: str = None) -> dict: """Upload original/raw PDF to Azure Blob Storage.""" if container_name is None: container_name = AZURE_CONTAINER_NAME try: client = get_blob_service_client() if not client: raise HTTPException( status_code=500, detail="Azure Blob Storage not configured") base_filename = os.path.splitext(filename)[0] safe_folder_name = re.sub(r'[<>:"/\\|?*]', '_', base_filename) blob_name = f"{ROOT_FOLDER}/{batch_id}/{safe_folder_name}/Raw/{filename}" blob_client = client.get_blob_client( container=container_name, blob=blob_name) print(f"📤 Uploading raw PDF to: {blob_name}") blob_client.upload_blob( pdf_bytes, overwrite=True, content_settings=ContentSettings(content_type='application/pdf'), metadata={ 'batch_id': batch_id, 'file_type': 'raw', 'uploaded_at': datetime.now().isoformat(), 'original_filename': filename } ) expiry_hours = 24 sas_token = generate_blob_sas( account_name=AZURE_STORAGE_ACCOUNT_NAME, container_name=container_name, blob_name=blob_name, account_key=AZURE_STORAGE_ACCOUNT_KEY, permission=BlobSasPermissions(read=True), expiry=datetime.utcnow() + timedelta(hours=expiry_hours) ) blob_url = blob_client.url download_url = f"{blob_url}?{sas_token}" expires_at = (datetime.utcnow() + timedelta(hours=expiry_hours)).isoformat() + "Z" print(f"✅ Uploaded raw PDF: {blob_name}") return { "blob_name": blob_name, "blob_url": blob_url, "download_url": download_url, "expires_at": expires_at, "expires_in_hours": expiry_hours, "storage": "azure_blob", "folder_type": "raw", "container": container_name, "size_bytes": len(pdf_bytes), "size_mb": round(len(pdf_bytes) / (1024 * 1024), 2) } except Exception as e: print(f"❌ Raw PDF upload failed: {e}") raise HTTPException( status_code=500, detail=f"Azure Blob upload failed: {str(e)}") def upload_split_pdf_to_blob(pdf_bytes: bytes, invoice_filename: str, original_filename: str, batch_id: str, container_name: str = None) -> dict: """Upload split invoice PDF to Azure Blob Storage.""" if container_name is None: container_name = AZURE_CONTAINER_NAME try: client = get_blob_service_client() if not client: raise HTTPException( status_code=500, detail="Azure Blob Storage not configured") base_filename = os.path.splitext(original_filename)[0] safe_folder_name = re.sub(r'[<>:"/\\|?*]', '_', base_filename) blob_name = f"{ROOT_FOLDER}/{batch_id}/{safe_folder_name}/Splitted/{invoice_filename}" blob_client = client.get_blob_client( container=container_name, blob=blob_name) blob_client.upload_blob( pdf_bytes, overwrite=True, content_settings=ContentSettings(content_type='application/pdf'), metadata={ 'batch_id': batch_id, 'file_type': 'split', 'uploaded_at': datetime.now().isoformat(), 'original_filename': original_filename, 'invoice_filename': invoice_filename } ) expiry_hours = 24 sas_token = generate_blob_sas( account_name=AZURE_STORAGE_ACCOUNT_NAME, container_name=container_name, blob_name=blob_name, account_key=AZURE_STORAGE_ACCOUNT_KEY, permission=BlobSasPermissions(read=True), expiry=datetime.utcnow() + timedelta(hours=expiry_hours) ) blob_url = blob_client.url download_url = f"{blob_url}?{sas_token}" expires_at = (datetime.utcnow() + timedelta(hours=expiry_hours)).isoformat() + "Z" return { "blob_name": blob_name, "blob_url": blob_url, "download_url": download_url, "expires_at": expires_at, "expires_in_hours": expiry_hours, "storage": "azure_blob", "folder_type": "split", "container": container_name, "size_bytes": len(pdf_bytes), "size_mb": round(len(pdf_bytes) / (1024 * 1024), 2) } except Exception as e: print(f"❌ Split PDF upload failed: {e}") raise HTTPException( status_code=500, detail=f"Azure Blob upload failed: {str(e)}") async def cleanup_old_blobs(batch_id: str, container_name: str = None): """Delete all blobs for a specific batch_id.""" if container_name is None: container_name = AZURE_CONTAINER_NAME try: client = get_blob_service_client() if not client: return container_client = client.get_container_client(container_name) prefix = f"{ROOT_FOLDER}/{batch_id}/" blobs = container_client.list_blobs(name_starts_with=prefix) deleted_count = 0 for blob in blobs: blob_client = container_client.get_blob_client(blob.name) blob_client.delete_blob() deleted_count += 1 print(f"🧹 Cleaned up {deleted_count} blobs for batch {batch_id}") except Exception as e: print(f"⚠️ Cleanup error: {e}") # ============================================================================ # OPTIMIZED GEMINI FUNCTIONS WITH ASYNC PROCESSING # ============================================================================ def get_gemini_model(): """Get or create Gemini model instance WITH QUOTA CHECK""" global gemini_model, current_model_index if not GEMINI_AVAILABLE: return None if not GEMINI_API_KEY: return None # ⭐ CHECK QUOTA BEFORE RETURNING MODEL model_config = get_current_model_config() if not can_use_model(model_config): logger.warning(f"⚠️ {model_config['name']} quota exhausted") # Try other models for i, alt_model in enumerate(GEMINI_MODELS): if i != current_model_index and can_use_model(alt_model): current_model_index = i model_config = alt_model logger.info(f"🔄 Switched to {model_config['name']}") gemini_model = None # Force recreation break else: # All models exhausted - wait success, new_index = wait_for_quota_renewal(MAX_WAIT_TIME) if success: current_model_index = new_index model_config = GEMINI_MODELS[new_index] gemini_model = None else: logger.error("❌ All models quota exhausted") return None # Create/recreate model if needed with model_lock: if gemini_model is None or not hasattr(gemini_model, '_model_name') or gemini_model._model_name != model_config['name']: try: genai.configure(api_key=GEMINI_API_KEY) gemini_model = genai.GenerativeModel(model_config['name']) gemini_model._model_name = model_config['name'] logger.info(f"✅ Using {model_config['name']}") except Exception as e: logger.error( f"❌ Failed to initialize {model_config['name']}: {e}") return None # ❌ REMOVE THIS LINE - Don't record request for model creation # record_model_request(model_config) # <-- DELETE THIS return gemini_model def extract_invoice_gemini_sync(page): model_config = get_current_model_config() try: pix = page.get_pixmap(matrix=fitz.Matrix( GEMINI_IMAGE_RESOLUTION, GEMINI_IMAGE_RESOLUTION)) img_bytes = pix.tobytes("png") pix = None prompt = """Look at this invoice image and extract ONLY the invoice number. Return ONLY the invoice number. If not found return NONE.""" text = call_gemini_25(model_config["name"], img_bytes, prompt) cleaned = text.strip().replace("Invoice Number:", "").strip() print(f"🤖 Gemini raw response: {text}") if cleaned and cleaned.upper() != "NONE" and len(cleaned) >= 3: cleaned = re.sub(r'[^A-Za-z0-9\-/]', '', cleaned) print(f"✅ Gemini extracted: {cleaned}") return cleaned.upper() # Fallback OCR ocr = call_gemini_25( model_config["name"], img_bytes, "Extract all visible text from this image") return try_extract_invoice_from_text(ocr) except Exception as e: print(f"❌ Gemini error: {e}") return None async def extract_invoices_batch_async(doc: fitz.Document, is_image_pdf: bool, batch_size: int = MAX_PARALLEL_GEMINI_CALLS) -> List[Optional[str]]: """Extract invoice numbers with parallel processing.""" page_invoice_nos = [] if not is_image_pdf: print(f" 📝 Text-based extraction (sequential)") for i in range(doc.page_count): if i % 50 == 0: print(f" Extracting... Page {i+1}/{doc.page_count}") page = doc.load_page(i) inv = extract_invoice_text_based(page) page_invoice_nos.append(inv) page = None if i % 100 == 0: gc.collect() return page_invoice_nos print(f" 🚀 Image-based extraction (parallel, batch_size={batch_size})") with ThreadPoolExecutor(max_workers=batch_size) as executor: futures = [] for i in range(doc.page_count): page = doc.load_page(i) text_result = extract_invoice_text_based(page) if text_result: futures.append((i, None, text_result)) else: future = executor.submit(extract_invoice_gemini_sync, page) futures.append((i, future, None)) page_invoice_nos = [None] * doc.page_count completed = 0 for i, future, text_result in futures: try: if text_result: page_invoice_nos[i] = text_result completed += 1 else: result = future.result(timeout=30) page_invoice_nos[i] = result completed += 1 if completed % 5 == 0: print( f" ✓ Processed {completed}/{doc.page_count} pages...") except Exception as e: print(f" ⚠️ Page {i+1} failed: {e}") page_invoice_nos[i] = None if completed % 20 == 0: gc.collect() print(f" ✅ Extraction complete: {completed}/{doc.page_count} pages") return page_invoice_nos def extract_invoices_smart_sampling(doc: fitz.Document, is_image_pdf: bool) -> List[Optional[str]]: """Smart sampling strategy for large PDFs.""" print(f" ⚡ Smart sampling mode (faster, ~95% accurate)") page_invoice_nos = [None] * doc.page_count page = doc.load_page(0) page_invoice_nos[0] = extract_invoice_no_from_page(page, is_image_pdf) print(f" ✓ Page 1: {page_invoice_nos[0]}") sample_interval = max(3, doc.page_count // 20) print(f" Sampling interval: every {sample_interval} pages") for i in range(sample_interval, doc.page_count, sample_interval): page = doc.load_page(i) inv = extract_invoice_no_from_page(page, is_image_pdf) page_invoice_nos[i] = inv if i % 10 == 0: print(f" Sampling page {i+1}/{doc.page_count}...") prev_known_idx = i - sample_interval while prev_known_idx >= 0 and page_invoice_nos[prev_known_idx] is None: prev_known_idx -= 1 if prev_known_idx >= 0 and inv != page_invoice_nos[prev_known_idx]: print(f" 🔍 Boundary detected near page {i+1}, refining...") for offset in range(-3, 4): idx = i + offset if 0 <= idx < doc.page_count and page_invoice_nos[idx] is None: page = doc.load_page(idx) page_invoice_nos[idx] = extract_invoice_no_from_page( page, is_image_pdf) if page_invoice_nos[-1] is None: page = doc.load_page(doc.page_count - 1) page_invoice_nos[-1] = extract_invoice_no_from_page(page, is_image_pdf) print(f" ✓ Last page: {page_invoice_nos[-1]}") last_known = page_invoice_nos[0] filled = 0 for i in range(len(page_invoice_nos)): if page_invoice_nos[i] is not None: last_known = page_invoice_nos[i] else: page_invoice_nos[i] = last_known filled += 1 print(f" ✅ Smart sampling complete: forward-filled {filled} pages") return page_invoice_nos # ============================================================================ # PDF PROCESSING FUNCTIONS # ============================================================================ def is_image_based_pdf(doc: fitz.Document, sample_pages: int = 3) -> Tuple[bool, float]: total_text_length = 0 pages_to_check = min(sample_pages, doc.page_count) for i in range(pages_to_check): text = doc.load_page(i).get_text("text") or "" total_text_length += len(text.strip()) avg_text_length = total_text_length / pages_to_check is_image_based = avg_text_length < 50 print(f" PDF Type: {'IMAGE-BASED' if is_image_based else 'TEXT-BASED'}") print(f" Avg text per page: {avg_text_length:.1f} chars") return is_image_based, avg_text_length def normalize_text_for_search(s: str) -> str: if not s: return s s = s.replace("\u00A0", " ") s = re.sub(r"[\r\n\t]+", " ", s) s = re.sub(r"[ ]{2,}", " ", s).strip() return s def is_valid_invoice_number(candidate: str) -> bool: if not candidate or len(candidate) < 3: return False if len(candidate) == 15 and re.match(r'^[0-9A-Z]{15}$', candidate.upper()): return False if re.match(r'^\d+$', candidate): return 6 <= len(candidate) <= 15 if re.match(r'^\d+\.\d{2,}$', candidate): return False has_letter = any(c.isalpha() for c in candidate) has_digit = any(c.isdigit() for c in candidate) return has_letter and has_digit def try_extract_invoice_from_text(text: str) -> Optional[str]: """Universal label-first extraction with smart prioritization""" if not text: return None text_norm = normalize_text_for_search(text) if len(text_norm) > 0: print(f"\n{'='*70}") print(f"🔍 ANALYZING TEXT (first 800 chars):") print(f"{text_norm[:800]}") print(f"{'='*70}\n") label_patterns = [ (r"Invoice\s*(?:No\.?|Number|Num)", "Invoice No", True), (r"Inv\s*(?:No\.?|Number)", "Inv No", True), (r"Bill\s*(?:No\.?|Number|Num)", "Bill No", True), (r"Tax\s*Invoice\s*(?:No\.?|Number)", "Tax Invoice No", True), (r"Document\s*(?:No\.?|Number)", "Document No", True), (r"Receipt\s*(?:No\.?|Number)", "Receipt No", False), (r"Voucher\s*(?:No\.?|Number)", "Voucher No", False), (r"Reference\s*(?:No\.?|Number)", "Reference No", False), (r"Ref\s*(?:No\.?|Number)", "Ref No", False), (r"Credit\s*(?:Note\s*)?(?:No\.?|Number)", "Credit No", False), (r"Debit\s*(?:Note\s*)?(?:No\.?|Number)", "Debit No", False), ] for label_pattern, label_name, is_invoice_label in label_patterns: header_text = text_norm[:2000] label_matches = list(re.finditer( label_pattern, header_text, re.IGNORECASE)) for label_match in label_matches: start_pos = label_match.end() text_after_label = header_text[start_pos:start_pos + 200] print( f"🔍 Found label '{label_name}' at position {label_match.start()}") print( f" Text after label (first 80 chars): '{text_after_label[:80]}...'") all_candidates = re.findall( r'\b([A-Z0-9][A-Z0-9\-\/]{2,20})\b', text_after_label, re.IGNORECASE) print( f" Found {len(all_candidates)} potential candidates: {all_candidates[:5]}") for pass_number in [1, 2]: if pass_number == 2 and len(all_candidates) > 0: print(f" 🔄 Second pass: Trying alphanumeric candidates...") for candidate in all_candidates: invoice_num = candidate.strip(".,;:-_") if len(invoice_num) < 3: continue is_pure_numeric = invoice_num.isdigit() is_ideal_invoice_length = 12 <= len(invoice_num) <= 14 if pass_number == 1: if not (is_pure_numeric and is_ideal_invoice_length): continue print( f" ✨ PRIORITY candidate (12-14 digit numeric): '{invoice_num}'") else: if is_pure_numeric and is_ideal_invoice_length: continue print(f" 🔍 Evaluating candidate: '{invoice_num}'") if invoice_num.upper() in ("ORDER", "REF", "NO", "NUMBER", "DATE", "DT", "AND", "INV", "BILL", "ACCOUNT", "PO", "COPY", "OF", "DOCUMENT", "DOC", "GST", "GSTIN", "ACK", "USER", "DIRECT", "TAB", "MG", "GM", "ML", "CREDIT"): print(f" ⚠️ Skipped: noise word") continue if not is_invoice_label: if re.match(r'^[A-Z]\d{6}$', invoice_num, re.IGNORECASE): print( f" ⚠️ Skipped: batch pattern (non-invoice context)") continue if re.match(r'^[A-Z]{2,3}-[A-Z0-9]+-\d+$', invoice_num, re.IGNORECASE): print(f" ⚠️ Skipped: license pattern") continue if re.match(r'^[A-Z]{2,4}-\d{10}$', invoice_num, re.IGNORECASE): print(f" ⚠️ Skipped: state code/UIN pattern") continue if re.search(rf"Ack\.?\s*(?:No\.?|Number)?\s*:?\s*{re.escape(invoice_num)}", text_norm, re.IGNORECASE): print(f" ⚠️ Skipped: ACK number") continue if re.search(rf"PH\.?\s*(?:No\.?|Number)?\s*:?\s*{re.escape(invoice_num)}", text_norm, re.IGNORECASE): print(f" ⚠️ Skipped: PH number") continue if re.search(rf"(?:UIN|UID|State\s*Code|D\.L\.No)\.?\s*:?\s*{re.escape(invoice_num)}", text_norm, re.IGNORECASE): print(f" ⚠️ Skipped: UIN/UID/State Code") continue if re.search(rf"A[\s\/]*C\s*(?:No\.?|Number)?\s*[\-:\.]?\s*{re.escape(invoice_num)}", text_norm, re.IGNORECASE): print(f" ⚠️ Skipped: A/C number") continue if re.match(r'^[0-9]{10,11}$', invoice_num): if len(invoice_num) == 10 and invoice_num[0] in '6789': print(f" ⚠️ Skipped: mobile number") continue if len(invoice_num) == 11 and invoice_num[0] == '0': print(f" ⚠️ Skipped: landline number") continue if re.match(r'^20\d{6}$', invoice_num): print(f" ⚠️ Skipped: date pattern") continue if re.match(r'^\d{2}[\/\-]\d{2}[\/\-]\d{4}$', invoice_num): print(f" ⚠️ Skipped: date format") continue if len(invoice_num) == 15 and re.match(r'^\d{2}[A-Z]{5}\d{4}[A-Z]\d[A-Z]\d$', invoice_num, re.IGNORECASE): print(f" ⚠️ Skipped: GST number") continue print(f" ✅✅✅ ACCEPTED: '{invoice_num}'") return invoice_num.upper() print(f" ⚠️ No valid candidates found after '{label_name}'") print("\n⚠️ No labeled invoice number found, trying fallback extraction...") top_text = text_norm[:1000] credit_match = re.search( r"CREDIT\s*(?:NO|NUMBER|#)?\s*[:\-]?\s*(\d{12,20})", text_norm, re.IGNORECASE) if credit_match: credit_num = credit_match.group(1).strip() if 12 <= len(credit_num) <= 20 and len(credit_num) != 14: print(f"✓ Fallback: Found CREDIT number: {credit_num}") return credit_num.upper() long_numerics = re.findall(r'\b(\d{12,20})\b', top_text) for num in long_numerics: if re.search(rf"(?:Ack|PH|A[\s\/]*C|UIN|UID|State\s*Code|D\.L\.No)\.?\s*(?:No\.?|Number)?\s*:?\s*{re.escape(num)}", text_norm, re.IGNORECASE): print(f"⚠️ Fallback: Skipping (labeled as ACK/PH/A/C/UIN): {num}") continue print(f"✓ Fallback: Found long numeric: {num}") return num.upper() medium_numerics = re.findall(r'\b(\d{10,15})\b', top_text) for num in medium_numerics: if len(num) == 10 and num[0] in '6789': continue if len(num) == 11 and num[0] == '0': continue if len(num) == 8 and num.startswith('20'): continue if re.search(rf"(?:Ack|PH|A[\s\/]*C|UIN|UID)\.?\s*(?:No\.?|Number)?\s*:?\s*{re.escape(num)}", text_norm, re.IGNORECASE): continue print(f"✓ Fallback: Found medium numeric: {num}") return num.upper() print("✗ No invoice number found (labeled or unlabeled)") return None def extract_invoice_text_based(page: fitz.Page) -> Optional[str]: """Extract invoice number from TEXT-BASED PDF with Zydus fallback""" text = page.get_text("text") or "" text_norm = normalize_text_for_search(text) header_text = text_norm[:2500] zydus_candidates = re.findall(r'\b(23\d{8})\b', header_text) if zydus_candidates: zydus_number = zydus_candidates[0] print(f" ✅ ZYDUS INVOICE DETECTED: {zydus_number}") return zydus_number.upper() inv = try_extract_invoice_from_text(text) if inv: if re.match(r'^10\d{12}$', inv): print(f" ⚠️ REJECTED Order ID (14-digit): {inv}") inv = None else: return inv for block in (page.get_text("blocks") or []): block_text = block[4] if len(block) > 4 else "" if block_text: inv = try_extract_invoice_from_text(block_text) if inv: if re.match(r'^10\d{12}$', inv): print(f" ⚠️ REJECTED Order ID from block: {inv}") continue else: return inv blocks = page.get_text("blocks") or [] sorted_blocks = sorted(blocks, key=lambda b: b[1] if len(b) > 1 else 0) for block in sorted_blocks[:15]: block_text = block[4] if len(block) > 4 else "" if block_text: block_norm = normalize_text_for_search(block_text) block_candidates = re.findall(r'\b(23\d{8})\b', block_norm) if block_candidates: number = block_candidates[0] print(f" ✅ ZYDUS BLOCK DETECTION: {number}") return number.upper() print(f" ⚠️ No valid invoice found on this page (will use forward-fill)") return None def extract_invoice_no_from_page(page: fitz.Page, is_image_pdf: bool) -> Optional[str]: """Extract invoice number from a single page""" text_result = extract_invoice_text_based(page) if text_result: return text_result if is_image_pdf: return extract_invoice_gemini_sync(page) return None def build_pdf_from_pages(src_doc: fitz.Document, page_indices: List[int]) -> bytes: out = fitz.open() try: for i in page_indices: out.insert_pdf(src_doc, from_page=i, to_page=i) pdf_bytes = out.tobytes(garbage=4, deflate=True) return pdf_bytes finally: out.close() def remove_file(path: str): try: if os.path.exists(path): os.remove(path) except Exception as e: print(f"⚠️ Cleanup warning: {e}") def merge_first_null_group(groups: List[Dict]) -> List[Dict]: """Merge null first page with first invoice""" if len(groups) >= 2: first_group = groups[0] second_group = groups[1] if first_group["invoice_no"] is None and second_group["invoice_no"] is not None: print(f"\n🔧 AUTO-FIX: Merging null first page(s) with first invoice") print( f" Null group: Pages {[p+1 for p in first_group['pages']]}") print( f" First invoice: {second_group['invoice_no']}, Pages {[p+1 for p in second_group['pages']]}") merged_pages = first_group["pages"] + second_group["pages"] second_group["pages"] = merged_pages groups.pop(0) print( f" ✅ Merged result: {second_group['invoice_no']}, Pages {[p+1 for p in merged_pages]}\n") return groups # ============================================================================ # API ENDPOINTS # ============================================================================ @app.post("/split-invoices") async def split_invoices( background_tasks: BackgroundTasks, file: UploadFile = File(...), batch_id: str = Form(..., description="Batch ID (required)"), use_blob_storage: bool = Form( True, description="Upload PDFs to Azure Blob Storage"), blob_container: Optional[str] = Form( None, description="Custom Azure container"), include_base64: bool = Form( False, description="Include base64 in response"), parallel_batch_size: int = Form( MAX_PARALLEL_GEMINI_CALLS, description="Parallel Gemini API calls"), use_smart_sampling: bool = Form( USE_SMART_SAMPLING, description="Use smart sampling"), max_file_size_mb: int = Form(200, description="Maximum file size in MB"), ): """Universal Invoice Splitter with RPM Management""" if not file.filename: raise HTTPException(status_code=400, detail="No filename provided") filename_lower = file.filename.lower() SUPPORTED_EXTENSIONS = ['.pdf', '.png', '.jpg', '.jpeg', '.tiff', '.tif', '.bmp'] file_extension = None for ext in SUPPORTED_EXTENSIONS: if filename_lower.endswith(ext): file_extension = ext break if not file_extension: raise HTTPException( status_code=400, detail=f"Unsupported file format. Supported: {', '.join(SUPPORTED_EXTENSIONS)}") is_image_file = file_extension in [ '.png', '.jpg', '.jpeg', '.tiff', '.tif', '.bmp'] if is_image_file and not GEMINI_AVAILABLE: raise HTTPException( status_code=500, detail="Image processing requires PIL") if use_blob_storage and not get_blob_service_client(): raise HTTPException( status_code=500, detail="Azure Blob Storage not configured") container_name = blob_container if blob_container else AZURE_CONTAINER_NAME if use_blob_storage: ensure_container_exists(container_name) max_size_bytes = max_file_size_mb * 1024 * 1024 fd, temp_path = tempfile.mkstemp(suffix=file_extension) os.close(fd) doc = None original_pdf_bytes = None start_time = datetime.now() pdf_path = temp_path original_filename = file.filename try: print(f"\n{'='*70}") print(f"📥 Processing: {file.filename}") print(f" File Type: {'IMAGE' if is_image_file else 'PDF'}") print(f" Batch ID: {batch_id}") print( f" Performance Mode: {'Smart Sampling' if use_smart_sampling else f'Parallel ({parallel_batch_size} workers)'}") print(f"{'='*70}") total_size = 0 with open(temp_path, "wb") as buffer: chunk_read_size = 5 * 1024 * 1024 while content := await file.read(chunk_read_size): total_size += len(content) if total_size > max_size_bytes: remove_file(temp_path) raise HTTPException( status_code=413, detail=f"File too large. Max: {max_file_size_mb}MB") buffer.write(content) file_size_mb = total_size / (1024 * 1024) print(f"💾 File size: {file_size_mb:.2f}MB") if is_image_file: print(f"🖼️ Converting image to PDF...") try: from PIL import Image as PILImage img = PILImage.open(temp_path) if img.mode != 'RGB': img = img.convert('RGB') pdf_path = temp_path.replace(file_extension, '.pdf') img.save(pdf_path, 'PDF', resolution=100.0) img.close() print(f"✅ Image converted to PDF") file.filename = file.filename.replace(file_extension, '.pdf') except Exception as e: print(f"❌ Image conversion failed: {e}") raise HTTPException( status_code=500, detail=f"Failed to convert image to PDF: {str(e)}") with open(pdf_path, "rb") as f: original_pdf_bytes = f.read() raw_pdf_info = None if use_blob_storage: try: print( f"\n📤 Uploading original {'PDF' if not is_image_file else 'converted PDF'} to Raw folder...") raw_pdf_info = upload_raw_pdf_to_blob( original_pdf_bytes, file.filename, batch_id, container_name) print(f"✅ Original PDF uploaded: {raw_pdf_info['blob_name']}") except Exception as e: print(f"⚠️ Failed to upload raw PDF: {e}") doc = fitz.open(pdf_path) if doc.page_count == 0: raise HTTPException(status_code=400, detail="Empty PDF") print(f"📄 Total pages: {doc.page_count}") is_image_pdf, _ = is_image_based_pdf(doc) if is_image_pdf and not get_gemini_model(): raise HTTPException( status_code=500, detail="Image PDF detected but Gemini not configured") print(f"\n📊 Extracting invoice numbers...") extraction_start = datetime.now() if use_smart_sampling and doc.page_count > 10: page_invoice_nos = extract_invoices_smart_sampling( doc, is_image_pdf) else: page_invoice_nos = await extract_invoices_batch_async(doc, is_image_pdf, batch_size=parallel_batch_size) extraction_time = (datetime.now() - extraction_start).total_seconds() print(f"✅ Extraction completed in {extraction_time:.1f} seconds") print(f" Speed: {doc.page_count / extraction_time:.1f} pages/second") print(f"\n🔧 Grouping invoices...") print(f"\n🔍 DEBUG - Raw extraction results:") for idx, inv in enumerate(page_invoice_nos[:min(10, len(page_invoice_nos))]): print(f" Page {idx+1}: {inv if inv else '(not found)'}") if len(page_invoice_nos) > 10: print( f" ... (showing first 10 of {len(page_invoice_nos)} pages)") page_invoice_nos_normalized = [] for v in page_invoice_nos: if v and v.upper().startswith("GST"): page_invoice_nos_normalized.append(None) elif v: normalized = v.upper().strip().replace(" ", "").replace("_", "") page_invoice_nos_normalized.append(normalized) else: page_invoice_nos_normalized.append(None) page_invoice_nos_filled = [] last_known_invoice = None for idx, inv in enumerate(page_invoice_nos_normalized): if inv is not None: last_known_invoice = inv page_invoice_nos_filled.append(inv) else: page_invoice_nos_filled.append(last_known_invoice) filled_count = sum(1 for i in range(len(page_invoice_nos_normalized)) if page_invoice_nos_normalized[i] is None and page_invoice_nos_filled[i] is not None) unique_invoices = set( [v for v in page_invoice_nos_filled if v is not None]) print(f"\n 📊 Found {len(unique_invoices)} unique invoice numbers:") for inv_no in sorted(unique_invoices) if unique_invoices else []: page_count = sum(1 for v in page_invoice_nos_filled if v == inv_no) print(f" • {inv_no}: {page_count} pages") groups = [] current_group = [] current_invoice = None for idx, inv in enumerate(page_invoice_nos_filled): if idx == 0: current_invoice = inv current_group = [idx] else: if inv != current_invoice: groups.append({"invoice_no": current_invoice, "pages": current_group[:]}) print( f" 📄 Group {len(groups)}: Invoice {current_invoice or 'UNKNOWN'} - Pages {current_group[0]+1}-{current_group[-1]+1} ({len(current_group)} pages)") current_invoice = inv current_group = [idx] else: current_group.append(idx) if current_group: groups.append({"invoice_no": current_invoice, "pages": current_group[:]}) print( f" 📄 Group {len(groups)}: Invoice {current_invoice or 'UNKNOWN'} - Pages {current_group[0]+1}-{current_group[-1]+1} ({len(current_group)} pages)") if len(groups) == 1 and groups[0]["invoice_no"] is None: groups = [{"invoice_no": None, "pages": list(range(doc.page_count))}] groups = merge_first_null_group(groups) print(f"\n✅ Created {len(groups)} invoice groups (after auto-merge)") print( f" Forward-filled {filled_count} pages with missing invoice numbers") print(f"\n🔨 Building and uploading split invoices...") all_parts = [] for idx, g in enumerate(groups): if (idx + 1) % 20 == 0: print(f" Processing {idx + 1}/{len(groups)} invoices...") part_bytes = build_pdf_from_pages(doc, g["pages"]) invoice_no = g["invoice_no"] if g["invoice_no"] else f"NO_NUMBER_{idx + 1}" safe_invoice_no = re.sub(r'[<>:"/\\|?*]', '_', invoice_no) invoice_filename = f"invoice_{safe_invoice_no}.pdf" invoice_info = { "invoice_no": g["invoice_no"], "pages": [p + 1 for p in g["pages"]], "page_range": f"{g['pages'][0]+1}-{g['pages'][-1]+1}" if len(g['pages']) > 1 else f"{g['pages'][0]+1}", "num_pages": len(g["pages"]), "size_bytes": len(part_bytes), "size_mb": round(len(part_bytes) / (1024 * 1024), 2), } if use_blob_storage: try: blob_info = upload_split_pdf_to_blob( part_bytes, invoice_filename, file.filename, batch_id, container_name) invoice_info["storage"] = blob_info invoice_info["pdf_url"] = blob_info["download_url"] invoice_info["blob_name"] = blob_info["blob_name"] invoice_info["expires_at"] = blob_info["expires_at"] except Exception as e: print(f" ⚠️ Failed to upload invoice {idx+1}: {e}") invoice_info["upload_error"] = str(e) if include_base64: invoice_info["pdf_base64"] = base64.b64encode( part_bytes).decode("ascii") all_parts.append(invoice_info) del part_bytes if idx % 50 == 0: gc.collect() print(f"✅ Processed all {len(all_parts)} invoices") total_pages_count = doc.page_count doc.close() doc = None remove_file(temp_path) if pdf_path != temp_path: remove_file(pdf_path) gc.collect() total_time = (datetime.now() - start_time).total_seconds() response_data = { "success": True, "batch_id": batch_id, "folder_structure": { "root": ROOT_FOLDER, "path": f"{ROOT_FOLDER}/{batch_id}/{os.path.splitext(file.filename)[0]}", "raw_folder": f"{ROOT_FOLDER}/{batch_id}/{os.path.splitext(file.filename)[0]}/Raw", "split_folder": f"{ROOT_FOLDER}/{batch_id}/{os.path.splitext(file.filename)[0]}/Splitted" }, "source_file": { "name": file.filename, "original_name": original_filename, "size_mb": round(file_size_mb, 2), "total_pages": total_pages_count, "pdf_type": "image-based" if is_image_pdf else "text-based", "was_converted": is_image_file, "raw_pdf": raw_pdf_info }, "summary": { "total_invoices": len(all_parts), "unique_invoice_numbers": len(unique_invoices), "extraction_method": "gemini" if is_image_pdf else "text", "pages_forward_filled": filled_count, "storage_type": "azure_blob" if use_blob_storage else "base64", "model_used": get_current_model_config()['name'] }, "performance": { "total_time_seconds": round(total_time, 2), "extraction_time_seconds": round(extraction_time, 2), "pages_per_second": round(total_pages_count / extraction_time, 2) if extraction_time > 0 else 0, "parallel_batch_size": parallel_batch_size, "smart_sampling_used": use_smart_sampling and total_pages_count > 10 }, "invoices": all_parts } print(f"\n{'='*70}") print(f"✅ SUCCESS!") print(f" Batch ID: {batch_id}") print(f" Original File: {original_filename}") if is_image_file: print(f" ✓ Image converted to PDF") print( f" Raw PDF: {raw_pdf_info['blob_name'] if raw_pdf_info else 'Not uploaded'}") print(f" Split invoices: {len(all_parts)}") print(f" Unique invoice numbers: {len(unique_invoices)}") print(f" Model used: {get_current_model_config()['name']}") print(f" Total time: {total_time:.1f}s") print( f" Extraction time: {extraction_time:.1f}s ({total_pages_count / extraction_time:.1f} pages/sec)") print(f"{'='*70}\n") return JSONResponse(response_data) except HTTPException: raise except Exception as e: print(f"\n✗ Error: {str(e)}") import traceback traceback.print_exc() raise HTTPException(status_code=500, detail=str(e)) finally: if doc: doc.close() remove_file(temp_path) if pdf_path != temp_path: remove_file(pdf_path) gc.collect() @app.post("/cleanup-batch/{batch_id}") async def cleanup_batch( batch_id: str, background_tasks: BackgroundTasks, container_name: Optional[str] = Form(None) ): """Delete all blobs for a specific batch""" if container_name is None: container_name = AZURE_CONTAINER_NAME background_tasks.add_task(cleanup_old_blobs, batch_id, container_name) return JSONResponse({ "success": True, "message": f"Cleanup started for batch {batch_id}", "batch_id": batch_id, "folder_path": f"{ROOT_FOLDER}/{batch_id}/", "container": container_name }) @app.get("/quota-status") def quota_status(): """Get quota status for all models""" status = [] for i, model in enumerate(GEMINI_MODELS): reset_model_quota_counters(model) with quota_manager_lock: status.append({ "model": model["name"], "is_current": i == current_model_index, "rpm": {"used": model["current_rpm"], "limit": model["max_requests_per_minute"]}, "rpd": {"used": model["current_rpd"], "limit": model["max_requests_per_day"]}, "available": can_use_model(model) }) return JSONResponse({"models": status, "timestamp": datetime.now().isoformat()}) @app.get("/") async def root(): return { "service": "Universal Invoice Splitter API with RPM Management", "version": "4.0", "status": "running", "features": { "multi_format_support": True, "zydus_healthcare_support": True, "auto_merge_null_groups": True, "azure_blob_storage": True, "parallel_processing": True, "rpm_management": True, "multi_model_fallback": True }, "models": [m["name"] for m in GEMINI_MODELS], "current_model": get_current_model_config()['name'] } @app.get("/health") async def health(): model_config = get_current_model_config() return { "status": "healthy", "timestamp": datetime.now().isoformat(), "gemini_configured": bool(GEMINI_API_KEY), "azure_configured": bool(AZURE_STORAGE_CONNECTION_STRING or (AZURE_STORAGE_ACCOUNT_NAME and AZURE_STORAGE_ACCOUNT_KEY)), "current_model": model_config['name'], "quota_available": can_use_model(model_config) } if __name__ == "__main__": import uvicorn # Initialize model quota tracking for model in GEMINI_MODELS: model["last_rpm_reset"] = datetime.now() model["last_rpd_reset"] = datetime.now() print("\n" + "="*80) print("🚀 Starting Universal Invoice Splitter API with RPM Management") print("="*80) print(f"✅ Supports ALL invoice types") print(f"✅ Zydus Healthcare fallback (23xxxxxxxx pattern)") print(f"✅ Auto-merge null first pages") print(f"✅ RPM/RPD quota management") print(f"✅ Multi-model fallback") print("="*80) print(f"📋 Model Chain:") for i, model in enumerate(GEMINI_MODELS): print(f" {i+1}. {model['name']}") print( f" RPM: {model['max_requests_per_minute']}, RPD: {model['max_requests_per_day']}") print("="*80) print(f"🌐 Server: http://127.0.0.1:8000") print("="*80 + "\n") uvicorn.run(app, host=HOST, port=PORT, log_level="info")