from google import genai from pydantic import BaseModel, Field from typing import List, Optional, Dict, Tuple import pdf2image import os from pathlib import Path import concurrent.futures from dataclasses import dataclass from functools import partial import logging from PIL import Image from dotenv import load_dotenv load_dotenv() class InvoiceItem(BaseModel): """Represents a single item in an invoice.""" product_name: str = Field(description="The name of the product") batch_number: str = Field(description="The batch number of the product") expiry_date: str = Field(description="The expiry date (format: MM/YY)") mrp: str = Field(description="Maximum Retail Price") quantity: int = Field(description="Product quantity") class InvoiceData(BaseModel): """Represents the complete invoice data including headers.""" headers: List[str] = Field( description="Column headers from the invoice table", default_factory=list ) items: List[InvoiceItem] = Field( description="List of extracted invoice items", default_factory=list ) class HeaderExtraction(BaseModel): """Model for extracting headers separately.""" headers: List[str] = Field( description="The column headers found in the invoice table" ) @dataclass class PageData: """Container for page processing data.""" idx: int image_path: str headers: List[str] items: List[InvoiceItem] def extract_headers(client: genai.Client, image_path: str, model_id: str) -> List[str]: """ Extract column headers from the first page of the invoice. Args: client: The Gemini API client image_path: Path to the image file model_id: The model ID to use for extraction Returns: List of column headers """ header_prompt = """ Extract only the column headers from this invoice table. Return them exactly as they appear, maintaining their order from left to right. Only extract the headers, not any data from the rows. """ image_file = client.files.upload( file=image_path, config={'display_name': 'invoice_header_page'} ) response = client.models.generate_content( model=model_id, contents=[header_prompt, image_file], config={ 'response_mime_type': 'application/json', 'response_schema': HeaderExtraction } ) return response.parsed.headers if response.parsed else [] def setup_client() -> genai.Client: """Create and return a Gemini API client.""" return genai.Client(api_key=os.getenv("GEMINI_API_KEY")) def save_image(image: Image, temp_dir: Path, idx: int) -> str: """ Save a single page image to disk. Args: image: The PDF page image (PIL Image) temp_dir: Directory to save the image idx: Page index Returns: Path to the saved image """ image_path = str(temp_dir / f"page_{idx+1}.jpg") image.save(image_path, "JPEG") return image_path def process_single_page( page_data: Tuple[int, Image.Image, Path, List[str], genai.Client, str] ) -> PageData: """ Process a single page of the PDF. Args: page_data: Tuple containing (page_index, page_image, temp_dir, headers, client, model_id) Returns: PageData object containing extracted information """ idx, image, temp_dir, headers, client, model_id = page_data # Save image image_path = save_image(image, temp_dir, idx) # First page: extract headers if idx == 0: headers = extract_headers(client, image_path, model_id) prompt = """ Extract product details from this invoice table. Use the exact column headers you see in the table. """ else: headers_str = ", ".join(headers) prompt = f""" Extract product details from this invoice table. This is page {idx + 1} of the same invoice. Use these column headers: {headers_str} Ensure the extracted data aligns with these columns in order. """ # Process image image_file = client.files.upload( file=image_path, config={'display_name': f'invoice_page_{idx+1}'} ) response = client.models.generate_content( model=model_id, contents=[prompt, image_file], config={ 'response_mime_type': 'application/json', 'response_schema': InvoiceData } ) items = response.parsed.items if response.parsed and response.parsed.items else [] return PageData(idx=idx, image_path=image_path, headers=headers, items=items) def process_pdf_with_headers(pdf_path: str, max_workers: int = 3) -> InvoiceData: """ Process a PDF invoice while preserving column header context using parallel processing. Args: pdf_path: Path to the PDF file max_workers: Maximum number of concurrent workers Returns: InvoiceData object containing headers and extracted items """ # Convert PDF pages to images images = pdf2image.convert_from_path(pdf_path) # Create temp directory temp_dir = Path("content/temp") temp_dir.mkdir(parents=True, exist_ok=True) # Initialize shared resources client = setup_client() model_id = "gemini-2.0-flash" headers: List[str] = [] # Prepare data for parallel processing page_data = [] try: # Process first page separately to get headers first_page = process_single_page((0, images[0], temp_dir, headers, client, model_id)) headers = first_page.headers all_items = first_page.items # Prepare remaining pages for parallel processing remaining_pages = [ (i, img, temp_dir, headers, client, model_id) for i, img in enumerate(images[1:], start=1) ] # Process remaining pages in parallel with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor: future_to_page = { executor.submit(process_single_page, page): page[0] for page in remaining_pages } # Collect results as they complete for future in concurrent.futures.as_completed(future_to_page): page_idx = future_to_page[future] try: page_result = future.result() all_items.extend(page_result.items) except Exception as e: logging.error(f"Error processing page {page_idx}: {str(e)}") finally: # Cleanup temporary files for file in temp_dir.glob("*.jpg"): try: file.unlink() except Exception as e: logging.warning(f"Failed to delete temporary file {file}: {str(e)}") return InvoiceData(headers=headers, items=all_items) def main(): """Main function to demonstrate usage.""" # Configure logging logging.basicConfig( level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s' ) try: invoice_data = process_pdf_with_headers( "/Users/krishnaadithya/Desktop/dev/invoice_processing_2.0/pdf_only/expiry_invoice/DR REDDYS PE 1194.pdf", max_workers=3 # Adjust based on your system and API limits ) # Print headers print("Column Headers:", ", ".join(invoice_data.headers)) print("\nExtracted Items:") # Print results for item in invoice_data.items: print(f"Product: {item.product_name}") print(f"Batch: {item.batch_number}") print(f"Expiry: {item.expiry_date}") print(f"MRP: {item.mrp}") print(f"Quantity: {item.quantity}") print("-" * 50) except Exception as e: logging.error(f"Error processing invoice: {str(e)}") if __name__ == "__main__": main()