Create main.py

main.py

@@ -0,0 +1,366 @@
+import os
+import json
+import logging
+import re
+import tempfile
+import time
+from datetime import datetime
+from flask import Flask, request, jsonify
+from flask_cors import CORS
+import pandas as pd
+import pypdf
+import google.generativeai as genai
+
+# Setup logging
+logging.basicConfig(level=logging.INFO)
+
+app = Flask(__name__)
+CORS(app)  # Enable CORS for all routes
+
+# Get API key securely
+api_key = os.getenv('Gemini')
+if not api_key:
+    raise ValueError("Gemini API key not found in environment variables")
+
+def configure_gemini(api_key):
+    """Configure Gemini AI model."""
+    try:
+        genai.configure(api_key=api_key)
+        return genai.GenerativeModel('gemini-2.0-flash')
+    except Exception as e:
+        logging.error(f"Error configuring Gemini: {str(e)}")
+        raise
+
+def read_pdf_pages(file_path):
+    """Extract text from each page of a PDF file as list of strings."""
+    try:
+        pages = []
+        with open(file_path, 'rb') as f:
+            reader = pypdf.PdfReader(f)
+            for page in reader.pages:
+                text = page.extract_text() or ""
+                pages.append(text)
+        return pages
+    except Exception as e:
+        logging.error(f"Error reading PDF: {str(e)}")
+        raise
+
+PROMPT = """Analyze this bank statement and extract transactions in JSON format with these fields:
+- Date (format DD/MM/YYYY)
+- Description
+- Amount (just the integer value)
+- Type (categorize into one of the following based on the transaction nature):
+* 'income' - money received from customers, sales, services rendered
+* 'expense' - operational costs, purchases, payments made
+* 'asset' - purchase of equipment, property, vehicles, or other assets
+* 'liability' - taking on debt, loans received, credit facilities
+* 'equity' - owner investments, capital contributions, retained earnings transfers
+* 'transfer' - money moved between own accounts, internal transfers
+* 'investment' - securities purchases, investment account funding, portfolio additions
+* 'loan_repayment' - paying back borrowed money, loan principal payments
+* 'capital_injection' - owner or investor adding money to the business
+- Customer Name (Only if Type is 'income' and if no name is extracted write 'general income'. For all other types, extract relevant party name or write 'N/A')
+- City (In address of bank statement)
+- Destination_of_funds (categorize based on Type and description):
+* If 'expense': Water and electricity, Salaries and wages, Repairs & Maintenance, Motor vehicle expenses, Projects Expenses, Hardware expenses, Refunds, Accounting fees, Loan interest, Bank charges, Insurance, SARS PAYE UIF, Advertising & Marketing, Logistics and distribution, Fuel, Website hosting fees, Rentals, Subscriptions, Computer internet and Telephone, Staff training, Travel and accommodation, Depreciation, Other expenses
+* If 'income': 'income'
+* If 'asset': Equipment, Property, Vehicles, Technology, Furniture, Other assets
+* If 'liability': Bank loan, Credit facility, Supplier credit, Other liabilities
+* If 'equity': Owner investment, Retained earnings, Share capital, Other equity
+* If 'transfer': Internal transfer
+* If 'investment': Securities, Mutual funds, Fixed deposits, Other investments
+* If 'loan_repayment': Loan repayment
+* If 'capital_injection': Capital injection
+- ignore opening or closing balances.
+- extract the amount in full including decimals.
+Return ONLY the raw JSON object, without any surrounding text, explanations, or markdown fences like ```json.
+
+Return ONLY valid JSON with this structure:
+{
+"transactions": [
+{
+"Date": "string",
+"Description": "string",
+"Customer_name": "string",
+"City": "string",
+"Amount": number,
+"Type": "string",
+"Destination_of_funds": "string"
+}
+]
+}"""
+
+def get_text_prompt_with_date():
+    """Generate TEXT_PROMPT with current date context."""
+    current_date = datetime.now().strftime("%d/%m/%Y")
+    return f"""IMPORTANT: Today's date is {current_date}. If the user does not specify a date for a transaction, use {current_date} as the default date.
+
+Analyze the following natural language text and extract transactions in JSON format with these fields:
+- Date (format DD/MM/YYYY) - USE {current_date} IF NO DATE IS SPECIFIED
+- Description
+- Amount (just the integer value)
+- Type (categorize into one of the following based on the transaction nature):
+* 'income' - money received from customers, sales, services rendered
+* 'expense' - operational costs, purchases, payments made
+* 'asset' - purchase of equipment, property, vehicles, or other assets
+* 'liability' - taking on debt, loans received, credit facilities
+* 'equity' - owner investments, capital contributions, retained earnings transfers
+* 'transfer' - money moved between own accounts, internal transfers
+* 'investment' - securities purchases, investment account funding, portfolio additions
+* 'loan_repayment' - paying back borrowed money, loan principal payments
+* 'capital_injection' - owner or investor adding money to the business
+- Customer Name (Only if Type is 'income' and if no name is extracted write 'general income'. For all other types, extract relevant party name or write 'N/A')
+- City (extract from any address information provided or write 'N/A' if not available)
+- Destination_of_funds (categorize based on Type and description):
+* If 'expense': Water and electricity, Salaries and wages, Repairs & Maintenance, Motor vehicle expenses, Projects Expenses, Hardware expenses, Refunds, Accounting fees, Loan interest, Bank charges, Insurance, SARS PAYE UIF, Advertising & Marketing, Logistics and distribution, Fuel, Website hosting fees, Rentals, Subscriptions, Computer internet and Telephone, Staff training, Travel and accommodation, Depreciation, Other expenses
+* If 'income': 'income'
+* If 'asset': Equipment, Property, Vehicles, Technology, Furniture, Other assets
+* If 'liability': Bank loan, Credit facility, Supplier credit, Other liabilities
+* If 'equity': Owner investment, Retained earnings, Share capital, Other equity
+* If 'transfer': Internal transfer
+* If 'investment': Securities, Mutual funds, Fixed deposits, Other investments
+* If 'loan_repayment': Loan repayment
+* If 'capital_injection': Capital injection
+- ignore opening or closing balances.
+
+Return ONLY valid JSON with this structure:
+{{
+"transactions": [
+{{
+"Date": "string",
+"Description": "string",
+"Customer_name": "string",
+"City": "string",
+"Amount": number,
+"Type": "string",
+"Destination_of_funds": "string"
+}}
+]
+}}
+important: Return an empty array if no transactions are in the text. Do not make up false data.
+"""
+
+def repair_json_with_gemini(model, broken_json_string):
+    """Makes a second API call to Gemini to fix a broken JSON string."""
+    logging.info("Attempting to repair broken JSON with another Gemini call...")
+    repair_prompt = f"""The following text is a JSON object that is syntactically incorrect. 
+It might have missing commas, brackets, or other errors. 
+Please fix the syntax to make it a valid JSON object. 
+Return ONLY the corrected, raw JSON object and nothing else.
+
+Broken JSON:
+{broken_json_string}
+"""
+    try:
+        resp = model.generate_content(repair_prompt)
+        return resp.text
+    except Exception as e:
+        logging.error(f"Error during JSON repair call: {e}")
+        raise ValueError("Failed to repair the JSON string.")
+
+def call_gemini_with_retry_custom(model, text, prompt, retries=3, backoff_factor=2):
+    """Call Gemini with retries, now with JSON repair logic."""
+    for attempt in range(1, retries + 1):
+        try:
+            resp = model.generate_content([prompt, text])
+            response_text = resp.text
+            try:
+                # First attempt to parse the original response
+                return extract_json_from_response(response_text)
+            except Exception:
+                # If parsing fails, trigger the repair process
+                logging.warning("Initial JSON parsing failed. Attempting repair.")
+                repaired_text = repair_json_with_gemini(model, response_text)
+                return extract_json_from_response(repaired_text) # Parse the repaired text
+        except Exception as e:
+            msg = str(e)
+            if '429' in msg or 'RateLimit' in msg and attempt < retries:
+                wait = backoff_factor ** attempt
+                logging.warning(f"Rate limit hit, retrying in {wait}s (attempt {attempt}/{retries})")
+                time.sleep(wait)
+            else:
+                logging.error(f"Error processing with Gemini after retries: {msg}")
+                raise
+
+def call_gemini_with_retry(model, text, retries=3, backoff_factor=2):
+    """Call Gemini with retries, now with JSON repair logic."""
+    # This function now simply calls the custom one with the default PROMPT
+    return call_gemini_with_retry_custom(model, text, PROMPT, retries, backoff_factor)
+
+def extract_json_from_response(response_text):
+    """Extract valid JSON from Gemini's response, even if it's embedded."""
+    # First, try to find a JSON block enclosed in markdown fences
+    match = re.search(r'```json\s*(\{.*?\})\s*```', response_text, re.DOTALL)
+    if match:
+        json_string = match.group(1)
+    else:
+        # If no markdown, find the first '{' and the last '}'
+        match = re.search(r'{.*}', response_text, re.DOTALL)
+        if not match:
+            raise ValueError("No valid JSON object found in the LLM response")
+        json_string = match.group(0)
+
+    try:
+        # Clean and load the extracted JSON string
+        return json.loads(json_string)
+    except json.JSONDecodeError as e:
+        logging.error(f"Failed to parse extracted JSON. Error: {e}")
+        logging.error(f"Problematic JSON string was: {json_string}")
+        raise ValueError(f"Could not parse JSON from LLM response: {e}")
+
+@app.route('/process-pdf', methods=['POST'])
+def process_pdf():
+    """Handle PDF upload, process it in page-chunks with Gemini, and aggregate results."""
+    try:
+        if 'file' not in request.files:
+            return jsonify({'error': 'No file uploaded'}), 400
+        file = request.files['file']
+        if file.filename == '' or not file.filename.lower().endswith('.pdf'):
+            return jsonify({'error': 'A valid PDF file must be uploaded'}), 400
+
+        # Save to temp file
+        with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as tmp:
+            file.save(tmp.name)
+            file_path = tmp.name
+
+        try:
+            model = configure_gemini(api_key)
+            pages = read_pdf_pages(file_path)
+            all_transactions = []
+
+            for idx, page_text in enumerate(pages, start=1):
+                if not page_text.strip():
+                    continue
+                logging.info(f"Processing page {idx}/{len(pages)}")
+                result = call_gemini_with_retry(model, page_text)
+                txs = result.get('transactions', [])
+                all_transactions.extend(txs)
+
+            return jsonify({'transactions': all_transactions})
+        finally:
+            os.remove(file_path)
+
+    except ValueError as ve:
+        logging.warning(f"Client error: {ve}")
+        return jsonify({'error': str(ve)}), 400
+    except Exception as e:
+        logging.error(f"Internal server error: {e}")
+        return jsonify({'error': 'Internal server error'}), 500
+
+@app.route('/process-text', methods=['POST'])
+def process_text():
+    """Handle text input and extract transactions using Gemini."""
+    try:
+        # Get JSON data from request
+        data = request.get_json()
+        if not data or 'text' not in data:
+            return jsonify({'error': 'No text provided'}), 400
+        
+        text_input = data['text']
+        
+        if not text_input.strip():
+            return jsonify({'error': 'Text input cannot be empty'}), 400
+        
+        # Configure Gemini model
+        model = configure_gemini(api_key)
+        
+        # Generate prompt with current date
+        text_prompt = get_text_prompt_with_date()
+        
+        # Process the text with Gemini
+        logging.info("Processing text input for transaction extraction")
+        result = call_gemini_with_retry_custom(model, text_input, text_prompt)
+        
+        transactions = result.get('transactions', [])
+        
+        return jsonify({'transactions': transactions})
+        
+    except ValueError as ve:
+        logging.warning(f"Client error: {ve}")
+        return jsonify({'error': str(ve)}), 400
+    except Exception as e:
+        logging.error(f"Internal server error: {e}")
+        return jsonify({'error': 'Internal server error'}), 500
+
+@app.route('/transaction-types', methods=['GET'])
+def get_transaction_types():
+    """Return available transaction types and their categories."""
+    transaction_types = {
+        "types": [
+            {
+                "type": "income",
+                "description": "Money received from customers, sales, services rendered",
+                "destination_categories": ["income"]
+            },
+            {
+                "type": "expense",
+                "description": "Operational costs, purchases, payments made",
+                "destination_categories": [
+                    "Water and electricity", "Salaries and wages", "Repairs & Maintenance",
+                    "Motor vehicle expenses", "Projects Expenses", "Hardware expenses",
+                    "Refunds", "Accounting fees", "Loan interest", "Bank charges",
+                    "Insurance", "SARS PAYE UIF", "Advertising & Marketing",
+                    "Logistics and distribution", "Fuel", "Website hosting fees",
+                    "Rentals", "Subscriptions", "Computer internet and Telephone",
+                    "Staff training", "Travel and accommodation", "Depreciation",
+                    "Other expenses"
+                ]
+            },
+            {
+                "type": "asset",
+                "description": "Purchase of equipment, property, vehicles, or other assets",
+                "destination_categories": [
+                    "Equipment", "Property", "Vehicles", "Technology", "Furniture", "Other assets"
+                ]
+            },
+            {
+                "type": "liability",
+                "description": "Taking on debt, loans received, credit facilities",
+                "destination_categories": [
+                    "Bank loan", "Credit facility", "Supplier credit", "Other liabilities"
+                ]
+            },
+            {
+                "type": "equity",
+                "description": "Owner investments, capital contributions, retained earnings transfers",
+                "destination_categories": [
+                    "Owner investment", "Retained earnings", "Share capital", "Other equity"
+                ]
+            },
+            {
+                "type": "transfer",
+                "description": "Money moved between own accounts, internal transfers",
+                "destination_categories": ["Internal transfer"]
+            },
+            {
+                "type": "investment",
+                "description": "Securities purchases, investment account funding, portfolio additions",
+                "destination_categories": [
+                    "Securities", "Mutual funds", "Fixed deposits", "Other investments"
+                ]
+            },
+            {
+                "type": "loan_repayment",
+                "description": "Paying back borrowed money, loan principal payments",
+                "destination_categories": ["Loan repayment"]
+            },
+            {
+                "type": "capital_injection",
+                "description": "Owner or investor adding money to the business",
+                "destination_categories": ["Capital injection"]
+            }
+        ]
+    }
+    return jsonify(transaction_types)
+
+@app.route('/health', methods=['GET'])
+def health_check():
+    """Health check endpoint."""
+    return jsonify({
+        'status': 'healthy',
+        'timestamp': datetime.now().isoformat(),
+        'version': '2.0.0'
+    })
+
+if __name__ == '__main__':
+    app.run(debug=True, host="0.0.0.0", port=7860)