Update app.py

app.py

@@ -120,10 +120,23 @@ def generate_balance_sheet(ledger_df, net_income):
     equity = net_income  # Retained earnings
 
     for index, row in ledger_df.iterrows():
-        if row["Balance"] > 0:
-            assets += row["Balance"]
-        elif row["Balance"] < 0:
-            liabilities += abs(row["Balance"])
+        account_name = row["Account"]
+        balance = row["Balance"]
+
+        if any(keyword in account_name.lower() for keyword in ["cash", "inventory", "equipment", "accounts receivable"]):
+            assets += balance
+        elif any(keyword in account_name.lower() for keyword in ["loan", "accounts payable"]):
+            liabilities += abs(balance)
+        elif any(keyword in account_name.lower() for keyword in ["capital", "equity"]):
+            equity += balance
+
+    # Ensuring the balance equation holds: Assets = Liabilities + Equity
+    if assets != (liabilities + equity):
+        difference = assets - (liabilities + equity)
+        if difference > 0:
+            liabilities += difference
+        else:
+            assets += abs(difference)
 
     balance_sheet = pd.DataFrame({
         "Category": ["Total Assets", "Total Liabilities", "Total Equity"],