Update app.py

app.py

@@ -4,6 +4,7 @@ import uuid
 import datetime
 import logging
 import re
+import gradio as gr
 
 # Setup logging
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
@@ -18,8 +19,8 @@ except ImportError:
     AutoModelForCausalLM = None
     AutoTokenizer = None
 
-# Initialize AI model (distilbert as placeholder; replace with fine-tuned model or Mistral-7B)
-model_name = "distilbert-base-uncased"  # Lightweight model for demo
+# Initialize AI model (t5-small for better accounting parsing; fallback to distilbert)
+model_name = "t5-small"  # Smaller model for free tier
 if TRANSFORMERS_AVAILABLE:
     try:
         tokenizer = AutoTokenizer.from_pretrained(model_name)
@@ -34,7 +35,7 @@ else:
     model = None
 
 # Database setup
-conn = sqlite3.connect("erp.db")
+conn = sqlite3.connect("erp.db", check_same_thread=False)
 cursor = conn.cursor()
 
 # Create tables
@@ -85,6 +86,7 @@ def initialize_chart_of_accounts():
         ("1.2.2", "Laptop", "Asset", "1.2", True, True),
         ("1.2.3", "Inventory", "Asset", "1.2", True, True),
         ("1.2.4", "Accounts Receivable", "Asset", "1.2", True, True),
+        ("1.2.5", "Bank", "Asset", "1.2", True, True),
         ("2", "Liabilities", "Liability", None, True, False),
         ("2.1", "Accounts Payable", "Liability", "2", True, True),
         ("2.2", "Loan Payable", "Liability", "2", True, True),
@@ -107,40 +109,37 @@ def initialize_chart_of_accounts():
     logging.info("Chart of accounts initialized.")
 
 # Enhanced fallback parser
-def parse_prompt(prompt):
+def parse_prompt(prompt, state):
     if model and tokenizer:
         try:
             input_text = f"""
             Parse the following accounting prompt into a JSON object with:
             - debit: {{account, type, amount}}
             - credit: {{account, type, amount}}
-            - payment_method: 'cash', 'credit', or null
+            - payment_method: 'cash', 'credit', 'bank', or null
             Prompt: {prompt}
             """
             inputs = tokenizer(input_text, return_tensors="pt")
             outputs = model.generate(**inputs, max_length=300)
             response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-            return json.loads(response)
+            return json.loads(response), state
         except Exception as e:
             logging.warning(f"Model parsing failed: {e}. Using fallback parser.")
 
     # Fallback parser
     prompt_lower = prompt.lower().strip()
     amount = None
-    # Extract amount
     match = re.search(r'\$[\d,.]+', prompt_lower)
     if match:
         try:
             amount = float(match.group().replace('$', '').replace(',', ''))
         except ValueError:
-            logging.error("Invalid amount format.")
-            return None
+            return {"error": "Invalid amount format."}, state
 
     if not amount:
-        logging.error("No amount found in prompt.")
-        return None
+        return {"error": "No amount found in prompt."}, state
 
-    # Map keywords to accounts and types
+    # Account mappings
     account_mappings = {
         "laptop": ("Laptop", "Asset"),
         "inventory": ("Inventory", "Asset"),
@@ -149,14 +148,17 @@ def parse_prompt(prompt):
         "plant": ("Plant", "Asset"),
         "office supplies": ("Office Supplies", "Expense"),
         "cash": ("Cash", "Asset"),
+        "bank": ("Bank", "Asset"),
         "receivable": ("Accounts Receivable", "Asset"),
         "sold goods": ("Sales Revenue", "Revenue"),
         "sales": ("Sales Revenue", "Revenue"),
         "rent": ("Rent Expense", "Expense"),
         "salary": ("Salary Expense", "Expense"),
         "paid": ("Cash", "Asset"),
-        "bought": ("Laptop", "Asset"),  # Default to Laptop; refine below
-        "purchased": ("Laptop", "Asset")
+        "bought": ("Laptop", "Asset"),
+        "purchased": ("Laptop", "Asset"),
+        "owner's draw": ("Drawings", "Equity"),
+        "loan": ("Loan Payable", "Liability")
     }
 
     debit_account = None
@@ -165,75 +167,81 @@ def parse_prompt(prompt):
     credit_type = None
     payment_method = None
 
