Update app.py

app.py

@@ -2,7 +2,7 @@ import re
 import json
 import os
 import time
-from datetime import datetime
+from datetime import datetime, date
 from io import BytesIO
 
 import pandas as pd
@@ -32,40 +32,50 @@ def read_pdf(file_obj):
 # Process PDF text with Gemini to extract transactions as JSON
 def process_with_gemini(model, text):
     prompt = """Analyze this bank statement and extract transactions in JSON format with these fields:
-- Date (format DD/MM/YYYY)
-- Description
-- Amount (just the integer value)
-- Type (is 'income' if 'credit amount', else 'expense')
-- Customer Name (Only If Type is 'income' and if no name is extracted write 'general income' and if type is not 'income' write 'expense')
-- City (In address of bank statement)
-
-Return ONLY valid JSON with this structure:
-{
-    "transactions": [
-        {
-            "Date": "string",
-            "Description": "string",
-            "Customer_name": "string",
-            "City": "string",
-            "Amount": number,
-            "Type": "string"
-        }
-    ]
-}"""
+    - Date (format DD/MM/YYYY)
+    - Description
+    - Amount (just the integer value)
+    - Type (is 'income' if 'credit amount', else 'expense')
+    - Customer Name (Only If Type is 'income' and if no name is extracted write 'general income' and if type is not 'income' write 'expense')
+    - City (In address of bank statement)
+    - Category_of_expense (a string, if transaction 'Type' is 'expense' categorize it based on description into: Water and electricity, Salaries and wages, Repairs & Maintenance, Motor vehicle expenses, Projects Expenses, Hardware expenses, Refunds, Accounting fees, Loan interest, Bank charges, Insurance, SARS PAYE UIF, Advertising & Marketing, Logistics and distribution, Fuel, Website hosting fees, Rentals, Subscriptions, Computer internet and Telephone, Staff training, Travel and accommodation, Depreciation, Other expenses. If no category matches, default to ‘Other expenses’. If 'Type' is 'income' set Destination_of_funds to ‘income’.)
+    - ignore opening or closing balances.
+    
+    Return ONLY valid JSON with this structure:
+    {
+        "transactions": [
+            {
+                "Date": "string",
+                "Description": "string",
+                "Customer_name": "string",
+                "City": "string",
+                "Amount": number,
+                "Type": "string",
+                "Category_of_expense": "string"
+            }
+        ]
+    }"""
     response = model.generate_content([prompt, text])
     time.sleep(7)  # Sleep for 7 seconds to work around rate limit
     return response.text
 
-# Generate financial report from aggregated JSON transactions and chosen sections
-def generate_financial_report(model, json_data, report_types):
+# Generate financial report from aggregated JSON transactions and chosen parameters
+def generate_financial_report(model, json_data, start_date, end_date, statement_type):
     prompt = f"""Based on the following transactions JSON data:
 {json.dumps(json_data)}
+For transactions between {start_date.strftime('%d/%m/%Y')} and {end_date.strftime('%d/%m/%Y')}, generate a detailed {statement_type} report presented in a standard accounting format relevant to Zimbabwe, but with improved readability and visual appeal.
 
-Generate a detailed financial report that includes the following sections: {', '.join(report_types)}.
-Ensure that each section is clearly formatted with headings and includes insights and summaries.
-Return the complete report as plain text."""
+Specific Formatting and Content Requirements:
+
+Standard Accounting Structure (Zimbabwe Focus): Organize the {statement_type} according to typical accounting practices followed in South Africa (e.g., for an Income Statement, clearly separate Revenue, Cost of Goods Sold, Gross Profit, Operating Expenses, and Net Income, considering local terminology where applicable). If unsure of specific local variations, adhere to widely accepted international accounting structures.
+Clear Headings and Subheadings: Use distinct and informative headings and subheadings in English to delineate different sections of the report. Ensure these are visually prominent.
+Consistent Formatting: Maintain consistent formatting for monetary values (e.g., using "ZWL" for Zimbabwe Dollar if applicable and discernible from the data, comma separators for thousands), dates, and alignment.
+Totals and Subtotals: Clearly display totals for relevant categories and subtotals where appropriate to provide a clear understanding of the financial performance or position.
+Descriptive Line Items: Use clear and concise descriptions for each transaction or aggregated account based on the provided JSON data.
+Key Insights: Include a brief section (e.g., "Key Highlights" or "Summary") that identifies significant trends, notable figures, or key performance indicators derived from the data within the statement. This should be written in plain, understandable English, potentially highlighting aspects particularly relevant to the economic context of Zimbabwe if discernible from the data.
+Concise Summary: Provide a concluding summary paragraph that encapsulates the overall financial picture presented in the {statement_type}."""
     response = model.generate_content([prompt])
