Update app.py

app.py

@@ -7,6 +7,7 @@
 # ────────────────────────────────────────────────────────────────────────────────
 import os, re, json, uuid, tempfile, asyncio
 from pathlib import Path
+import contextlib
 
 import pandas as pd
 import streamlit as st
@@ -14,7 +15,8 @@ import matplotlib
 matplotlib.use("Agg")
 import matplotlib.pyplot as plt
 
-from fpdf import FPDF
+from fpdf import FPDF, HTMLMixin                          # ➜ use HTML engine
+from markdown import markdown                             # ➜ md → html
 
 # Google AI & ADK
 from google import genai
@@ -23,47 +25,33 @@ from google.adk.runners import Runner
 from google.adk.sessions import InMemorySessionService
 from google.genai import types
 
-# LangChain helper for pandas + Gemini
+# LangChain helper
 from langchain_experimental.agents import create_pandas_dataframe_agent
 from langchain_google_genai import ChatGoogleGenerativeAI
 
-# ------------------------------------------------------------------------------
+# ──────────────────────────────────────────────────────────
 # 1️⃣  Environment & font setup
-# ------------------------------------------------------------------------------
+# ──────────────────────────────────────────────────────────
 os.environ["STREAMLIT_CONFIG_DIR"] = tempfile.gettempdir()
 os.environ["MPLCONFIGDIR"]         = tempfile.gettempdir()
 
-FONT_DIR          = Path(__file__).parent          # adjust if fonts in sub-dir
-FONT_NAME         = "NotoSans"
-FONT_REGULAR_TTF  = FONT_DIR / "NotoSans-Regular.ttf"
-FONT_BOLD_TTF     = FONT_DIR / "NotoSans-Bold.ttf"
+HERE                = Path(__file__).parent
+FONT_NAME           = "NotoSans"
+FONT_REGULAR_TTF    = HERE / "NotoSans-Regular.ttf"
+FONT_BOLD_TTF       = HERE / "NotoSans-Bold.ttf"
 
-# ------------------------------------------------------------------------------
-# 2️⃣  Streamlit UI
-# ------------------------------------------------------------------------------
-st.set_page_config(page_title="AI Business Report", layout="wide")
-st.title("📊 AI-Generated Business Report")
+# ──────────────────────────────────────────────────────────
+# 2️⃣  Little helpers
+# ──────────────────────────────────────────────────────────
+class PDF(FPDF, HTMLMixin):
+    """FPDF subclass with HTML support and sensible defaults."""
+    pass
 
-api_key = os.getenv("GEMINI_API_KEY")
-if not api_key:
-    st.error("❌ Environment variable `GEMINI_API_KEY` not found.")
-    st.stop()
 
-# Quick client check
-try:
-    genai.Client(api_key=api_key)
-except Exception as e:
-    st.exception(e)
-    st.stop()
+CHART_TAG_REGEX = r'<generate_chart:\s*"([^"]+)"\s*>'
 
-uploaded   = st.file_uploader("Upload CSV or XLSX dataset", ["csv", "xlsx"])
-user_ctx   = st.text_area("Optional additional business context")
-run_button = st.button("🚀  Generate Report")
 
-# ------------------------------------------------------------------------------
-# 3️⃣  Async helper to drive ADK Runner
-# ------------------------------------------------------------------------------
-async def run_with_runner_async(root_agent, context_data):
+async def run_with_runner_async(root_agent, context):
     svc        = InMemorySessionService()
     app_name   = "business_report_app"
     user_id    = "user1"
@@ -78,24 +66,56 @@ async def run_with_runner_async(root_agent, context_data):
                      session_service=svc)
 
     content = types.Content(role="user",
-                            parts=[types.Part(text=json.dumps(context_data))])
+                            parts=[types.Part(text=json.dumps(context))])
 
     async for event in runner.run_async(user_id=user_id,
                                         session_id=session_id,
                                         new_message=content):
         if event.is_final_response():
             return event.content.parts[0].text
