app.py

import gradio as gr
import json
import os
from pathlib import Path
import base64
from typing import List, Dict, Any
import google.generativeai as genai
from PIL import Image
import PyPDF2
import io

# Configure Gemini API
# GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")
GEMINI_API_KEY = "AIzaSyB2b80YwNHs3Yj6RZOTL8wjXk2YhxCluOA"
if GEMINI_API_KEY:
    genai.configure(api_key=GEMINI_API_KEY)

EXTRACTION_PROMPT = """You are a shipping document data extraction specialist. Extract structured data from the provided shipping/logistics documents.

Extract the following fields into a JSON format:

{
    "poNumber": "Purchase Order Number",
    "shipFrom": "Origin/Ship From Location",
    "carrierType": "Transportation type (RAIL/TRUCK/etc)",
    "originCarrier": "Carrier name (CN/CPRS/etc)",
    "railCarNumber": "Rail car identifier",
    "totalQuantity": "Total quantity as number",
    "totalUnits": "Unit type (UNIT/MBF/MSFT/etc)",
    "accountName": "Customer/Account name",
    "inventories": {
        "items": [
            {
                "quantityShipped": "Quantity as number",
                "inventoryUnits": "Unit type",
                "productName": "Full product description",
                "productCode": "Product code/SKU",
                "product": {
                    "category": "Product category (OSB/Lumber/etc)",
                    "unit": "Unit count as number",
                    "pcs": "Pieces per unit",
                    "mbf": "Thousand board feet (if applicable)",
                    "sf": "Square feet (if applicable)",
                    "pcsHeight": "Height in inches",
                    "pcsWidth": "Width in inches",
                    "pcsLength": "Length in feet"
                },
                "customFields": [
                    "Mill||Mill Name",
                    "Vendor||Vendor Name"
                ]
            }
        ]
    }
}

IMPORTANT INSTRUCTIONS:
1. Extract ALL products/items found in the document
2. Convert text numbers to actual numbers (e.g., "54" → 54)
3. Parse dimensions carefully (e.g., "2X6X14" means height=2, width=6, length=14)
4. Calculate MBF/SF when possible from dimensions and piece count
5. If a field is not found, use null (not empty string)
6. For multiple products, create separate items in the inventories.items array
7. Extract custom fields like Mill, Vendor from document metadata

Return ONLY valid JSON, no markdown formatting or explanations."""


def extract_text_from_pdf(pdf_file) -> str:
    """Extract text from PDF file"""
    try:
        pdf_reader = PyPDF2.PdfReader(pdf_file)
        text = ""
        for page in pdf_reader.pages:
            text += page.extract_text() + "\n"
        return text
    except Exception as e:
        return f"Error extracting PDF text: {str(e)}"


def convert_pdf_to_images(pdf_file) -> List[Image.Image]:
    """Convert PDF pages to images"""
    try:
        from pdf2image import convert_from_path
        images = convert_from_path(pdf_file)
        return images
    except ImportError:
        # Fallback if pdf2image not available
        return []
    except Exception as e:
        print(f"Error converting PDF to images: {e}")
        return []


def process_files(files: List[str]) -> Dict[str, Any]:
    """Process uploaded files and extract text/images"""
    processed_data = {
        "files": [],
        "combined_text": "",
        "images": []
    }
    
    if not files:
        return processed_data
    
    for file_path in files:
        file_name = Path(file_path).name
        file_ext = Path(file_path).suffix.lower()
        
        file_data = {
            "filename": file_name,
            "type": file_ext,
            "content": ""
        }
        
        try:
            if file_ext == '.pdf':
                # Extract text from PDF
                text = extract_text_from_pdf(file_path)
                file_data["content"] = text
                processed_data["combined_text"] += f"\n--- {file_name} ---\n{text}\n"
                
                # Try to convert PDF to images for visual extraction
                images = convert_pdf_to_images(file_path)
                processed_data["images"].extend(images)
                
            elif file_ext in ['.jpg', '.jpeg', '.png', '.gif', '.bmp']:
                # Load image
                img = Image.open(file_path)
                processed_data["images"].append(img)
                file_data["content"] = f"Image file: {file_name}"
                processed_data["combined_text"] += f"\n--- {file_name} (Image) ---\n"
                
            elif file_ext in ['.txt']:
                # Read text file
                with open(file_path, 'r', encoding='utf-8') as f:
                    text = f.read()
                file_data["content"] = text
                processed_data["combined_text"] += f"\n--- {file_name} ---\n{text}\n"
            
            processed_data["files"].append(file_data)
            
        except Exception as e:
            file_data["content"] = f"Error processing file: {str(e)}"
            processed_data["files"].append(file_data)
    
    return processed_data


