Spaces:
Paused
Paused
| 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 | |
| # ============================================================================ | |
| 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() | |
| 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 | |
| }) | |
| 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()}) | |
| 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'] | |
| } | |
| 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") | |