memra 0.2.13__py3-none-any.whl → 0.2.15__py3-none-any.whl

memra-ops/logic/invoice_tools.py

@@ -1,668 +0,0 @@
-"""
-Invoice processing tools for the Memra API server
-"""
-
-import os
-import logging
-import json
-import tempfile
-from typing import Dict, Any, Optional, List
-import subprocess
-from PIL import Image
-import base64
-import io
-import uuid
-from pathlib import Path
-import requests
-
-logger = logging.getLogger(__name__)
-
-class PDFProcessor:
-    """Process PDF files and extract content using vision model"""
-    
-    def __init__(self):
-        self.upload_dir = "/tmp/uploads"
-        self.screenshots_dir = "/tmp/screenshots"
-        # Ensure directories exist
-        os.makedirs(self.screenshots_dir, exist_ok=True)
-    
-    def process_pdf(self, file_path: str, schema: Dict[str, Any] = None) -> Dict[str, Any]:
-        """Process a PDF file and extract invoice data using vision model with schema"""
-        try:
-            if not file_path:
-                return {
-                    "success": False,
-                    "error": "No file path provided"
-                }
-            
-            # Handle uploaded files
-            if file_path.startswith('/uploads/'):
-                full_path = os.path.join(self.upload_dir, os.path.basename(file_path))
-            else:
-                full_path = file_path
-            
-            if not os.path.exists(full_path):
-                return {
-                    "success": False,
-                    "error": f"PDF file not found: {file_path}"
-                }
-            
-            logger.info(f"Processing PDF: {file_path}")
-            
-            # Step 1: Create invoice-specific directory
-            invoice_id = str(uuid.uuid4())
-            invoice_dir = os.path.join(self.screenshots_dir, invoice_id)
-            os.makedirs(invoice_dir, exist_ok=True)
-            
-            # Step 2: Convert PDF pages to screenshots
-            logger.info("Creating screenshots...")
-            screenshot_paths = self._create_screenshots(full_path, invoice_dir)
-            if not screenshot_paths:
-                return {
-                    "success": False,
-                    "error": "Failed to create screenshots from PDF (timeout or error)"
-                }
-            
-            # Step 3: Send screenshots + prompt + schema to vision model
-            logger.info(f"Sending {len(screenshot_paths)} screenshots to vision model with schema...")
-            
-            # Construct the comprehensive prompt with schema
-            vision_prompt = self._build_schema_prompt(schema)
-            
-            # Log and print the prompt being sent to vision model
-            logger.info(f"Vision Model Prompt: {vision_prompt}")
-            print(f"\n🔎 VISION MODEL PROMPT:")
-            print("=" * 60)
-            print(vision_prompt)
-            print("=" * 60)
-            
-            # Send to vision model and get JSON response
-            vision_response = self._call_vision_model_with_schema(screenshot_paths[0], vision_prompt)
-            
-            # Log and print the JSON response from vision model
-            logger.info(f"Vision Model JSON Response: {vision_response}")
-            print(f"\n📝 VISION MODEL JSON RESPONSE:")
-            print("=" * 60)
-            print(vision_response)
-            print("=" * 60)
-            
-            # Step 4: Parse the JSON response
-            try:
-                # Clean the response - remove markdown code blocks if present
-                cleaned_response = vision_response.strip()
-                if cleaned_response.startswith('```json'):
-                    cleaned_response = cleaned_response[7:]  # Remove ```json
-                if cleaned_response.endswith('```'):
-                    cleaned_response = cleaned_response[:-3]  # Remove ```
-                cleaned_response = cleaned_response.strip()
-                
-                extracted_data = json.loads(cleaned_response)
-                logger.info(f"Successfully parsed JSON response: {extracted_data}")
-                
-                # Convert to MCP bridge expected format
-                mcp_format_data = self._convert_to_mcp_format(extracted_data)
-                
-            except json.JSONDecodeError as e:
-                logger.error(f"Failed to parse JSON response: {e}")
-                # If it's an error response, create a structured error
-                if "error" in vision_response.lower():
-                    mcp_format_data = {
-                        "headerSection": {"vendorName": "", "subtotal": 0},
-                        "billingDetails": {"invoiceNumber": "", "invoiceDate": "", "dueDate": ""},
-                        "chargesSummary": {"document_total": 0, "secondary_tax": 0, "lineItemsBreakdown": []},
-                        "status": "vision_model_error",
-                        "error_message": vision_response
-                    }
-                else:
-                    mcp_format_data = {
-                        "headerSection": {"vendorName": "", "subtotal": 0},
-                        "billingDetails": {"invoiceNumber": "", "invoiceDate": "", "dueDate": ""},
-                        "chargesSummary": {"document_total": 0, "secondary_tax": 0, "lineItemsBreakdown": []},
-                        "status": "json_parse_error",
-                        "raw_response": vision_response
-                    }
-            
-            return {
-                "success": True,
-                "data": {
-                    "file_path": file_path,
-                    "invoice_id": invoice_id,
-                    "screenshots_dir": invoice_dir,
-                    "screenshot_count": len(screenshot_paths),
-                    "vision_prompt": vision_prompt,
-                    "vision_response": vision_response,
-                    "extracted_data": mcp_format_data
-                }
-            }
-            
-        except Exception as e:
-            logger.error(f"PDF processing failed: {str(e)}")
-            return {
-                "success": False,
-                "error": str(e)
-            }
-    
-    def _create_screenshots(self, pdf_path: str, output_dir: str) -> List[str]:
-        """Create high-resolution screenshots of PDF pages"""
-        try:
-            # Use pdftoppm to convert PDF to images with lower resolution for speed
-            cmd = [
-                'pdftoppm', 
-                '-png',           # Output format
-                '-r', '100',      # Very low resolution (100 DPI) for maximum speed
-                '-cropbox',       # Use crop box for consistent sizing
-                '-f', '1',        # Start from page 1
-                '-l', '1',        # Only process first page for speed
-                pdf_path,         # Input PDF
-                os.path.join(output_dir, 'page')  # Output prefix
-            ]
-            
-            # Add timeout to prevent hanging
-            result = subprocess.run(cmd, capture_output=True, text=True, timeout=15)
-            if result.returncode != 0:
-                logger.error(f"pdftoppm failed: {result.stderr}")
-                return []
-            
-            # Find generated image files
-            screenshot_paths = []
-            for file in sorted(os.listdir(output_dir)):
-                if file.endswith('.png'):
-                    image_path = os.path.join(output_dir, file)
-                    screenshot_paths.append(image_path)
-            
-            logger.info(f"Created {len(screenshot_paths)} screenshots in {output_dir}")
-            return screenshot_paths
-            
-        except subprocess.TimeoutExpired:
-            logger.error(f"Screenshot creation timed out after 15 seconds")
-            return []
-        except Exception as e:
-            logger.error(f"Screenshot creation failed: {str(e)}")
-            return []
-    
-    def _build_schema_prompt(self, schema: Dict[str, Any]) -> str:
-        """Build a prompt that includes the database schema"""
-        
-        logger.info(f"_build_schema_prompt called with schema type: {type(schema)}")
-        logger.info(f"Schema content: {schema}")
-        
-        # Default base prompt with essential fields
-        base_prompt = '''Extract invoice data from this image and return ONLY a JSON object with these specific fields:
-- vendor_name: The company name at the top of the invoice
-- invoice_number: The invoice number or ID
-- invoice_date: The date the invoice was issued (YYYY-MM-DD format)
-- total_amount: The total invoice amount
-- line_items: Array of items with descriptions and amounts
-
-Look specifically for the company/vendor name prominently displayed on the invoice.
-
-Return ONLY valid JSON with no additional text or explanation.'''
-        
-        # If no schema provided, return the base prompt
-        if not schema:
-            logger.info("No schema provided, returning base prompt")
-            return base_prompt
-        
-        # Handle different schema formats
-        columns = None
-        if isinstance(schema, list):
-            # Client sends array of column objects directly
-            columns = schema
-        elif isinstance(schema, dict) and "columns" in schema:
-            # Standard format with columns array
-            columns = schema["columns"]
-        else:
-            # Unknown format, use base prompt
-            return base_prompt
-        
-        # Build field descriptions from schema
-        field_descriptions = []
-        logger.info(f"Building prompt from {len(columns)} columns")
-        for col in columns:
-            # Handle both formats: {"column_name": "x"} and {"name": "x"}
-            name = col.get("column_name") or col.get("name", "")
-            col_type = col.get("data_type") or col.get("type", "")
-            logger.info(f"Processing column: {name} ({col_type})")
-            
-            # Skip system fields
-            if name and name not in ["id", "created_at", "updated_at", "status", "raw_json"]:
-                # Add helpful descriptions for key fields
-                if name == "vendor_name":
-                    field_descriptions.append(f"- {name}: The company name at the top of the invoice")
-                elif name == "invoice_number":
-                    field_descriptions.append(f"- {name}: The invoice number or ID")
-                elif name == "invoice_date":
-                    field_descriptions.append(f"- {name}: The date the invoice was issued (YYYY-MM-DD format)")
-                elif name == "total_amount":
-                    field_descriptions.append(f"- {name}: The total invoice amount")
-                elif name == "due_date":
-                    field_descriptions.append(f"- {name}: The invoice due date (YYYY-MM-DD format)")
-                elif name == "tax_amount":
-                    field_descriptions.append(f"- {name}: The tax amount on the invoice")
-                elif name == "line_items":
-                    field_descriptions.append(f"- {name}: Array of items with descriptions and amounts")
-                else:
-                    field_descriptions.append(f"- {name}: {col_type}")
-        
-        # If we have field descriptions, use them; otherwise use base prompt
-        logger.info(f"Built {len(field_descriptions)} field descriptions")
-        if field_descriptions:
-            schema_text = "\n".join(field_descriptions)
-            full_prompt = f'''Extract invoice data from this image and return ONLY a JSON object with these specific fields:
-{schema_text}
-
-Look specifically for the company/vendor name prominently displayed on the invoice.
-
-Return ONLY valid JSON with no additional text or explanation.'''
-            logger.info(f"Returning schema-based prompt with {len(field_descriptions)} fields")
-            return full_prompt
-        else:
-            logger.info("No field descriptions built, returning base prompt")
-            return base_prompt
-    
-    def _call_vision_model_with_schema(self, image_path: str, prompt: str) -> str:
-        """Call vision model with image and comprehensive prompt using Hugging Face"""
-        try:
-            # Import Hugging Face client
-            from huggingface_hub import InferenceClient
-            
-            # Get API key from environment
-            api_key = os.getenv("HUGGINGFACE_API_KEY")
-            
-            if not api_key:
-                logger.error("HUGGINGFACE_API_KEY environment variable is not set")
-                return json.dumps({
-                    "error": "Hugging Face API key not configured",
-                    "message": "Please set HUGGINGFACE_API_KEY environment variable",
-                    "expected_structure": {
-                        "vendor_name": "string",
-                        "invoice_number": "string", 
-                        "invoice_date": "YYYY-MM-DD",
-                        "due_date": "YYYY-MM-DD",
-                        "amount": 0.0,
-                        "tax_amount": 0.0,
-                        "line_items": "[]"
-                    }
-                })
-            
-            # Log the API key format for debugging (first few characters)
-            logger.info(f"Using Hugging Face API key: {api_key[:10]}...")
-            
-            # Create Hugging Face client with correct parameter name
-            client = InferenceClient(token=api_key)
-            
-            # Encode image to base64
-            def encode_image(image_path):
-                with open(image_path, "rb") as f:
-                    return base64.b64encode(f.read()).decode("utf-8")
-            
-            base64_image = encode_image(image_path)
-            
-            # Log the request details for debugging
-            logger.info(f"Making request to Hugging Face with model: meta-llama/Llama-4-Maverick-17B-128E-Instruct")
-            logger.info(f"Prompt length: {len(prompt)} characters")
-            logger.info(f"Image base64 length: {len(base64_image)} characters")
-            
-            # Call the model using the working approach - exactly as in your example
-            response = client.chat.completions.create(
-                model="meta-llama/Llama-4-Maverick-17B-128E-Instruct",
-                messages=[
-                    {
-                        "role": "user",
-                        "content": [
-                            {"type": "text", "text": prompt},
-                            {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{base64_image}"}}
-                        ]
-                    }
-                ],
-                max_tokens=500,
-            )
-            
-            # Extract the response content
-            extracted_text = response.choices[0].message.content
-            
-            logger.info(f"Hugging Face API call successful")
-            logger.info(f"Response length: {len(extracted_text)} characters")
-            
-            return extracted_text
-            
-        except Exception as e:
-            logger.error(f"Vision model call failed: {str(e)}")
-            return json.dumps({
-                "error": f"Vision model processing failed - {str(e)}"
-            })
-
-    def _convert_to_mcp_format(self, extracted_data: Dict[str, Any]) -> Dict[str, Any]:
-        """Convert extracted data to MCP bridge expected format"""
-        try:
-            # Handle nested response structure from vision model
-            # The vision model might return data in a nested structure like:
-            # {"data": {"invoice_number": "123", "vendor_name": "ABC Corp"}}
-            # or directly: {"invoice_number": "123", "vendor_name": "ABC Corp"}
-            
-            # If the data is nested, extract it
-            if isinstance(extracted_data, dict) and "data" in extracted_data:
-                actual_data = extracted_data["data"]
-            else:
-                actual_data = extracted_data
-            
-            # Handle both expected format and actual vision model output format
-            # Vision model might return: InvoiceNumber, InvoiceDate, InvoiceTotal, etc.
-            # Expected format: invoice_number, invoice_date, amount, etc.
-            
-            # Extract invoice number (try both formats)
-            invoice_number = (
-                actual_data.get("invoice_number") or 
-                actual_data.get("InvoiceNumber") or 
-                actual_data.get("invoiceNumber") or 
-                ""
-            )
-            
-            # Extract invoice date (try both formats)
-            invoice_date = (
-                actual_data.get("invoice_date") or 
-                actual_data.get("InvoiceDate") or 
-                actual_data.get("invoiceDate") or 
-                ""
-            )
-            
-            # Convert date format if needed
-            if invoice_date:
-                # Convert MM/DD/YY to YYYY-MM-DD format
-                if "/" in invoice_date and len(invoice_date.split("/")) == 3:
-                    parts = invoice_date.split("/")
-                    month, day, year = parts[0], parts[1], parts[2]
-                    if len(year) == 2:
-                        year = "20" + year
-                    invoice_date = f"{year}-{month.zfill(2)}-{day.zfill(2)}"
-            
-            # Extract due date (try both formats)
-            due_date = (
-                actual_data.get("due_date") or 
-                actual_data.get("DueDate") or 
-                actual_data.get("dueDate") or 
-                ""
-            )
-            
-            # Convert due date format if needed
-            if due_date:
-                # Convert MM/DD/YY to YYYY-MM-DD format
-                if "/" in due_date and len(due_date.split("/")) == 3:
-                    parts = due_date.split("/")
-                    month, day, year = parts[0], parts[1], parts[2]
-                    if len(year) == 2:
-                        year = "20" + year
-                    due_date = f"{year}-{month.zfill(2)}-{day.zfill(2)}"
-            
-            # Extract amount (try both formats)
-            amount = (
-                actual_data.get("total_amount") or  # Add this - matches the prompt
-                actual_data.get("amount") or 
-                actual_data.get("InvoiceTotal") or 
-                actual_data.get("invoiceTotal") or 
-                actual_data.get("total") or 
-                0
-            )
-            
-            # Extract vendor name (try both formats)
-            vendor_name = (
-                actual_data.get("vendor_name") or 
-                actual_data.get("VendorName") or 
-                actual_data.get("vendorName") or 
-                actual_data.get("Company") or 
-                actual_data.get("company") or 
-                ""
-            )
-            
-            # Extract tax amount (try both formats)
-            tax_amount = (
-                actual_data.get("tax_amount") or 
-                actual_data.get("TaxAmount") or 
-                actual_data.get("taxAmount") or 
-                0
-            )
-            
-            # Extract line items (try both formats)
-            line_items = (
-                actual_data.get("line_items") or 
-                actual_data.get("Order") or 
-                actual_data.get("order") or 
-                actual_data.get("LineItems") or 
-                actual_data.get("lineItems") or 
-                []
-            )
-            
-            if isinstance(line_items, str):
-                try:
-                    line_items = json.loads(line_items)
-                except:
-                    line_items = []
-            
-            # Convert to MCP bridge format
-            mcp_format = {
-                "headerSection": {
-                    "vendorName": vendor_name,
-                    "subtotal": float(amount)
-                },
-                "billingDetails": {
-                    "invoiceNumber": invoice_number,
-                    "invoiceDate": invoice_date,
-                    "dueDate": due_date
-                },
-                "chargesSummary": {
-                    "document_total": float(amount),
-                    "secondary_tax": float(tax_amount),
-                    "lineItemsBreakdown": line_items
-                },
-                "status": "processed"
-            }
-            
-            return mcp_format
-            
-        except Exception as e:
-            logger.error(f"Error converting to MCP format: {str(e)}")
-            return {
-                "headerSection": {"vendorName": "", "subtotal": 0},
-                "billingDetails": {"invoiceNumber": "", "invoiceDate": "", "dueDate": ""},
-                "chargesSummary": {"document_total": 0, "secondary_tax": 0, "lineItemsBreakdown": []},
-                "status": "conversion_error"
-            }
-
-class DatabaseQueryTool:
-    """Query database schemas and data"""
-    
-    def __init__(self, credentials: Dict[str, Any]):
-        self.credentials = credentials
-    
-    def get_schema(self, table_name: str) -> Dict[str, Any]:
-        """Get database schema for a table"""
-        # Mock schema for now
-        return {
-            "success": True,
-            "data": {
-                "table": table_name,
-                "columns": [
-                    {"name": "id", "type": "integer", "primary_key": True},
-                    {"name": "vendor_name", "type": "text"},
-                    {"name": "invoice_number", "type": "text"},
-                    {"name": "invoice_date", "type": "date"},
-                    {"name": "amount", "type": "decimal"},
-                    {"name": "created_at", "type": "timestamp"}
-                ]
-            }
-        }
-
-class OCRTool:
-    """Perform OCR on images and documents"""
-    
-    def extract_text(self, input_data: Dict[str, Any]) -> Dict[str, Any]:
-        """Extract text from document"""
-        return {
-            "success": True,
-            "data": {
-                "extracted_text": "Sample extracted text from document"
-            }
-        }
-
-class InvoiceExtractionWorkflow:
-    """Extract structured data from invoices"""
-    
-    def __init__(self):
-        pass
-    
-    def extract_data(self, text: str, schema: Dict[str, Any]) -> Dict[str, Any]:
-        """Extract structured data from invoice text or JSON"""
-        try:
-            # Check if the input is already JSON (from vision model)
-            if isinstance(text, dict):
-                # Input is already structured data from vision model
-                invoice_data = text
-            else:
-                # Try to parse as JSON first
-                try:
-                    invoice_data = json.loads(text)
-                except json.JSONDecodeError:
-                    # Fall back to text parsing
-                    invoice_data = self._parse_text_to_data(text)
-            
-            # If we still have empty data, this might be a case where we should
-            # use the data from a previous tool (PDFProcessor)
-            if not invoice_data.get("vendor_name") and not invoice_data.get("invoice_number"):
-                logger.warning("No invoice data found in input - this might be a workflow issue")
-                return {
-                    "success": False,
-                    "data": {
-                        "extracted_data": {
-                            "vendor_name": "",
-                            "invoice_number": "",
-                            "invoice_date": "",
-                            "amount": 0.0,
-                            "tax_amount": 0.0,
-                            "line_items": "[]",
-                            "status": "no_data_from_previous_tool"
-                        }
-                    }
-                }
-            
-            # Convert date format if needed
-            if invoice_data.get("invoice_date"):
-                invoice_data["invoice_date"] = self._convert_date_format(invoice_data["invoice_date"])
-            
-            # Ensure line_items is a JSON string
-            if isinstance(invoice_data.get("line_items"), list):
-                invoice_data["line_items"] = json.dumps(invoice_data["line_items"])
-            
-            # Set status
-            invoice_data["status"] = "processed"
-            
-            return {
-                "success": True,
-                "data": {
-                    "extracted_data": invoice_data
-                }
-            }
-            
-        except Exception as e:
-            logger.error(f"Invoice extraction failed: {str(e)}")
-            return {
-                "success": False,
-                "data": {
-                    "extracted_data": {
-                        "vendor_name": "",
-                        "invoice_number": "",
-                        "invoice_date": "",
-                        "amount": 0.0,
-                        "tax_amount": 0.0,
-                        "line_items": "[]",
-                        "status": "error"
-                    }
-                }
-            }
-    
-    def _parse_text_to_data(self, text: str) -> Dict[str, Any]:
-        """Parse text to extract invoice data (fallback method)"""
-        lines = text.split('\n')
-        invoice_data = {
-            "vendor_name": "",
-            "invoice_number": "",
-            "invoice_date": "",
-            "amount": 0.0,
-            "tax_amount": 0.0,
-            "line_items": "[]",
-            "status": "processed"
-        }
-        
-        # Extract data from the text using real parsing
-        for line in lines:
-            line = line.strip()
-            if "Invoice Number:" in line:
-                invoice_data["invoice_number"] = line.split(":")[1].strip()
-            elif "Invoice Date:" in line:
-                invoice_data["invoice_date"] = line.split(":")[1].strip()
-            elif "Order total:" in line:
-                amount_str = line.split(":")[1].strip()
-                try:
-                    invoice_data["amount"] = float(amount_str)
-                except:
-                    pass
-            elif "GST - HST / TPS -TVH:" in line:
-                tax_str = line.split(":")[1].strip()
-                try:
-                    invoice_data["tax_amount"] = float(tax_str)
-                except:
-                    pass
-            elif "SUPERIOR PROPANE" in line:
-                invoice_data["vendor_name"] = "SUPERIOR PROPANE"
-            elif "CHEP CANADA INC" in line:
-                invoice_data["vendor_name"] = "CHEP CANADA INC"
-        
-        return invoice_data
-    
-    def _convert_date_format(self, date_str: str) -> str:
-        """Convert date from MM/DD/YY to YYYY-MM-DD format"""
-        try:
-            # Handle MM/DD/YY format
-            if "/" in date_str and len(date_str.split("/")) == 3:
-                parts = date_str.split("/")
-                month, day, year = parts[0], parts[1], parts[2]
-                
-                # Convert 2-digit year to 4-digit
-                if len(year) == 2:
-                    year = "20" + year
-                
-                # Ensure proper formatting
-                return f"{year}-{month.zfill(2)}-{day.zfill(2)}"
-            
-            return date_str
-        except:
-            return date_str
-
-class DataValidator:
-    """Validate data against schemas"""
-    
-    def validate(self, data: Dict[str, Any], schema: Dict[str, Any]) -> Dict[str, Any]:
-        """Validate data against schema"""
-        return {
-            "success": True,
-            "data": {
-                "valid": True,
-                "errors": []
-            }
-        }
-
-class PostgresInsert:
-    """Insert data into PostgreSQL database"""
-    
-    def __init__(self, credentials: Dict[str, Any]):
-        self.credentials = credentials
-    
-    def insert_record(self, table: str, data: Dict[str, Any]) -> Dict[str, Any]:
-        """Insert a record into the database"""
-        return {
-            "success": True,
-            "data": {
-                "table": table,
-                "inserted_id": 123,
-                "message": "Record inserted successfully"
-            }
-        }