Update app.py

app.py

@@ -1,222 +1,47 @@
-import os
-import re
-import json
-import io
-import pandas as pd
-from typing import List, Optional
-from pydantic import BaseModel,Field
-from openai import OpenAI
 import gradio as gr
+from src.invoice_processor import InvoiceProcessor
 
-# File processing imports
-import PyPDF2
-from docx import Document
-from openpyxl import load_workbook
-from PIL import Image
-import pytesseract
-from striprtf.striprtf import rtf_to_text
+class InvoiceUI:
+    def __init__(self):
+        self.processor = InvoiceProcessor()
 
-from azure.core.credentials import AzureKeyCredential
-from azure.ai.documentintelligence import DocumentIntelligenceClient
-from azure.ai.documentintelligence.models import AnalyzeDocumentRequest, DocumentContentFormat, AnalyzeResult
-
-endpoint = "https://invoice-processing3141.cognitiveservices.azure.com/"
-credential = AzureKeyCredential(os.environ.get("AZURE_APT_KEY"))
-document_intelligence_client = DocumentIntelligenceClient(endpoint, credential)
-
-
-# Pydantic Models
-class InvoiceItem(BaseModel):
-    material: str = Field(description="The main product item in the row for which the details are being extracted.")
-    Packing: str = Field(description="The packing size of the product. This is usually the quantity present in one product")
-    quantity: str = Field(description="The number of products mentioned in the invoice")
-
-class Invoice(BaseModel):
-    po_date: str = Field(description="PO Date is usually the date on which the order was processed.")
-    order_value: float = Field(description="The total order value present at the end of the invoice product details.")
-    company_name: str = Field(description="The company sending in the invoice details.")
-    items: List[InvoiceItem] = Field(description="List of all Invoice Items present in the document.")
-
-# OpenAI Client Setup
-client = OpenAI(
-    base_url="https://api.studio.nebius.ai/v1/",
-    api_key=os.environ.get("NEBIUS_API_KEY")
-)
-
-# File Processing Functions
-def process_pdf(file_path: str) -> str:
-    text = ''
-    with open(file_path, 'rb') as file:
-        reader = PyPDF2.PdfReader(file)
-        for page in reader.pages:
-            text += page.extract_text()
-
-    if not text.strip():
-        for page in reader.pages:
-            if '/XObject' in page['/Resources']:
-                for obj in page['/Resources']['/XObject'].get_object():
-                    if page['/Resources']['/XObject'][obj]['/Subtype'] == '/Image':
-                        xObject = page['/Resources']['/XObject'][obj]
-                        data = xObject.get_data()
-                        image = Image.open(io.BytesIO(data))
-                        text += pytesseract.image_to_string(image)
-    return text
-
-def process_doc(file_path: str) -> str:
-    if file_path.lower().endswith('.doc'):
-        with open(file_path, "r", encoding='utf-8', errors='ignore') as file:
-            return rtf_to_text(file.read())
-    else:
-        doc = Document(file_path)
-        return '\n'.join([paragraph.text for paragraph in doc.paragraphs])
-
-def process_xlsx(file_path: str) -> str:
-    wb = load_workbook(file_path)
-    text = ''
-    for sheet in wb.sheetnames:
-        ws = wb[sheet]
-        for row in ws.iter_rows(values_only=True):
-            text += ' '.join(str(cell) for cell in row if cell is not None) + '\n'
-    return text
-
-def process_txt(file_path: str) -> str:
-    with open(file_path, 'r') as file:
-        return file.read()
-        
-def process_azure(file_path: str) -> str:
-    with open(file_path, "rb") as f:
-        poller = document_intelligence_client.begin_analyze_document("prebuilt-layout", body=f, output_content_format=DocumentContentFormat.MARKDOWN)
-        result: AnalyzeResult = poller.result()
-    return result.content
-
-def process_invoice_txt(file_path: str) -> str:
-    _, file_extension = os.path.splitext(file_path)
-
-    if file_extension.lower() == '.pdf':
-        return process_azure(file_path) #process_pdf(file_path)
-    elif file_extension.lower() in ['.doc']:
-        return process_doc(file_path)
-    elif file_extension.lower() in ['.docx']:
-        return process_azure(file_path)
-    elif file_extension.lower() in ['.xls', '.xlsx']:
-        return process_xlsx(file_path)
-    elif file_extension.lower() == '.txt':
-        return process_txt(file_path)
-    else:
-        raise ValueError(f"Unsupported file format: {file_extension}")
-
-# Invoice Processing Functions
-def process_invoice(invoice_data: str) -> str:
-    """Process invoice data using OpenAI API with guided JSON."""
-    system_prompt = """
-    You are an expert invoice data extraction assistant. Extract key information from invoices 
-    and return it in the specified JSON format. Handle variations in invoice layouts and infer 
-    information when not explicitly stated.
-    """
-
-    user_prompt = f"""
-    You are an expert at extracting invoice details from documents. Follow the below instructions carefully and return the items in a structured schema.
-
-    1. Extract the PO Date from the document. This is usually present before the product details.
-    2. Extract the total order value from the document. This would be the total value present at the end of the product details table.
-    3. Extract the name of the company sending in the order processing details.
-    4. Extract the details of the invoice. The invoice items might vary for each document but look for the main indicators needed for processing.
-        - Material: This is usually the item name or the product name in the table.
-        - Packing: This refers to the packing size of the product. The packing size can be in grams(GM), liters(ML) and capsules per pack ('S') for example 15'S is the packing size for a medicinal product.
-        - Quantity: This is the number of the material/product mentioned in the document.
-
-    Ensure that these instructions are followed and the data is extracted verbatim. Do not add your own comments or make up sentences.
-
-    Invoice:
-    {invoice_data}
-    """
-
-    try:
-        completion = client.chat.completions.create(
-            model="meta-llama/Meta-Llama-3.1-8B-Instruct-fast",
-            messages=[
-                {"role": "system", "content": system_prompt},
-                {"role": "user", "content": user_prompt}
-            ],
-            extra_body={"guided_json": Invoice.model_json_schema()}
-        )
-        
-        if hasattr(completion.choices[0].message, 'refusal') and completion.choices[0].message.refusal:
-            raise ValueError(f"API refused to process: {completion.choices[0].message.refusal}")
-        
-        return completion.choices[0].message.content.strip()
-    except Exception as e:
-        raise ValueError(f"Error processing with OpenAI API: {str(e)}")
-
-def process_and_save_to_excel(file_path: str) -> tuple:
-    """Process invoice file and save results to Excel."""
-    if not file_path:
-        raise ValueError("No file uploaded")
-        
-    temp_dir = "temp"
-    os.makedirs(temp_dir, exist_ok=True)
-
-    try:
-        # Process directly from Gradio-provided temp file path
-        invoice_text = process_invoice_txt(file_path)
-        json_output = process_invoice(invoice_text)
-        data = Invoice.model_validate_json(json_output)
-        
-        # Convert to DataFrame
-        rows = [{
-            'PO Date': data.po_date,
-            'Order Value': data.order_value,
-            'Company Name': data.company_name,
-            'Material': item.material,
-            'Packing': item.Packing,
-            'Quantity': item.quantity,
-        } for item in data.items]
-        
-        df = pd.DataFrame(rows)
-        
-        # Create output filename with timestamp
-        file_name = os.path.basename(file_path).split('.')[0]
-        output_filename = os.path.join(temp_dir, f"{file_name}.xlsx")
-        
-        # Save to Excel
-        df.to_excel(output_filename, index=False, engine='openpyxl')
-        
-        return output_filename, data.model_dump()
+    def create_interface(self) -> gr.Blocks:
+        """Create and return the Gradio interface."""
+        with gr.Blocks() as app:
+            gr.Markdown("# Invoice Data Processor")
             