-    time.sleep(7)  # Sleep for 7 seconds to work around rate limit
+    time.sleep(5)  # Sleep for 7 seconds to work around rate limit
     return response.text
 
 # Create a PDF file from the report text
@@ -87,84 +97,96 @@ def main():
     st.title("Quantitlytix AI")
     st.markdown("*Bank Statement Parser & Financial Report Generator*")
     
-    # Allow multiple PDF uploads
-    uploaded_files = st.file_uploader("Upload PDF bank statements", type="pdf", accept_multiple_files=True)
+    # Sidebar: Select input type: Bulk PDF or CSV Upload
+    input_type = st.sidebar.radio("Select Input Type", ("Bulk Bank Statement Upload", "CSV Upload"))
     
-    # Debug: Show number of files uploaded
-    if uploaded_files:
-        st.write(f"{len(uploaded_files)} file(s) uploaded.")
+    all_transactions = []
     
-    if uploaded_files:
-        try:
-            # Initialize the Gemini model
-            model = configure_gemini(api_key)
-            
-            all_transactions = []
-            for uploaded_file in uploaded_files:
-                # Read PDF text directly from the uploaded file
-                pdf_text = read_pdf(uploaded_file)
-                if not pdf_text:
-                    st.warning(f"No text found in {uploaded_file.name}.")
-                    continue
-                
-                with st.spinner(f"Processing {uploaded_file.name}..."):
-                    json_response = process_with_gemini(model, pdf_text)
-                    # Extract valid JSON from the response
-                    start_idx = json_response.find('{')
-                    end_idx = json_response.rfind('}') + 1
-                    if start_idx == -1 or end_idx == -1:
-                        st.warning(f"Invalid JSON response for {uploaded_file.name}.")
+    if input_type == "Bulk Bank Statement Upload":
+        uploaded_files = st.file_uploader("Upload PDF bank statements", type="pdf", accept_multiple_files=True)
+        if uploaded_files:
+            st.write(f"{len(uploaded_files)} PDF file(s) uploaded.")
+            try:
+                model = configure_gemini(api_key)
+                for uploaded_file in uploaded_files:
+                    pdf_text = read_pdf(uploaded_file)
+                    if not pdf_text:
+                        st.warning(f"No text found in {uploaded_file.name}.")
                         continue
