Spaces:
Sleeping
Sleeping
| import re | |
| import json | |
| import os | |
| import time | |
| from datetime import datetime, date | |
| from io import BytesIO | |
| import pandas as pd | |
| import streamlit as st | |
| import google.generativeai as genai | |
| import pypdf | |
| from fpdf import FPDF | |
| from google.api_core import exceptions | |
| import markdown | |
| from bs4 import BeautifulSoup | |
| # Configure API key for Gemini - Ensure this is set in your environment variables | |
| api_key = os.getenv('Gemini') | |
| def configure_gemini(api_key): | |
| """ | |
| Configures the Gemini model for transaction extraction as specified by the user. | |
| """ | |
| st.info("Configuring Gemini API for transaction extraction...") | |
| genai.configure(api_key=api_key) | |
| return genai.GenerativeModel('gemini-2.0-flash-thinking-exp') | |
| def configure_gemini1(api_key): | |
| """ | |
| Configures the Gemini model for report generation as specified by the user. | |
| """ | |
| st.info("Configuring Gemini API for report generation...") | |
| genai.configure(api_key=api_key) | |
| return genai.GenerativeModel('gemini-2.5-pro') | |
| def read_pdf_pages(file_obj): | |
| st.info(f"Reading PDF pages from {file_obj.name}...") | |
| file_obj.seek(0) | |
| pdf_reader = pypdf.PdfReader(file_obj) | |
| total_pages = len(pdf_reader.pages) | |
| st.info(f"Found {total_pages} pages in PDF.") | |
| return pdf_reader, total_pages | |
| def extract_page_text(pdf_reader, page_num): | |
| if page_num < len(pdf_reader.pages): | |
| text = pdf_reader.pages[page_num].extract_text() | |
| if not text or not text.strip(): | |
| st.warning(f"Page {page_num + 1} appears to be empty or contains no extractable text.") | |
| return text if text else "" | |
| return "" | |
| 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) | |
| - 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, charts and analysis. | |
| 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" | |
| } | |
| ] | |
| }""" | |
| try: | |
| response = model.generate_content([prompt, text]) | |
| time.sleep(6) | |
| return response.text | |
| except exceptions.GoogleAPICallError as e: | |
| st.error(f"A Google API call error occurred during transaction extraction: {e}") | |
| return None | |
| except Exception as e: | |
| st.error(f"An unexpected error occurred during Gemini transaction extraction: {e}") | |
| return None | |
| def process_pdf_pages(model, pdf_reader, total_pages, progress_callback=None): | |
| all_transactions = [] | |
| for page_num in range(total_pages): | |
| if progress_callback: | |
| progress_callback(page_num / total_pages, f"Processing page {page_num + 1} of {total_pages}") | |
| page_text = extract_page_text(pdf_reader, page_num) | |
| if not page_text.strip(): | |
| continue | |
| json_response = process_with_gemini(model, page_text) | |
| if json_response: | |
| match = re.search(r'\{.*\}', json_response, re.DOTALL) | |
| if not match: | |
| continue | |
| json_str = match.group(0) | |
| try: | |
| data = json.loads(json_str) | |
| transactions = data.get('transactions', []) | |
| if transactions: | |
| all_transactions.extend(transactions) | |
| except json.JSONDecodeError: | |
| continue | |
| return all_transactions | |
| def aggregate_financial_data(transactions: list, statement_type: str): | |
| st.info(f"Performing local financial aggregation for {len(transactions)} transactions...") | |
| if not transactions: | |
| return None | |
| df = pd.DataFrame(transactions) | |
| if 'Amount' not in df.columns: | |
| return None | |
| df['Amount'] = df['Amount'].astype(str).str.replace(r'[^\d.]', '', regex=True) | |
| df['Amount'] = pd.to_numeric(df['Amount'], errors='coerce').fillna(0) | |
| df['Type'] = df['Type'].str.lower() | |
| total_income = df[df['Type'] == 'income']['Amount'].sum() | |
| total_expenses = df[df['Type'] == 'expense']['Amount'].sum() | |
| net_position = total_income - total_expenses | |
| aggregated_data = { | |
| "total_income": round(total_income, 2), | |
| "total_expenses": round(total_expenses, 2), | |
| "net_position": round(net_position, 2), | |
| "transaction_count": len(df) | |
| } | |
| if statement_type == "Income Statement": | |
| aggregated_data["expense_breakdown"] = df[df['Type'] == 'expense'].groupby('Category_of_expense')['Amount'].sum().round(2).to_dict() | |
| aggregated_data["income_breakdown"] = df[df['Type'] == 'income'].groupby('Customer_name')['Amount'].sum().round(2).to_dict() | |
| st.success("Local financial aggregation complete.") | |
| return aggregated_data | |
| def generate_financial_report(model, aggregated_data, start_date, end_date, statement_type): | |
| """ | |
| Generates a financial report using a simplified, high-level prompt that | |
| trusts the model to create the correct structure and avoids using any | |
| Markdown characters that could break rendering. | |
| """ | |
| st.info(f"Preparing to generate {statement_type} with pre-aggregated data...") | |
| prompt = f"""You are an expert financial analyst. Your task is to generate a professional Income Statement in Markdown format using the pre-aggregated JSON data provided below. | |
| Here is the financial data: | |
| {json.dumps(aggregated_data, indent=2)} | |
| Your instructions for the report are: | |
| The main title of the report is "Income Statement". | |
| The reporting period is from {start_date.strftime('%d %B %Y')} to {end_date.strftime('%d %B %Y')}. | |
| The currency is South African Rand (ZAR). | |
| The report must contain sections for Revenue, Operating Expenses, and Net Income or Loss. Each of these sections must be a clear table. | |
| The Revenue section must contain a single table showing only the 'Total Revenue'. Do NOT create a table that itemizes or lists individual income sources. | |
| The report must also include a "Key Highlights" section with bullet points and a final "Summary" paragraph. | |
| Use the provided JSON data for all financial figures. | |
| For the Net Income or Loss table, if the net position is negative, display the amount in parentheses. | |
| Separate the major sections with a horizontal rule. | |
| """ | |
| try: | |
| st.info("Sending request to Gemini for final report formatting...") | |
| response = model.generate_content([prompt]) | |
| st.success("Successfully received formatted financial report from Gemini.") | |
| return response.text | |
| except exceptions.GoogleAPICallError as e: | |
| st.error(f"A Google API call error occurred during report generation: {e}") | |
| return None | |
| except Exception as e: | |
| st.error(f"An unexpected error occurred during Gemini report generation: {e}") | |
| return None | |
| def create_pdf_report(report_text): | |
| """ | |
| Creates a PDF from markdown text. This version includes fixes for both the | |
| ValueError from incorrect int conversion and the table rendering logic. | |
| """ | |
| if not report_text: | |
| st.warning("Report text is empty, skipping PDF generation.") | |
| raise ValueError("Input report_text cannot be empty.") | |
| try: | |
| st.info("Starting PDF generation from markdown report...") | |
| cleaned_md = re.sub(r'```markdown|```', '', report_text, flags=re.MULTILINE).strip() | |
| html_content = markdown.markdown(cleaned_md, extensions=['tables']) | |
| soup = BeautifulSoup(html_content, 'html.parser') | |
| pdf = FPDF() | |
| pdf.set_auto_page_break(auto=True, margin=15) | |
| pdf.set_left_margin(15) | |
| pdf.set_right_margin(15) | |
| pdf.add_page() | |
| for element in soup.find_all(True): | |
| if element.name in ['h1', 'h2', 'h3']: | |
| level = int(element.name[1]) | |
| font_size = {1: 16, 2: 14, 3: 12}.get(level, 11) | |
| pdf.set_font('helvetica', 'B', font_size) | |
| pdf.multi_cell(0, 10, element.get_text().strip()) | |
| pdf.ln(level * 2) | |
| elif element.name == 'p': | |
| pdf.set_font('helvetica', '', 11) | |
| pdf.multi_cell(0, 6, element.get_text().strip()) | |
| pdf.ln(4) | |
| elif element.name == 'i': | |
| pdf.set_font('helvetica', 'I', 11) | |
| pdf.multi_cell(0, 6, element.get_text().strip()) | |
| pdf.ln(4) | |
| elif element.name == 'hr': | |
| pdf.line(pdf.get_x(), pdf.get_y(), pdf.w - pdf.r_margin, pdf.get_y()) | |
| pdf.ln(5) | |
| elif element.name == 'ul': | |
| pdf.ln(2) | |
| for li in element.find_all('li'): | |
| pdf.set_font('helvetica', '', 11) | |
| item_text = li.get_text().strip().replace('•', '-') | |
| pdf.multi_cell(0, 5, f" - {item_text}") | |
| pdf.ln(1) | |
| pdf.ln(4) | |
| elif element.name == 'table': | |
| header = [th.get_text().strip() for th in element.find_all('th')] | |
| rows = [[td.get_text().strip() for td in tr.find_all('td')] for tr in element.find_all('tr')[1:]] | |
| if header: | |
| pdf.set_font('helvetica', 'B', 10) | |
| pdf.set_fill_color(230, 230, 230) | |
| col_widths = [ (pdf.w - pdf.l_margin - pdf.r_margin) * 0.6, (pdf.w - pdf.l_margin - pdf.r_margin) * 0.4 ] | |
| for i, header_text in enumerate(header): | |
| pdf.cell(col_widths[i], 8, header_text, border=1, fill=True, align='C') | |
| pdf.ln() | |
| pdf.set_font('helvetica', '', 10) | |
| for row in rows: | |
| is_total_row = any('Total' in cell for cell in row) | |
| if is_total_row: | |
| pdf.set_font('helvetica', 'B', 10) | |
| if len(row) == len(col_widths): | |
| pdf.cell(col_widths[0], 7, row[0], border=1) | |
| pdf.cell(col_widths[1], 7, row[1], border=1, align='R') | |
| pdf.ln() | |
| if is_total_row: | |
| pdf.set_font('helvetica', '', 10) | |
| pdf.ln(6) | |
| st.info("Content added to PDF. Outputting PDF to buffer...") | |
| pdf_output = pdf.output() | |
| st.success("PDF report generated successfully.") | |
| return BytesIO(pdf_output) | |
| except Exception as e: | |
| st.error(f"Failed to generate PDF: {e}") | |
| st.exception(e) | |
| raise | |
| def main(): | |
| st.title("Quantitlytix AI") | |
| st.markdown("*Bank Statement Parser & Financial Report Generator*") | |
| if 'min_date' not in st.session_state: | |
| st.session_state['min_date'] = date(2024, 1, 1) | |
| if 'max_date' not in st.session_state: | |
| st.session_state['max_date'] = date.today() | |
| if 'transactions' not in st.session_state: | |
| st.session_state['transactions'] = [] | |
| input_type = st.sidebar.radio("Select Input Type", ("Bulk Bank Statement Upload", "CSV Upload")) | |
| 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: | |
| model = configure_gemini(api_key) | |
| progress_bar = st.progress(0) | |
| all_transactions = [] | |
| for i, file in enumerate(uploaded_files): | |
| st.text(f"Processing {file.name}...") | |
| pdf_reader, total_pages = read_pdf_pages(file) | |
| if total_pages > 0: | |
| file_transactions = process_pdf_pages(model, pdf_reader, total_pages) | |
| all_transactions.extend(file_transactions) | |
| progress_bar.progress((i + 1) / len(uploaded_files)) | |
| st.session_state['transactions'] = all_transactions | |
| st.success(f"All PDF files processed. Total transactions collected: {len(st.session_state['transactions'])}.") | |
| elif input_type == "CSV Upload": | |
| uploaded_csv = st.file_uploader("Upload CSV of transactions", type="csv") | |
| if uploaded_csv: | |
| df = pd.read_csv(uploaded_csv) | |
| df = df.loc[:, ~df.columns.str.startswith('Unnamed:')] | |
| st.session_state['transactions'] = df.to_dict(orient='records') | |
| st.success(f"Successfully loaded {len(st.session_state['transactions'])} transactions from CSV.") | |
| if st.session_state['transactions']: | |
| df = pd.DataFrame(st.session_state['transactions']) | |
| # Clean and validate data | |
| df['Date'] = pd.to_datetime(df['Date'], errors='coerce', dayfirst=True) | |
| df.dropna(subset=['Date'], inplace=True) | |
| # Perform the same robust cleaning for Amount as in the aggregation function | |
| if 'Amount' in df.columns: | |
| df['Amount'] = df['Amount'].astype(str).str.replace(r'[^\d.]', '', regex=True) | |
| df['Amount'] = pd.to_numeric(df['Amount'], errors='coerce').fillna(0) | |
| if not df.empty: | |
| st.session_state['min_date'] = df['Date'].min().date() | |
| st.session_state['max_date'] = df['Date'].max().date() | |
| st.write("### Extracted Transactions") | |
| st.dataframe(df.astype(str)) | |
| # --- NEW FEATURE: Download Processed CSV Button --- | |
| st.download_button( | |
| label="Download Processed CSV", | |
| data=df.to_csv(index=False).encode('utf-8'), | |
| file_name='processed_transactions.csv', | |
| mime='text/csv', | |
| ) | |
| # --- END OF NEW FEATURE --- | |
| st.write("### Generate Financial Report") | |
| col1, col2 = st.columns(2) | |
| with col1: | |
| start_date = st.date_input("Start Date", st.session_state['min_date']) | |
| with col2: | |
| end_date = st.date_input("End Date", st.session_state['max_date']) | |
| statement_type = st.selectbox("Select Financial Statement", ["Income Statement"]) | |
| if st.button("Generate Financial Report"): | |
| if not st.session_state['transactions']: | |
| st.error("No transactions available to generate report. Please upload files first.") | |
| else: | |
| # Re-create dataframe from session state to ensure it's fresh for filtering | |
| df_report = pd.DataFrame(st.session_state['transactions']) | |
| df_report['Date'] = pd.to_datetime(df_report['Date'], errors='coerce', dayfirst=True) | |
| mask = (df_report['Date'] >= pd.to_datetime(start_date)) & (df_report['Date'] <= pd.to_datetime(end_date)) | |
| filtered_df = df_report.loc[mask] | |
| if filtered_df.empty: | |
| st.warning("No transactions found within the selected date range.") | |
| else: | |
| st.info(f"Found {len(filtered_df)} transactions within the selected date range.") | |
| filtered_transactions_list = filtered_df.to_dict(orient='records') | |
| try: | |
| with st.spinner("Aggregating financial data locally..."): | |
| aggregated_summary = aggregate_financial_data(filtered_transactions_list, statement_type) | |
| if aggregated_summary: | |
| with st.spinner("Generating formatted report with Gemini..."): | |
| model1 = configure_gemini1(api_key) | |
| report_text = generate_financial_report(model1, aggregated_summary, start_date, end_date, statement_type) | |
| if report_text: | |
| st.success("Financial report generated successfully!") | |
| st.markdown("### Financial Report Preview") | |
| st.markdown(report_text, unsafe_allow_html=True) | |
| pdf_buffer = create_pdf_report(report_text) | |
| st.download_button( | |
| label="Download Financial Report as PDF", | |
| data=pdf_buffer, | |
| file_name=f"{statement_type.replace(' ', '_')}_{datetime.now().strftime('%Y%m%d')}.pdf", | |
| mime="application/pdf" | |
| ) | |
| except Exception as e: | |
| st.error(f"An unexpected error occurred during the report generation process: {e}") | |
| st.exception(e) | |
| if __name__ == "__main__": | |
| main() |