-    # Determine debit account
-    for keyword, (account, acc_type) in account_mappings.items():
-        if keyword in prompt_lower:
-            if keyword in ["bought", "purchased"]:
-                # Check for specific asset/expense
-                for asset in ["laptop", "inventory", "machinery", "building", "plant", "office supplies"]:
-                    if asset in prompt_lower:
-                        debit_account, debit_type = account_mappings[asset]
-                        break
-                if not debit_account:
+    # Check if this is a follow-up response
+    if state.get("pending_prompt"):
+        follow_up = prompt_lower
+        if follow_up in ["cash", "credit", "bank"]:
+            payment_method = follow_up
+            parsed = state["pending_parsed"]
+            debit_account = parsed["debit"]["account"]
+            debit_type = parsed["debit"]["type"]
+            amount = parsed["debit"]["amount"]
+            if payment_method == "cash":
+                credit_account, credit_type = "Cash", "Asset"
+            elif payment_method == "bank":
+                credit_account, credit_type = "Bank", "Asset"
+            elif payment_method == "credit":
+                credit_account, credit_type = "Accounts Payable", "Liability"
+            state = {}  # Clear state
+        else:
+            return {"error": "Please respond with 'cash', 'credit', or 'bank'."}, state
+
+    else:
+        # Parse new prompt
+        for keyword, (account, acc_type) in account_mappings.items():
+            if keyword in prompt_lower:
+                if keyword in ["bought", "purchased"]:
+                    for asset in ["laptop", "inventory", "machinery", "building", "plant", "office supplies"]:
+                        if asset in prompt_lower:
+                            debit_account, debit_type = account_mappings[asset]
+                            break
+                    if not debit_account:
+                        debit_account, debit_type = account, acc_type
+                elif keyword in ["rent", "salary", "office supplies"]:
                     debit_account, debit_type = account, acc_type
-            elif keyword in ["rent", "salary", "office supplies"]:
-                debit_account, debit_type = account, acc_type
-            elif keyword in ["sold goods", "sales"]:
-                debit_account, debit_type = "Accounts Receivable", "Asset"
-                credit_account, credit_type = account, acc_type
-            break
+                elif keyword in ["sold goods", "sales"]:
+                    debit_account, debit_type = "Accounts Receivable", "Asset"
+                    credit_account, credit_type = account, acc_type
+                elif keyword == "owner's draw":
+                    debit_account, debit_type = "Drawings", "Equity"
+                    credit_account, credit_type = "Cash", "Asset"
+                break
 
-    # Determine credit account and payment method
-    if "cash" in prompt_lower:
-        credit_account, credit_type = "Cash", "Asset"
-        payment_method = "cash"
-    elif "credit" in prompt_lower:
-        credit_account, credit_type = "Accounts Payable", "Liability"
-        payment_method = "credit"
-    elif debit_account and not credit_account:
-        # Ambiguous payment method
-        return {"debit": {"account": debit_account, "type": debit_type, "amount": amount}, "credit": None, "payment_method": None}
+        if "cash" in prompt_lower:
+            credit_account, credit_type = "Cash", "Asset"
+            payment_method = "cash"
+        elif "bank" in prompt_lower:
+            credit_account, credit_type = "Bank", "Asset"
+            payment_method = "bank"
+        elif "credit" in prompt_lower:
+            credit_account, credit_type = "Accounts Payable", "Liability"
+            payment_method = "credit"
+        elif debit_account and not credit_account:
+            return {
+                "status": "clarify",
+                "message": "Was this bought on cash, credit, or bank?",
+                "pending_parsed": {"debit": {"account": debit_account, "type": debit_type, "amount": amount}}
+            }, {"pending_prompt": prompt, "pending_parsed": {"debit": {"account": debit_account, "type": debit_type, "amount": amount}}}
 
     if debit_account and credit_account:
         return {
             "debit": {"account": debit_account, "type": debit_type, "amount": amount},
             "credit": {"account": credit_account, "type": credit_type, "amount": amount},
             "payment_method": payment_method
-        }
-    logging.error("Prompt not recognized.")
-    return None
+        }, state
+    return {"error": "Prompt not recognized. Try 'Bought a laptop for $200' or 'Paid rent $400'."}, state
 
 # Generate journal entry
-def generate_journal_entry(prompt, follow_up_response=None):
-    parsed = parse_prompt(prompt)
-    if not parsed:
-        return "Unable to parse prompt. Please provide more details."
+def generate_journal_entry(parsed, state):
+    if "error" in parsed:
+        return parsed["error"], state
+    if parsed.get("status") == "clarify":
+        return parsed["message"], state
 
     debit_account = parsed["debit"]["account"]
     amount = parsed["debit"]["amount"]