-                    json_str = json_response[start_idx:end_idx]
-                    json_str = json_str.replace('```json', '').replace('```', '')
-                    data = json.loads(json_str)
-                    transactions = data.get('transactions', [])
-                    all_transactions.extend(transactions)
-            
-            # Combine transactions into one JSON object
-            combined_json = {"transactions": all_transactions}
-            
-            # Display extracted transactions in a DataFrame if available
-            if all_transactions:
-                df = pd.DataFrame(all_transactions)
-                # Convert amounts to numeric and format
-                df['Amount'] = pd.to_numeric(df['Amount'], errors='coerce')
-                df['Amount'] = df['Amount'].apply(lambda x: f"R {x:,.2f}" if x >= 0 else f"R ({abs(x):,.2f})")
-                df['Date'] = pd.to_datetime(df['Date'], format='%d/%m/%Y', errors='coerce').dt.strftime('%d/%m/%Y')
-                st.success("Extraction complete!")
-                st.write("### Extracted Transactions")
-                st.dataframe(df)
-            else:
-                st.warning("No transactions were extracted from the uploaded files.")
-            
-            # Allow user to select financial report sections
-            st.write("### Generate Financial Report")
-            report_options = st.multiselect(
-                "Select financial report sections to include",
-                ["By Date", "Income Statement", "Cashflow Statement", "Balance Sheet"],
-                default=["By Date", "Income Statement", "Cashflow Statement", "Balance Sheet"]
-            )
-            
-            if st.button("Generate Financial Report"):
+                    with st.spinner(f"Processing {uploaded_file.name}..."):
+                        json_response = process_with_gemini(model, pdf_text)
+                        start_idx = json_response.find('{')
+                        end_idx = json_response.rfind('}') + 1
+                        if start_idx == -1 or end_idx == -1:
+                            st.warning(f"Invalid JSON response for {uploaded_file.name}.")
+                            continue
+                        json_str = json_response[start_idx:end_idx]
+                        json_str = json_str.replace('```json', '').replace('```', '')
+                        data = json.loads(json_str)
+                        transactions = data.get('transactions', [])
+                        all_transactions.extend(transactions)
+            except Exception as e:
+                st.error(f"Error processing PDF documents: {str(e)}")
+                st.error("Please ensure you're using valid bank statement PDFs and a valid API key")
+    elif input_type == "CSV Upload":
+        uploaded_csv = st.file_uploader("Upload CSV of transactions", type="csv")
+        if uploaded_csv:
+            try:
+                df = pd.read_csv(uploaded_csv)
+                st.write("CSV Data Preview:")
+                st.dataframe(df.head())
+                # Convert dataframe to list of transaction dictionaries
+                transactions = df.to_dict(orient='records')
+                all_transactions.extend(transactions)
+            except Exception as e:
+                st.error(f"Error processing CSV file: {str(e)}")
+    
+    # If transactions are loaded, show DataFrame
+    if all_transactions:
+        df = pd.DataFrame(all_transactions)
+        try:
+            df['Amount'] = pd.to_numeric(df['Amount'], errors='coerce')
+            df['Amount'] = df['Amount'].apply(lambda x: f"R {x:,.2f}" if x >= 0 else f"R ({abs(x):,.2f})")
+            df['Date'] = pd.to_datetime(df['Date'], format='%d/%m/%Y', errors='coerce').dt.strftime('%d/%m/%Y')
+        except Exception as e:
+            st.warning("Some data could not be formatted correctly.")
+        st.success("Transactions loaded successfully!")
+        st.write("### Extracted Transactions")
+        st.dataframe(df)
+    else:
+        st.info("No transactions loaded yet.")
+    
+    # Financial report generation parameters
+    st.write("### Generate Financial Report")
+    col1, col2 = st.columns(2)
+    with col1:
+        start_date = st.date_input("Start Date", date(2024, 1, 1))
+    with col2:
+        end_date = st.date_input("End Date", date(2024, 12, 31))
+    
+    statement_type = st.selectbox("Select Financial Statement", ["Income Statement", "Cashflow Statement", "Balance Sheet"])
+    
+    if st.button("Generate Financial Report"):
+        if not all_transactions:
+            st.error("No transactions available to generate report.")
+        else:
+            try:
+                model = configure_gemini(api_key)
+                combined_json = {"transactions": all_transactions}
                 with st.spinner("Generating financial report..."):
-                    report_text = generate_financial_report(model, combined_json, report_options)
+                    report_text = generate_financial_report(model, combined_json, start_date, end_date, statement_type)
                 st.success("Financial report generated!")
                 st.text_area("Financial Report", report_text, height=300)
                 
-                # Create PDF from the report text
                 pdf_buffer = create_pdf_report(report_text)
-                
-                # Provide a download button for the PDF report
                 st.download_button(
                     label="Download Financial Report as PDF",
                     data=pdf_buffer,
                     file_name=f"financial_report_{datetime.now().strftime('%Y%m%d_%H%M%S')}.pdf",
                     mime="application/pdf"
                 )
-            
-        except Exception as e:
-            st.error(f"Error processing documents: {str(e)}")
-            st.error("Please ensure you're using valid bank statement PDFs and a valid API key")
+            except Exception as e:
+                st.error(f"Error generating financial report: {str(e)}")
 
 if __name__ == "__main__":
     main()