Update app.py

app.py

@@ -4,10 +4,15 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
 import uuid
 import datetime
 
-# Initialize AI model (Mistral-7B)
+# Initialize AI model (Mistral-7B placeholder; replace with actual model)
 model_name = "mistralai/Mixtral-8x7B-Instruct-v0.1"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name)
+try:
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    model = AutoModelForCausalLM.from_pretrained(model_name)
+except:
+    # Fallback for demo (mock AI response)
+    tokenizer = None
+    model = None
 
 # Database setup
 conn = sqlite3.connect("erp.db")
@@ -17,21 +22,24 @@ cursor = conn.cursor()
 cursor.execute("""
 CREATE TABLE IF NOT EXISTS chart_of_accounts (
     account_id TEXT PRIMARY KEY,
-    account_name TEXT,
-    account_type TEXT,
+    account_name TEXT NOT NULL,
+    account_type TEXT NOT NULL,
     parent_id TEXT,
     allow_budgeting BOOLEAN,
-    allow_posting BOOLEAN
+    allow_posting BOOLEAN,
+    FOREIGN KEY (parent_id) REFERENCES chart_of_accounts(account_id)
 )
 """)
 cursor.execute("""
 CREATE TABLE IF NOT EXISTS journal_entries (
     entry_id TEXT PRIMARY KEY,
-    date TEXT,
-    debit_account_id TEXT,
-    credit_account_id TEXT,
-    amount REAL,
-    description TEXT
+    date TEXT NOT NULL,
+    debit_account_id TEXT NOT NULL,
+    credit_account_id TEXT NOT NULL,
+    amount REAL NOT NULL,
+    description TEXT,
+    FOREIGN KEY (debit_account_id) REFERENCES chart_of_accounts(account_id),
+    FOREIGN KEY (credit_account_id) REFERENCES chart_of_accounts(account_id)
 )
 """)
 conn.commit()
@@ -45,14 +53,25 @@ ACCOUNT_RULES = {
     "Expense": {"increase": "Debit", "decrease": "Credit"}
 }
 
-# Sample chart of accounts
+# Initialize chart of accounts
 def initialize_chart_of_accounts():
     accounts = [
-        ("1", "Fixed Assets", "Asset", None, True, False),
-        ("1.1", "Plant", "Asset", "1", True, True),
-        ("1.2", "Machinery", "Asset", "1", True, True),
-        ("2", "Cash", "Asset", None, True, True),
-        ("3", "Accounts Payable", "Liability", None, True, True)
+        ("1", "Assets", "Asset", None, True, False),
+        ("1.1", "Fixed Assets", "Asset", "1", True, False),
+        ("1.1.1", "Plant", "Asset", "1.1", True, True),
+        ("1.1.2", "Machinery", "Asset", "1.1", True, True),
+        ("1.1.3", "Building", "Asset", "1.1", True, True),
+        ("1.2", "Current Assets", "Asset", "1", True, False),
+        ("1.2.1", "Cash", "Asset", "1.2", True, True),
+        ("1.2.2", "Laptop", "Asset", "1.2", True, True),
+        ("2", "Liabilities", "Liability", None, True, False),
+        ("2.1", "Accounts Payable", "Liability", "2", True, True),
+        ("3", "Equity", "Equity", None, True, False),
+        ("3.1", "Owner's Capital", "Equity", "3", True, True),
+        ("4", "Revenue", "Revenue", None, True, False),
+        ("4.1", "Sales", "Revenue", "4", True, True),
+        ("5", "Expenses", "Expense", None, True, False),
+        ("5.1", "Operating Expenses", "Expense", "5", True, True)
     ]
     cursor.executemany("""
     INSERT OR REPLACE INTO chart_of_accounts
@@ -61,61 +80,103 @@ def initialize_chart_of_accounts():
     """, accounts)
     conn.commit()
 
-# Parse prompt using AI model
+# Parse prompt using AI model (or fallback)
 def parse_prompt(prompt):