def extract_with_gemini(processed_data: Dict[str, Any], api_key: str) -> Dict[str, Any]:
    """Extract structured data using Gemini API"""
    
    if not api_key:
        return {"error": "Gemini API key not provided"}
    
    try:
        # Configure Gemini with provided API key
        genai.configure(api_key=api_key)
        
        # Use Gemini Pro Vision if images available, otherwise Pro
        if processed_data["images"]:
            model = genai.GenerativeModel('gemini-1.5-flash')
            
            # Prepare content with images and text
            content = [EXTRACTION_PROMPT]
            
            # Add text content
            if processed_data["combined_text"]:
                content.append(f"\nDocument Text:\n{processed_data['combined_text']}")
            
            # Add images (limit to first 5 for token management)
            for img in processed_data["images"][:5]:
                content.append(img)
            
            response = model.generate_content(content)
        else:
            # Text-only extraction
            model = genai.GenerativeModel('gemini-1.5-flash')
            prompt = f"{EXTRACTION_PROMPT}\n\nDocument Text:\n{processed_data['combined_text']}"
            response = model.generate_content(prompt)
        
        # Parse response
        response_text = response.text.strip()
        
        # Remove markdown code blocks if present
        if response_text.startswith("```json"):
            response_text = response_text[7:]
        if response_text.startswith("```"):
            response_text = response_text[3:]
        if response_text.endswith("```"):
            response_text = response_text[:-3]
        
        response_text = response_text.strip()
        
        # Parse JSON
        extracted_data = json.loads(response_text)
        
        return {
            "success": True,
            "data": extracted_data,
            "raw_response": response_text
        }
        
    except json.JSONDecodeError as e:
        return {
            "success": False,
            "error": f"JSON parsing error: {str(e)}",
            "raw_response": response.text if 'response' in locals() else None
        }
    except Exception as e:
        return {
            "success": False,
            "error": f"Extraction error: {str(e)}"
        }


def process_documents(files, api_key):
    """Main processing function"""
    
    if not files:
        return "⚠️ Please upload at least one document.", None, None
    
    if not api_key:
        return "⚠️ Please provide your Gemini API key.", None, None
    
    try:
        # Step 1: Process files
        status = "📄 Processing files..."
        processed_data = process_files(files)
        
        # Step 2: Extract with Gemini
        status = "🤖 Extracting data with Gemini AI..."
        result = extract_with_gemini(processed_data, api_key)
        
        if result.get("success"):
            # Format JSON output
            json_output = json.dumps(result["data"], indent=2)
            
            # Create summary
            data = result["data"]
            summary = f"""✅ **Extraction Successful!**

**Shipment Details:**
- PO Number: {data.get('poNumber', 'N/A')}
- Ship From: {data.get('shipFrom', 'N/A')}
- Carrier: {data.get('originCarrier', 'N/A')} ({data.get('carrierType', 'N/A')})
- Rail Car: {data.get('railCarNumber', 'N/A')}
- Total Quantity: {data.get('totalQuantity', 'N/A')} {data.get('totalUnits', '')}
- Account: {data.get('accountName', 'N/A')}

**Products Found:** {len(data.get('inventories', {}).get('items', []))}
"""
            
            return summary, json_output, None
        else:
            error_msg = f"❌ **Extraction Failed**\n\nError: {result.get('error', 'Unknown error')}"
            raw = result.get('raw_response', 'No response')
            return error_msg, None, raw
            
    except Exception as e:
        return f"❌ **Processing Error**\n\n{str(e)}", None, None


# Create Gradio Interface
with gr.Blocks(title="Shipping Document Extractor", theme=gr.themes.Soft()) as demo:
    gr.Markdown("""
    # 📦 Shipping Document Data Extractor
    
    Upload shipping documents (PDFs, images, etc.) to automatically extract structured data.
    Powered by Google Gemini AI.
    """)
    
    with gr.Row():
        with gr.Column(scale=1):
            api_key_input = gr.Textbox(
                label="🔑 Gemini API Key",
                placeholder="Enter your Gemini API key",
                type="password",
                value=GEMINI_API_KEY or ""
            )
            gr.Markdown("[Get your API key](https://makersuite.google.com/app/apikey)")
            
            file_input = gr.File(
                label="📁 Upload Documents",
                file_count="multiple",
                file_types=[".pdf", ".jpg", ".jpeg", ".png", ".txt"]
            )
            
            process_btn = gr.Button("🚀 Extract Data", variant="primary", size="lg")
            
            gr.Markdown("""
            ### Supported Files:
            - PDF documents
            - Images (JPG, PNG)
            - Text files
            """)
    
    with gr.Column(scale=2):
        status_output = gr.Markdown(label="Status")
        
        with gr.Tabs():
            with gr.Tab("📊 Extracted JSON"):
                json_output = gr.Code(
                    label="Structured Data (JSON)",
                    language="json",
                    lines=20
                )
                
                download_btn = gr.DownloadButton(
                    label="💾 Download JSON",
                    visible=False
                )
            
            with gr.Tab("🔍 Raw Response"):
                raw_output = gr.Code(
                    label="Raw API Response",
                    language="text",
                    lines=20
                )
    
    # Event handlers
    def update_download(json_str):
        if json_str:
            return gr.DownloadButton(visible=True, value=json_str)
        return gr.DownloadButton(visible=False)
    
    process_btn.click(
        fn=process_documents,
        inputs=[file_input, api_key_input],
        outputs=[status_output, json_output, raw_output]
    )
    
    json_output.change(
        fn=update_download,
        inputs=[json_output],
        outputs=[download_btn]
    )
    
    gr.Markdown("""
    ---
    ### 💡 Tips:
    - Upload multiple documents at once for batch processing
    - Better quality images/PDFs = better extraction accuracy
    - The AI extracts: PO numbers, carrier info, products, quantities, dimensions, and more
    """)

# Launch the app
if __name__ == "__main__":
    demo.launch()