-    payment_method = parsed.get("payment_method")
-
-    # Handle ambiguous payment method
-    if not payment_method and not follow_up_response:
-        return {"status": "clarify", "message": "Was this bought on cash or credit? (cash/credit)", "original_prompt": prompt}
-
-    # Determine credit account
-    credit_account = None
-    credit_type = None
-    if follow_up_response:
-        follow_up_lower = follow_up_response.lower()
-        if follow_up_lower == "cash":
-            credit_account, credit_type = "Cash", "Asset"
-        elif follow_up_lower == "credit":
-            credit_account, credit_type = "Accounts Payable", "Liability"
-        else:
-            return "Invalid response. Please specify 'cash' or 'credit'."
-    elif payment_method == "cash":
-        credit_account, credit_type = parsed["credit"]["account"], parsed["credit"]["type"]
-    elif payment_method == "credit":
-        credit_account, credit_type = "Accounts Payable", "Liability"
-    else:
-        return "Invalid payment method specified."
+    credit_account = parsed["credit"]["account"]
+    credit_type = parsed["credit"]["type"]
 
     # Validate accounts
     cursor.execute("SELECT account_id, account_type, allow_posting FROM chart_of_accounts WHERE account_name = ?", (debit_account,))
@@ -242,25 +250,24 @@ def generate_journal_entry(prompt, follow_up_response=None):
     credit_result = cursor.fetchone()
 
     if not debit_result or not credit_result:
-        return "One or both accounts not found in chart of accounts."
+        return "One or both accounts not found.", state
     if not debit_result[2] or not credit_result[2]:
-        return "Posting not allowed for one or both accounts."
-
-    # Validate account types
+        return "Posting not allowed for one or both accounts.", state
     if debit_result[1] != parsed["debit"]["type"] or credit_result[1] != credit_type:
-        return "Account type mismatch."
+        return "Account type mismatch.", state
 
     # Create journal entry
     entry_id = str(uuid.uuid4())
     date = datetime.datetime.now().isoformat()
+    description = state.get("pending_prompt", "Transaction")
     cursor.execute("""
     INSERT INTO journal_entries (entry_id, date, debit_account_id, credit_account_id, amount, description)
     VALUES (?, ?, ?, ?, ?, ?)
-    """, (entry_id, date, debit_result[0], credit_result[0], amount, prompt))
+    """, (entry_id, date, debit_result[0], credit_result[0], amount, description))
     conn.commit()
     logging.info(f"Journal entry created: Debit {debit_account} ${amount}, Credit {credit_account} ${amount}")
 
-    return f"Journal Entry Created: Debit {debit_account} ${amount}, Credit {credit_account} ${amount}"
+    return f"Journal Entry Created: Debit {debit_account} ${amount}, Credit {credit_account} ${amount}", {}
 
 # Generate T-account
 def generate_t_account(account_name):
@@ -298,46 +305,41 @@ def generate_t_account(account_name):
     
     return t_account
 
-# Main CLI loop
-def main():
+# Gradio chat function
+def chat_function(message, history, state=None):
+    if state is None:
+        state = {}
+    
+    # Initialize chart of accounts
     initialize_chart_of_accounts()
-    print("AI ERP System: Enter accounting prompts or 't-account <account_name>' to view T-accounts. Type 'exit' to quit.")
     
-    pending_prompt = None
-    while True:
-        try:
-            prompt = input("> ").strip()
-            if prompt.lower() == "exit":
-                break
-            if prompt.lower().startswith("t-account "):
-                account_name = prompt[10:].strip()
-                if account_name:
-                    print(generate_t_account(account_name))
-                else:
-                    print("Please specify an account name.")
-                continue
-            
-            # Handle follow-up response
-            if pending_prompt:
-                result = generate_journal_entry(pending_prompt, prompt)
-                pending_prompt = None
-            else:
-                result = generate_journal_entry(prompt)
-            
-            if isinstance(result, dict) and result["status"] == "clarify":
-                print(result["message"])
-                pending_prompt = result["original_prompt"]
-            else:
-                print(result)
-        
-        except KeyboardInterrupt:
-            print("\nExiting...")
-            break
-        except Exception as e:
-            logging.error(f"Error processing prompt: {e}")
-            print("An error occurred. Please try again.")
+    # Handle T-account request
+    if message.lower().startswith("t-account "):
+        account_name = message[10:].strip()
+        if account_name:
+            return generate_t_account(account_name), state
+        return "Please specify an account name.", state
+    
+    # Parse prompt and generate entry
+    parsed, state = parse_prompt(message, state)
+    response, state = generate_journal_entry(parsed, state)
+    
+    return response, state
+
+# Gradio interface
+with gr.Blocks() as demo:
+    gr.Markdown("# AI ERP System")
+    gr.Markdown("Enter accounting prompts like 'Bought a laptop for $200' or 't-account Laptop'. The system will ask for clarification if needed.")
+    chatbot = gr.Chatbot()
+    msg = gr.Textbox(placeholder="Type your prompt here...")
+    clear = gr.Button("Clear")
+    
+    # Maintain state
+    state = gr.State({})
     
-    conn.close()
+    msg.submit(chat_function, [msg, chatbot, state], [chatbot, state])
+    clear.click(lambda: None, None, chatbot, queue=False)
 
+# Launch Gradio
 if __name__ == "__main__":
-    main()
+    demo.launch(server_name="0.0.0.0", server_port=7860)