Update app.py

app.py

@@ -10,7 +10,9 @@ import io
 import base64
 import requests
 import google.generativeai as genai
-from fpdf import FPDF  # New import for PDF generation
+from fpdf import FPDF
+import markdown2
+import re
 
 # API Endpoint and payload
 API_URL = "https://irisplus.elixir.co.zw/public/api/profile/reporting/stock-card/genericReports"
@@ -39,25 +41,102 @@ def fetch_data():
         st.error(f"Error fetching data: {response.status_code} - {response.text}")
         return None
 
-# Configure Gemini API
-gemini_api_key = os.environ.get('GOOGLE_API_KEY')
-if not gemini_api_key:
-    st.error("GOOGLE_API_KEY environment variable not set.")
-    st.stop()
+# Function to render a table in the PDF
+def render_table(pdf, df):
+    pdf.set_font("Arial", "B", 12)
+    # Column Widths
+    col_widths = [40] * len(df.columns)
+    
+    # Header Row
+    for i, col in enumerate(df.columns):
+        pdf.cell(col_widths[i], 10, col, border=1, align="C")
+    pdf.ln()
+    
+    # Data Rows
+    pdf.set_font("Arial", "", 12)
+    for _, row in df.iterrows():
+        for i, col in enumerate(row):
+            pdf.cell(col_widths[i], 10, col, border=1, align="C")
+        pdf.ln()
 
-genai.configure(api_key=gemini_api_key)
+# Function to handle Markdown conversion to PDF
+def md_to_pdf(md_text, pdf):
+    lines = md_text.split("\n")
+    for line in lines:
+        line = line.strip()
 
-generation_config = {
-  "temperature": 0.2,
-  "top_p": 0.95,
-  "max_output_tokens": 5000,
-}
+        # Convert Headings
+        if line.startswith("### "):  # H3
+            pdf.set_font("Arial", "B", 14)
+            pdf.cell(0, 10, line[4:], ln=True)
+        elif line.startswith("## "):  # H2
+            pdf.set_font("Arial", "B", 16)
+            pdf.cell(0, 10, line[3:], ln=True)
+        elif line.startswith("# "):  # H1
+            pdf.set_font("Arial", "B", 18)
+            pdf.cell(0, 10, line[2:], ln=True)
+
+        # Convert Bold Text (**bold** → bold)
+        elif "**" in line:
+            bold_text = re.sub(r"\*\*(.*?)\*\*", r"\1", line)
+            pdf.set_font("Arial", "B", 12)
+            pdf.multi_cell(0, 10, bold_text)
+
+        # Convert Bullet Points
+        elif line.startswith("- "):
+            pdf.set_font("Arial", "", 12)
+            pdf.cell(10)  # Indent bullet points
+            pdf.multi_cell(0, 10, f"• {line[2:]}")
+
+        # Detect and Convert Tables (Markdown tables)
+        elif "|" in line and "-" not in line:  # Skip separator lines
+            table_data.append(line.split("|")[1:-1])  # Remove empty edges
 
-model = genai.GenerativeModel(
-  model_name="gemini-2.0-flash-thinking-exp",
-  generation_config=generation_config,
-)
+        # Normal Text
+        else:
+            pdf.set_font("Arial", "", 12)
+            pdf.multi_cell(0, 10, line)
+
+    pdf.ln(5)  # Add spacing
 
+# Modified generate_pdf function to handle markdown
+def generate_pdf(report_text):
+    """
+    Generate a PDF from Markdown content and return the bytes.
+    """
+    pdf = FPDF()
+    pdf.add_page()
+
+    # Convert Markdown text to PDF
+    md_to_pdf(report_text, pdf)
+
+    # Output PDF to a byte string
+    pdf_bytes = pdf.output(dest="S").encode("latin1")
+    return pdf_bytes
+
+# Extract Tables from Markdown
+def extract_tables(md_text):
+    tables = []
+    lines = md_text.split("\n")
+    table_lines = []
+    
+    for line in lines:
+        if "|" in line:  # Identify table lines
+            table_lines.append(line)
+        elif table_lines:  # Table ended
+            tables.append(table_lines)
+            table_lines = []
+    
+    return tables
+
+# Convert Extracted Markdown Table to DataFrame
+def table_to_dataframe(table_lines):
+    headers = table_lines[0].split("|")[1:-1]  # Extract column headers
+    data = [row.split("|")[1:-1] for row in table_lines[2:]]  # Extract row data
+    df = pd.DataFrame(data, columns=headers)
+    return df
+
+# --- Chat Tab Functions ---
 class StreamLitResponse(ResponseParser):
     def __init__(self, context):
         super().__init__(context)
@@ -173,21 +252,6 @@ def handle_userinput(question, df):
     except Exception as e:
         st.error(f"Error processing input: {e}")
 
-def generate_pdf(report_text):
-    """
-    Generate a PDF in memory from the report text.
-    This function uses FPDF to create a PDF and returns its bytes.
-    """
-    pdf = FPDF()
-    pdf.add_page()
-    pdf.set_font("Arial", size=12)
-    # Split report text into lines for writing
-    for line in report_text.split('\n'):
-        pdf.multi_cell(0, 10, line)
-    # Output PDF to a byte string
-    pdf_bytes = pdf.output(dest="S").encode("latin1")
-    return pdf_bytes
-
 def main():
     st.set_page_config(page_title="AI Chat with Your Data", page_icon="📊")
 
@@ -240,6 +304,8 @@ def main():
                     prompt = f"""
                         "You are an expert business analyst. Analyze the following data and generate a comprehensive and insightful business report 
                         including key performance indicators and recommendations.\n\nData:\n" + {str(filtered_df.to_json(orient='records'))}
+
+                        Instructions: Answer directly, without any introductory remarks, salutations, or extra commentary. Just provide the solution.
                     """
                     response = model.generate_content(prompt)
                     # Display the report text; adjust according to your response format.