-    return None  # fall-through (shouldn’t happen)
+    return None
+
 
+def add_markdown_fragment(pdf: PDF, md_text: str):
+    """Render a chunk of Markdown into the PDF using HTMLMixin."""
+    html = markdown(md_text, extensions=["tables", "fenced_code"])
+    pdf.write_html(html)
 
 
-# ------------------------------------------------------------------------------
-# 4️⃣  Main execution block
-# ------------------------------------------------------------------------------
+def save_current_figure() -> Path:
+    """Save the current matplotlib figure to a PNG and return the path."""
+    path = Path(tempfile.gettempdir()) / f"{uuid.uuid4()}.png"
+    plt.tight_layout()
+    plt.savefig(path, dpi=300, bbox_inches="tight")
+    return path
+
+
+# ──────────────────────────────────────────────────────────
+# 3️⃣  Streamlit UI
+# ──────────────────────────────────────────────────────────
+st.set_page_config(page_title="📊 AI Business Report", layout="wide")
+st.title("📊 AI-Generated Business Report")
+
+api_key = os.getenv("GEMINI_API_KEY")
+if not api_key:
+    st.error("❌ Environment variable `GEMINI_API_KEY` not found.")
+    st.stop()
+
+# Quick client check
+with contextlib.suppress(Exception):
+    genai.Client(api_key=api_key)
+
+uploaded   = st.file_uploader("Upload CSV or XLSX dataset", ["csv", "xlsx"])
+user_ctx   = st.text_area("Optional additional business context")
+run_button = st.button("🚀 Generate Report")
+
+# ──────────────────────────────────────────────────────────
+# 4️⃣  Main logic
+# ──────────────────────────────────────────────────────────
 if run_button:
-    # 4.1  Load dataset
+    # ›› 4.1 read file
     if not uploaded:
-        st.warning("⚠️  Please upload a dataset first.")
+        st.warning("⚠️ Please upload a dataset first.")
         st.stop()
 
     try:
@@ -107,24 +127,24 @@ if run_button:
         st.error(f"Failed to read file: {e}")
         st.stop()
 
