src/gemini_processing.py

from google import genai
import json
import os

class GeminiInvoiceProcessor:
    def __init__(self, model_name="gemini-2.0-flash"):
        """Initialize the Gemini processor.
        
        Args:
            model_name: Optional model name. If not provided, uses the one from settings.
        """            
        api_key = os.getenv('GEMINI_API_KEY')
        if not api_key:
            raise ValueError("GEMINI_API_KEY not found in environment variables")
            
        self.client = genai.Client(api_key=api_key)
        self.model_name = model_name
    
    def _generate_prompt(self, user_fields):
        """Generate the prompt for invoice processing.
        
        Args:
            user_fields: List of fields to extract from the order details.
        """
        gemini_prompt = f"""
        Given the following document extract details about the following in JSON format:
        1. purchase_order_details - Usually contains details about the company placing the purchase order. The details would include company name, address, GSTIN details among others.
        2. supplier_details - Usually contains details about the company providing the supplies. The details would include company name, address, GSTIN details and others.
        3. order_details - This information is usually presented in a table format. Extract the below information for each order item along with its serial number from the table
            - {user_fields}
        4. Total Order Amount - Extract the total order amount for this invoice.

        Maintain absolutely precise field names, exactly as specified, without any alterations whatsoever."""

        return gemini_prompt
    
    def process_invoice(self, content, user_fields):
        """
        Process an invoice file and extract details.
        
        Args:
            file_path: Path to the invoice file
            user_fields: List of fields to extract from the invoice
        """
    
        prompt = self._generate_prompt(user_fields)

        response = self.client.models.generate_content(model=self.model_name,contents=[prompt, content]).text
        res = response.strip("```").replace('json', '')
        try:
            json_out = json.loads(res)
            return json_out
        except json.JSONDecodeError as e:
            print('ERROR PARSING JSON - MANUAL PARSING REQUIRED')
            raise e


if __name__ == "__main__":
    # Test the processor
    DATA_DIR = "data/001405_24_OR_24000798277 (1).pdf"
    test_file = os.path.join(DATA_DIR, 'dr03.Doc')
    user_fields = ['Company', 'Item', 'Quantity', 'Cost']
    
    # Read the file content
    with open(test_file, 'r') as f:
        invoice_text = f.read()
    
    invoice_processor = GeminiInvoiceProcessor()
    json_out = invoice_processor.process_invoice(invoice_text, user_fields)
    
    print(json_out)