-    except Exception as e:
-        raise ValueError(f"Failed to process invoice: {str(e)}")
-        
-# Gradio Interface
-def create_gradio_app():
-    with gr.Blocks() as app:
-        gr.Markdown("# Invoice Data Processor")
-        
-        with gr.Row():
-            file_input = gr.File(
-                label="Upload Invoice File (PDF, DOC, XLSX, or TXT)",
-                file_types=[".pdf", ".doc", ".docx", ".xlsx", ".txt"]
-            )
-        
-        with gr.Row():
-            process_button = gr.Button("Process Invoice")
-        
-        with gr.Row():
-            excel_output = gr.File(
-                label="Download Excel File",
-                file_types=[".xlsx"],
-                type="binary"  # Important for proper file downloading
+            with gr.Row():
+                file_input = gr.File(
+                    label="Upload Invoice File (PDF, DOC, XLSX, or TXT)",
+                    file_types=[".pdf", ".doc", ".docx", ".xlsx", ".txt"]
+                )
+            
+            with gr.Row():
+                process_button = gr.Button("Process Invoice")
+            
+            with gr.Row():
+                excel_output = gr.File(
+                    label="Download Excel File",
+                    file_types=[".xlsx"],
+                    type="binary"
+                )
+                json_output = gr.JSON(label="Extracted Data Preview")
+            
+            process_button.click(
+                fn=self.processor.process_and_save_to_excel,
+                inputs=[file_input],
+                outputs=[excel_output, json_output]
             )
-            json_output = gr.JSON(label="Extracted Data Preview")
         
-        process_button.click(
-            fn=process_and_save_to_excel,
-            inputs=[file_input],
-            outputs=[excel_output, json_output]
-        )
-    
-    return app
+        return app 
 
-if __name__ == "__main__":
-    app = create_gradio_app()
+def main():
+    """Initialize and launch the Gradio interface"""
+    ui = InvoiceUI()
+    app = ui.create_interface()
     app.launch(debug=True)
 
+if __name__ == "__main__":
+    main()
+
+