-    # 4.2  Build the report-writing agent
+    # ›› 4.2 build agents
     instruction = """
-You are a senior business analyst. Write an executive-level Markdown report
+You are a senior business analyst. Write an executive-level **Markdown** report
 covering descriptive statistics, key insights, and recommendations.
-Insert placeholder tags for visualisations like:
-    <generate_chart: "bar chart of total_sales by region">
-    """
+
+When you need a visual, insert a placeholder exactly like:
+<generate_chart: "bar chart of total_sales by region">
+"""
     report_agent = LlmAgent(
         name="ReportAgent",
         model="gemini-2.5-flash",
         description="Creates an executive business analysis report in Markdown",
-        instruction=instruction
+        instruction=instruction,
     )
-    root_agent = SequentialAgent(name="ReportPipeline",
-                                 sub_agents=[report_agent])
+    root_agent = SequentialAgent("ReportPipeline", [report_agent])
 
-    # 4.3  Context passed to Gemini
-    context_data = {
+    # ›› 4.3 context
+    ctx = {
         "dataset_info": {
             "shape": df.shape,
             "columns": list(df.columns),
@@ -132,59 +152,69 @@ Insert placeholder tags for visualisations like:
             "missing": df.isna().sum().to_dict(),
         },
         "user_context": user_ctx,
-        "preview": df.head().to_dict()
+        "preview": df.head().to_dict(),
     }
 
-    # 4.4  Run the AI pipeline
-    with st.spinner("🔍  Running AI analysis…"):
-        report_md = asyncio.run(
-            run_with_runner_async(root_agent, context_data)
-        )
+    with st.spinner("🔍 Running AI analysis …"):
+        report_md = asyncio.run(run_with_runner_async(root_agent, ctx))
 
     if not report_md:
-        st.error("⚠️  No response from AI.")
+        st.error("⚠️ No response from AI.")
         st.stop()
 
-    st.subheader("🔖 Draft Report")
-    st.markdown(report_md, unsafe_allow_html=True)
-
-    # 4.5  Parse chart placeholders & generate with pandas agent
-    chart_prompts = re.findall(r'<generate_chart:\s*"([^"]+)"\s*>', report_md)
-    if chart_prompts:
-        st.subheader("📈 Generated Charts")
-        pandas_agent = create_pandas_dataframe_agent(
-            llm=ChatGoogleGenerativeAI(model="gemini-2.0-flash",
-                                       google_api_key=api_key),
-            df=df,
-            verbose=False,
-            allow_dangerous_code=True
-        )
-        for desc in chart_prompts:
-            _ = pandas_agent.run(f"Create a {desc} using matplotlib")
-            fig = plt.gcf()
-            st.pyplot(fig, use_container_width=True)
-            plt.close()
-
-    # 4.6  Build PDF with Noto Sans (Unicode-safe)
-    pdf = FPDF()
-    pdf.set_auto_page_break(auto=True, margin=15)
+    # ›› 4.4 initialise helpers
+    pandas_agent = create_pandas_dataframe_agent(
+        llm=ChatGoogleGenerativeAI(model="gemini-2.0-flash",
+                                   google_api_key=api_key),
+        df=df,
+        verbose=False,
+        allow_dangerous_code=True,
+    )
 
-    pdf.add_font(FONT_NAME, "", str(FONT_REGULAR_TTF), uni=True)
+    pdf = PDF()
+    pdf.set_auto_page_break(True, margin=15)
+    pdf.add_font(FONT_NAME, "",  str(FONT_REGULAR_TTF), uni=True)
     pdf.add_font(FONT_NAME, "B", str(FONT_BOLD_TTF), uni=True)
     pdf.set_fallback_fonts([FONT_NAME])
-
     pdf.add_page()
+
+    # ›› 4.5 write title
     pdf.set_font(FONT_NAME, "B", 16)
     pdf.cell(0, 10, "AI-Generated Business Report", ln=True)
-    pdf.ln(5)
-
-    pdf.set_font(FONT_NAME, "", 11)
-    pdf.multi_cell(0, 5, report_md)
-
-    pdf_bytes = pdf.output(dest="S")          # returns bytearray / bytes
-# If you insist on a bytes object:
-    pdf_bytes = bytes(pdf_bytes)              # no-op when it’s already bytes
-    st.download_button("⬇️  Download PDF",
-                       pdf_bytes,
-                       file_name="business_report.pdf",
-                       mime="application/pdf")
+    pdf.ln(4)
+
+    # ›› 4.6 walk through the Markdown, replacing chart tags in-place
+    segments = re.split(CHART_TAG_REGEX, report_md)
+
+    for idx, segment in enumerate(segments):
+        # even indices → text, odd → chart description
+        if idx % 2 == 0:
+            if segment.strip():
+                st.markdown(segment, unsafe_allow_html=True)
+                add_markdown_fragment(pdf, segment)
+        else:
+            desc = segment.strip()
+            # Generate the chart
+            _ = pandas_agent.run(f"Create a {desc} using matplotlib")
+            fig = plt.gcf()
+            st.pyplot(fig, use_container_width=True)
+
+            # Save & insert into PDF
+            img_path = save_current_figure()
+            pdf.image(str(img_path), w=pdf.epw)
+            pdf.ln(4)
+            plt.close()
+
+    # ›› 4.7 stream PDF to user
+    raw_pdf = pdf.output(dest="S")
+    pdf_bytes = (
+        raw_pdf.encode("latin-1", "replace")      # PyFPDF 1.x
+        if isinstance(raw_pdf, str)    else raw_pdf                              # fpdf2 2.x
+    )
+
+    st.download_button(
+        "⬇️ Download PDF",
+        pdf_bytes,
+        file_name="business_report.pdf",
+        mime="application/pdf",
+    )