-    input_text = f"Parse the following accounting prompt and return a JSON object with debit and credit accounts, amount, and payment method (cash or credit). Prompt: {prompt}"
-    inputs = tokenizer(input_text, return_tensors="pt")
-    outputs = model.generate(**inputs, max_length=200)
-    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    if model and tokenizer:
+        input_text = f"""
+        Parse the following accounting prompt into a JSON object with:
+        - debit: {account, type, amount}
+        - credit: {account, type, amount}
+        - payment_method: 'cash' or 'credit' 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)
+        try:
+            return json.loads(response)
+        except:
+            pass
+
+    # Fallback parsing for common scenarios
+    prompt_lower = prompt.lower()
+    amount = None
+    for word in prompt_lower.split():
+        if word.startswith("$"):
+            try:
+                amount = float(word[1:])
+                break
+            except:
+                pass
     
-    # Assume model returns JSON-like string
-    try:
-        parsed = json.loads(response)
-        return parsed
-    except:
-        # Fallback parsing
-        if "laptop" in prompt.lower() and "cash" in prompt.lower():
-            return {
-                "debit": {"account": "Laptop", "type": "Asset", "amount": 200},
-                "credit": {"account": "Cash", "type": "Asset", "amount": 200},
-                "payment_method": "cash"
-            }
+    if not amount:
         return None
 
+    if "laptop" in prompt_lower:
+        debit_account = "Laptop"
+        debit_type = "Asset"
+        if "cash" in prompt_lower:
+            credit_account = "Cash"
+            credit_type = "Asset"
+            payment_method = "cash"
+        elif "credit" in prompt_lower:
+            credit_account = "Accounts Payable"
+            credit_type = "Liability"
+            payment_method = "credit"
+        else:
+            return {"debit": {"account": "Laptop", "type": "Asset", "amount": amount}, "credit": None, "payment_method": None}
+        return {
+            "debit": {"account": debit_account, "type": debit_type, "amount": amount},
+            "credit": {"account": credit_account, "type": credit_type, "amount": amount},
+            "payment_method": payment_method
+        }
+    return None
+
 # 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 clarify."
+        return "Unable to parse prompt. Please provide more details."
 
     debit_account = parsed["debit"]["account"]
-    credit_account = parsed["credit"]["account"]
     amount = parsed["debit"]["amount"]
     payment_method = parsed.get("payment_method")
 
-    # Follow-up for payment method
-    if not payment_method:
-        return "Wonderful, did you buy on credit? (Yes/No)"
+    # Handle ambiguous payment method
+    if not payment_method and not follow_up_response:
+        return {"status": "clarify", "message": "Wonderful, did you buy on credit? (Yes/No)"}
 
-    # Adjust credit account for credit purchases
+    # Determine credit account
+    credit_account = None
+    credit_type = None
     if follow_up_response and follow_up_response.lower() == "yes":
         credit_account = "Accounts Payable"
         credit_type = "Liability"
-    else:
+    elif payment_method == "cash":
+        credit_account = parsed["credit"]["account"]
         credit_type = parsed["credit"]["type"]
+    elif payment_method == "credit":
+        credit_account = "Accounts Payable"
+        credit_type = "Liability"
+    else:
+        return "Invalid payment method specified."
 
-    # Find account IDs
-    cursor.execute("SELECT account_id, allow_posting FROM chart_of_accounts WHERE account_name = ?", (debit_account,))
+    # Validate accounts
+    cursor.execute("SELECT account_id, account_type, allow_posting FROM chart_of_accounts WHERE account_name = ?", (debit_account,))
     debit_result = cursor.fetchone()
-    cursor.execute("SELECT account_id, allow_posting FROM chart_of_accounts WHERE account_name = ?", (credit_account,))
+    cursor.execute("SELECT account_id, account_type, allow_posting FROM chart_of_accounts WHERE account_name = ?", (credit_account,))
     credit_result = cursor.fetchone()
 
     if not debit_result or not credit_result:
-        return "Invalid account names."
-    if not debit_result[1] or not credit_result[1]:
+        return "One or both accounts not found in chart of accounts."
+    if not debit_result[2] or not credit_result[2]:
         return "Posting not allowed for one or both accounts."
 
-    # Save journal entry
+    # Validate account types
+    if debit_result[1] != parsed["debit"]["type"] or credit_result[1] != credit_type:
+        return "Account type mismatch."
+
+    # Create journal entry
     entry_id = str(uuid.uuid4())
     date = datetime.datetime.now().isoformat()
     cursor.execute("""
@@ -138,15 +199,21 @@ def generate_t_account(account_name):
     SELECT amount, description, 'Debit' as type FROM journal_entries WHERE debit_account_id = ?
     UNION
     SELECT amount, description, 'Credit' as type FROM journal_entries WHERE credit_account_id = ?
+    ORDER BY date
     """, (account_id, account_id))
     entries = cursor.fetchall()
     
-    t_account = f"T-Account for {account_name}\n{'-'*40}\nDebit | Credit\n{'-'*40}\n"
+    t_account = f"T-Account for {account_name}\n{'='*50}\n{'Debit':<20} | {'Credit':<20} | Description\n{'-'*50}\n"
+    debit_total = 0
+    credit_total = 0
     for amount, desc, entry_type in entries:
         if entry_type == "Debit":
-            t_account += f"${amount} |       | {desc}\n"
+            t_account += f"${amount:<19} | {'':<20} | {desc}\n"
+            debit_total += amount
         else:
-            t_account += f"       | ${amount} | {desc}\n"
+            t_account += f"{'':<20} | ${amount:<19} | {desc}\n"
+            credit_total += amount
+    t_account += f"{'-'*50}\nTotal Debit: ${debit_total:<10} | Total Credit: ${credit_total}\n"
     
     return t_account
 
@@ -154,19 +221,22 @@ def generate_t_account(account_name):
 if __name__ == "__main__":
     initialize_chart_of_accounts()
     
-    # Test prompt
+    # Test prompt (cash purchase)
     prompt = "Bought a laptop for $200 on cash"
     result = generate_journal_entry(prompt)
     print(result)
     
+    # Test prompt (ambiguous payment method)
+    prompt = "Bought a laptop for $300"
+    result = generate_journal_entry(prompt)
+    print(result)
+    if isinstance(result, dict) and result["status"] == "clarify":
+        result = generate_journal_entry(prompt, "Yes")
+        print(result)
+    
     # Test T-account
     t_account = generate_t_account("Laptop")
     print(t_account)
     
-    # Test follow-up
-    prompt = "Bought a laptop for $200"
-    result = generate_journal_entry(prompt)
-    print(result)
-    if "did you buy on credit" in result.lower():
-        result = generate_journal_entry(prompt, "Yes")
-        print(result)
+    # Clean up
+    conn.close()