Update app.py

app.py

@@ -1,6 +1,7 @@
 import streamlit as st
 import pandas as pd
 from huggingface_hub import InferenceClient
+import re
 
 # Initialize hosted inference client
 client = InferenceClient(model="google/flan-t5-base")
@@ -9,60 +10,111 @@ client = InferenceClient(model="google/flan-t5-base")
 account_map = {
     "rent": "60001",
     "utilities": "60002",
+    "capital": "30000",
     "cash": "10001",
-    "bank": "10002"
+    "bank": "10002",
+    "sales": "40001",
+    "supplies": "50001",
+    "salary": "50002"
 }
 
 # Simulated business segments
 segment = {
     "company": "01",
-    "business_type": "102",  # Grocery Store
+    "business_type": "102",
     "location": "001",
     "cost_center": "001",
     "future": "000"
 }
 
+# Session state to store entries
+if "gl_entries" not in st.session_state:
+    st.session_state.gl_entries = []
+
 def parse_prompt(prompt):
-    response = client.text_generation(prompt=f"Extract accounting entry: {prompt}", max_new_tokens=50)
-    return response
+    return client.text_generation(prompt=f"Extract accounting entry: {prompt}", max_new_tokens=50).strip()
 
 def handle_gl_entry(prompt):
-    # Simulate parsing response
-    if "rent" in prompt.lower():
+    prompt_lower = prompt.lower()
+    amount = 0
+    account_name = ""
+
+    # Extract amount using regex
+    amount_match = re.search(r'(\d{1,3}(,\d{3})*|\d+)(\.\d{1,2})?', prompt)
+    if amount_match:
+        amount = float(amount_match.group().replace(',', ''))
+
+    # Identify transaction type
+    if any(word in prompt_lower for word in ["invest", "capital", "start"]):
+        account_name = "capital"
+        description = "Owner Capital Contribution"
+        debit_account = "cash"
+        credit_account = account_name
+    elif "rent" in prompt_lower:
         account_name = "rent"
-        amount = 500  # Normally extracted from LLM
-    else:
+        description = "Rent Expense"
+        debit_account = account_name
+        credit_account = "cash"
+    elif "utilities" in prompt_lower:
         account_name = "utilities"
-        amount = 300
-
-    expense_account = account_map[account_name]
-    cash_account = account_map["cash"]
+        description = "Utilities Expense"
+        debit_account = account_name
+        credit_account = "cash"
+    elif any(word in prompt_lower for word in ["sale", "revenue"]):
+        account_name = "sales"
+        description = "Sales Revenue"
+        debit_account = "cash"
+        credit_account = account_name
+    elif "supplies" in prompt_lower:
+        account_name = "supplies"
+        description = "Supplies Purchase"
+        debit_account = account_name
+        credit_account = "cash"
+    elif "salary" in prompt_lower or "payroll" in prompt_lower:
+        account_name = "salary"
+        description = "Salary Expense"
+        debit_account = account_name
+        credit_account = "cash"
+    else:
+        description = "Unrecognized Entry"
+        return pd.DataFrame([{"Date": "2025-04-01", "Description": description, "Account Code": "N/A", "Debit": 0, "Credit": 0}])
 
-    expense_account_code = f"{segment['company']}-{segment['business_type']}-{segment['location']}-{segment['cost_center']}-{expense_account}-{segment['future']}"
-    cash_account_code = f"{segment['company']}-{segment['business_type']}-{segment['location']}-{segment['cost_center']}-{cash_account}-{segment['future']}"
+    debit_code = f"{segment['company']}-{segment['business_type']}-{segment['location']}-{segment['cost_center']}-{account_map[debit_account]}-{segment['future']}"
+    credit_code = f"{segment['company']}-{segment['business_type']}-{segment['location']}-{segment['cost_center']}-{account_map[credit_account]}-{segment['future']}"
 
-    entry = pd.DataFrame([
+    entry = [
         {
             "Date": "2025-04-01",
-            "Description": f"{account_name.title()} Expense",
-            "Account Code": expense_account_code,
+            "Description": description,
+            "Account Code": debit_code,
             "Debit": amount,
             "Credit": 0
         },
         {
             "Date": "2025-04-01",
-            "Description": f"Payment for {account_name}",
-            "Account Code": cash_account_code,
+            "Description": f"Offset for {description.lower()}",
+            "Account Code": credit_code,
             "Debit": 0,
             "Credit": amount
         }
-    ])
-    return entry
+    ]
+    st.session_state.gl_entries.extend(entry)
+    return pd.DataFrame(entry)
 
 # Streamlit UI
 st.title("AI ERP Chat - MVP")
 prompt = st.text_input("Enter your accounting instruction:")
 
+delete_records = st.button("Delete All Records")
+if delete_records:
+    st.session_state.gl_entries = []
+    st.success("All records have been deleted.")
+
 if prompt:
     result = handle_gl_entry(prompt)
-    st.dataframe(result)  # Show result as a DataFrame
+    st.dataframe(result)
+
+# Show saved entries
+if st.session_state.gl_entries:
+    st.subheader("All Recorded Entries")
+    st.dataframe(pd.DataFrame(st.session_state.gl_entries))