Update main.py

main.py

@@ -102,6 +102,153 @@ If the text below does not specify a year or date, reasonable assume {current_da
 {FINANCIAL_DOC_PROMPT}
 """
 
+# -------------------------------------------------------------------------
+# CATEGORIZATION LOGIC - TYPE-BASED (FIX FOR THE BUG)
+# -------------------------------------------------------------------------
+
+def categorize_transaction(transaction):
+    """
+    Categorizes a transaction based strictly on its Type field.
+    This prevents keyword-based misclassification.
+    
+    Args:
+        transaction: dict with keys including 'Type', 'Description', 'Destination_of_funds'
+    
+    Returns:
+        dict with added 'Account_Category' field
+    """
+    tx_type = transaction.get('Type', '').lower()
+    description = transaction.get('Description', '').lower()
+    destination = transaction.get('Destination_of_funds', '').lower()
+    
+    # Add the categorized account field
+    account_category = "Uncategorized"
+    
+    # ========== INCOME TYPE ==========
+    if tx_type == 'income':
+        # All income should map to revenue accounts, NOT expenses
+        if any(keyword in description for keyword in ['sales', 'service', 'revenue', 'invoice']):
+            account_category = "Sales Revenue"
+        elif any(keyword in description for keyword in ['interest', 'dividend']):
+            account_category = "Interest Income"
+        elif any(keyword in description for keyword in ['transfer', 'deposit', 'payment']):
+            # This fixes the "Income Trap" - transfers FROM others are income
+            account_category = "Other Income"
+        else:
+            account_category = "Other Income"
+    
+    # ========== EXPENSE TYPE ==========
+    elif tx_type == 'expense':
+        # Map based on Destination_of_funds or description keywords
+        # This is TYPE-FIRST, so "cash" in description won't make it an asset
+        
+        # Specific expense categories based on your system
+        if 'salaries' in destination or 'wages' in destination or 'salary' in description:
+            account_category = "Salaries and Wages"
+        elif 'water' in destination or 'electricity' in destination:
+            account_category = "Water and Electricity"
+        elif 'fuel' in destination or 'petrol' in description:
+            account_category = "Fuel"
+        elif 'rental' in destination or 'rent' in description:
+            account_category = "Rentals"
+        elif 'marketing' in destination or 'advertising' in destination:
+            account_category = "Advertising & Marketing"
+        elif 'repair' in destination or 'maintenance' in destination:
+            account_category = "Repairs & Maintenance"
+        elif 'vehicle' in destination or 'motor' in destination:
+            account_category = "Motor Vehicle Expenses"
+        elif 'hardware' in destination:
+            account_category = "Hardware Expenses"
+        elif 'accounting' in destination:
+            account_category = "Accounting Fees"
+        elif 'insurance' in destination:
+            account_category = "Insurance"
+        elif 'bank' in destination and 'charge' in destination:
+            account_category = "Bank Charges"
+        elif 'loan' in destination and 'interest' in destination:
+            account_category = "Loan Interest"
+        elif 'subscription' in destination:
+            account_category = "Subscriptions"
+        elif 'internet' in destination or 'telephone' in destination:
+            account_category = "Computer Internet and Telephone"
+        elif 'training' in destination:
+            account_category = "Staff Training"
+        elif 'travel' in destination or 'accommodation' in destination:
+            account_category = "Travel and Accommodation"
+        elif 'depreciation' in destination:
+            account_category = "Depreciation"
+        
+        # Special cases based on description (but still respecting expense type)
+        elif 'atm' in description and 'cash' in description:
+            # This fixes the "Cash Trap" - ATM withdrawals are drawings, not assets
+            account_category = "Owner's Drawings"
+        elif 'payment to' in description:
+            # Payment to suppliers/vendors
+            if any(word in description for word in ['fabric', 'printing', 'material']):
+                account_category = "Cost of Sales"
+            else:
+                account_category = "Miscellaneous Expense"
+        else:
+            account_category = "Miscellaneous Expense"
+    
+    # ========== ASSET TYPE ==========
+    elif tx_type == 'asset':
+        if 'equipment' in destination or 'equipment' in description:
+            account_category = "Equipment"
+        elif 'vehicle' in destination or 'vehicle' in description:
+            account_category = "Vehicles"
+        elif 'property' in destination or 'property' in description:
+            account_category = "Property"
+        elif 'technology' in destination or 'computer' in description:
+            account_category = "Technology"
+        elif 'furniture' in destination:
+            account_category = "Furniture"
+        else:
+            account_category = "Other Assets"
+    
+    # ========== LIABILITY TYPE ==========
+    elif tx_type == 'liability':
+        if 'bank loan' in destination or 'loan' in description:
+            account_category = "Bank Loan"
+        elif 'credit' in destination:
+            account_category = "Credit Facility"
+        else:
+            account_category = "Other Liabilities"
+    
+    # ========== EQUITY TYPE ==========
+    elif tx_type == 'equity':
+        if 'owner' in destination or 'capital' in description:
+            account_category = "Owner Investment"
+        elif 'retained' in destination:
+            account_category = "Retained Earnings"
+        else:
+            account_category = "Other Equity"
+    
+    # ========== TRANSFER TYPE ==========
+    elif tx_type == 'transfer':
+        account_category = "Internal Transfer"
+    
+    # ========== INVESTMENT TYPE ==========
+    elif tx_type == 'investment':
+        if 'securities' in destination or 'stock' in description:
+            account_category = "Securities"
+        elif 'mutual' in destination:
+            account_category = "Mutual Funds"
+        else:
+            account_category = "Other Investments"
+    
+    # ========== LOAN REPAYMENT TYPE ==========
+    elif tx_type == 'loan_repayment':
+        account_category = "Loan Repayment"
+    
+    # ========== CAPITAL INJECTION TYPE ==========
+    elif tx_type == 'capital_injection':
+        account_category = "Capital Injection"
+    
+    # Add the category to the transaction
+    transaction['Account_Category'] = account_category
+    return transaction
+
 # -------------------------------------------------------------------------
 # HELPER FUNCTIONS
 # -------------------------------------------------------------------------
@@ -149,11 +296,27 @@ def call_gemini_with_retry(model, content, prompt, retries=2):
             response = model.generate_content([prompt, content])
             
             try:
-                return extract_json_from_response(response.text)
+                result = extract_json_from_response(response.text)
+                
+                # POST-PROCESSING: Categorize each transaction based on Type
+                if 'transactions' in result:
+                    result['transactions'] = [
+                        categorize_transaction(tx) for tx in result['transactions']
+                    ]
+                
+                return result
             except ValueError as ve:
                 # Value error here contains the broken JSON string
                 broken_json = str(ve)
-                return repair_json_with_gemini(model, broken_json)
+                repaired = repair_json_with_gemini(model, broken_json)
+                
+                # Categorize repaired transactions too
+                if 'transactions' in repaired:
+                    repaired['transactions'] = [
+                        categorize_transaction(tx) for tx in repaired['transactions']
+                    ]
+                
+                return repaired
                 
         except Exception as e:
             if "429" in str(e) or "ResourceExhausted" in str(e):
@@ -414,7 +577,7 @@ def health_check():
     return jsonify({
         'status': 'healthy',
         'timestamp': datetime.now().isoformat(),
-        'version': '2.1.0',
+        'version': '2.2.0',
         'vision_support': PDF_IMAGE_SUPPORT
     })