Update app.py

app.py

@@ -2,7 +2,7 @@ import re
 import json
 import os
 import time
-from datetime import datetime, date
+from datetime import datetime, date, timedelta
 from io import BytesIO
 
 import pandas as pd
@@ -10,9 +10,10 @@ import streamlit as st
 import google.generativeai as genai
 import pypdf
 from fpdf import FPDF
+from google.api_core import exceptions
 
 # Configure API key for Gemini
-api_key = os.environ['Gemini']
+api_key = os.environ.get('Gemini')
 
 def configure_gemini(api_key):
     genai.configure(api_key=api_key)
@@ -40,7 +41,7 @@ def process_with_gemini(model, text):
     - 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": [
@@ -55,39 +56,54 @@ def process_with_gemini(model, text):
             }
         ]
     }"""
-    response = model.generate_content([prompt, text])
-    time.sleep(5)  # Sleep for 5 seconds to work around rate limit
-    return response.text
+    try:
+        response = model.generate_content([prompt, text])
+        time.sleep(5)  # Sleep for 5 seconds to work around rate limit
+        return response.text
+    except exceptions.ServiceUnavailable as e:
+        if e.response.status_code == 504:
+            st.error("Error generating report: Gemini API timed out (504). Please try reducing the time period for the report.")
+            return None
+        else:
+            raise
 
 # 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 {statement_type} report for the period from {start_date.strftime('%d/%m/%Y')} to {end_date.strftime('%d/%m/%Y')}. Present the report in a standard accounting format relevant to Zimbabwe, but with improved readability and visual appeal.
 
 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.
+Standard Accounting Structure (South Africa 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(5)  # Sleep for 7 seconds to work around rate limit
-    return response.text
+    try:
+        response = model.generate_content([prompt])
+        time.sleep(7)  # Sleep for 7 seconds to work around rate limit
+        return response.text
+    except exceptions.ServiceUnavailable as e:
+        if e.response.status_code == 504:
+            st.error("Error generating report: Gemini API timed out (504). Please try reducing the time period for the report.")
+            st.session_state['last_error'] = "504" # Store the error in session state
+            return None
+        else:
+            raise
 
 # Create a PDF file from the report text
 def create_pdf_report(report_text):
     pdf = FPDF()
     pdf.add_page()
     pdf.set_font("Arial", size=12)
-    
+
     # Split report text into lines and add them to the PDF
     for line in report_text.split('\n'):
         pdf.multi_cell(0, 10, line)
-    
+
     pdf_buffer = BytesIO()
     pdf.output(pdf_buffer)
     pdf_buffer.seek(0)
@@ -96,12 +112,16 @@ def create_pdf_report(report_text):
 def main():
     st.title("Quantitlytix AI")
     st.markdown("*Bank Statement Parser & Financial Report Generator*")
-    
+
+    # Initialize session state for last error
+    if 'last_error' not in st.session_state:
+        st.session_state['last_error'] = None
+
     # Sidebar: Select input type: Bulk PDF or CSV Upload
     input_type = st.sidebar.radio("Select Input Type", ("Bulk Bank Statement Upload", "CSV Upload"))
-    
+
     all_transactions = []
-    
+
     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:
@@ -115,16 +135,20 @@ def main():
                         continue
                     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)
+                        if json_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}.")
+                                continue
+                            json_str = json_response[start_idx:end_idx]
+                            json_str = json_str.replace('```json', '').replace('```', '')
+                            try:
+                                data = json.loads(json_str)
+                                transactions = data.get('transactions', [])
+                                all_transactions.extend(transactions)
+                            except json.JSONDecodeError as e:
+                                st.error(f"Error decoding JSON for {uploaded_file.name}: {e}")
             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")
@@ -140,7 +164,7 @@ def main():
                 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)
@@ -155,7 +179,7 @@ def main():
         st.dataframe(df)
     else:
         st.info("No transactions loaded yet.")
-    
+
     # Financial report generation parameters
     st.write("### Generate Financial Report")
     col1, col2 = st.columns(2)
@@ -163,30 +187,54 @@ def main():
         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, start_date, end_date, statement_type)
-                st.success("Financial report generated!")
-                st.text_area("Financial Report", report_text, height=300)
-                
-                pdf_buffer = create_pdf_report(report_text)
-                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 generating financial report: {str(e)}")
+            # Filter transactions by date
+            filtered_transactions = []
+            for transaction in all_transactions:
+                try:
+                    transaction_date = datetime.strptime(transaction.get('Date'), '%d/%m/%Y').date()
+                    if start_date <= transaction_date <= end_date:
+                        filtered_transactions.append(transaction)
+                except (ValueError, TypeError):
+                    st.warning(f"Could not parse date for transaction: {transaction}")
+                    continue
+
+            if not filtered_transactions:
+                st.warning("No transactions found within the selected date range.")
+            else:
+                try:
+                    model = configure_gemini(api_key)
+                    combined_json = {"transactions": filtered_transactions}
+                    with st.spinner("Generating financial report..."):
+                        report_text = generate_financial_report(model, combined_json, start_date, end_date, statement_type)
+                        if report_text:
+                            st.success("Financial report generated!")
+                            st.text_area("Financial Report", report_text, height=300)
+
+                            pdf_buffer = create_pdf_report(report_text)
+                            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 exceptions.ServiceUnavailable as e:
+                    if e.response.status_code == 504:
+                        st.error("Error generating report: Gemini API timed out (504). Please try reducing the time period for the report.")
+                    else:
+                        st.error(f"Error generating financial report: {str(e)}")
+                except Exception as e:
+                    st.error(f"Error generating financial report: {str(e)}")
+                    if "504" in str(e):
+                        st.info("The Gemini API might be overloaded. Consider reducing the time period for the report.")
+                    elif len(filtered_transactions) > 500: # Example threshold, adjust as needed
+                        st.info("For very large datasets, consider generating reports for smaller time periods (e.g., monthly) and combining them manually if a single comprehensive report fails.")
 
 if __name__ == "__main__":
     main()