Update app.py

app.py

@@ -1,13 +1,5 @@
-# ────────────────────────────────────────────────────────────────────────────────
-# app.py  –  Streamlit AI Business-Report demo (Google Gemini + ADK)
-# Requires: fpdf2 >= 2.7, streamlit, pandas, matplotlib, google-adk, 
-#           langchain-google-genai, langchain-experimental
-# Fonts:    NotoSans-Regular.ttf & NotoSans-Bold.ttf are expected to live in the
-#           same folder as this file (or adjust FONT_DIR below).
-# ────────────────────────────────────────────────────────────────────────────────
 import os, re, json, uuid, tempfile, asyncio
 from pathlib import Path
-import contextlib
 
 import pandas as pd
 import streamlit as st
@@ -15,43 +7,52 @@ import matplotlib
 matplotlib.use("Agg")
 import matplotlib.pyplot as plt
 
-from fpdf import FPDF, HTMLMixin                          # ➜ use HTML engine
-from markdown import markdown                             # ➜ md → html
+from fpdf import FPDF, HTMLMixin
+from markdown import markdown              # NEW
+from markdown_it import MarkdownIt         # prettier HTML (optional)
 
-# Google AI & ADK
-from google import genai
-from google.adk.agents import LlmAgent, SequentialAgent
-from google.adk.runners import Runner
-from google.adk.sessions import InMemorySessionService
-from google.genai import types
-
-# 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()
 
-HERE                = Path(__file__).parent
-FONT_NAME           = "NotoSans"
-FONT_REGULAR_TTF    = HERE / "NotoSans-Regular.ttf"
-FONT_BOLD_TTF       = HERE / "NotoSans-Bold.ttf"
+FONT_DIR          = Path(__file__).parent
+FONT_NAME         = "NotoSans"
+FONT_REGULAR_TTF  = FONT_DIR / "NotoSans-Regular.ttf"
+FONT_BOLD_TTF     = FONT_DIR / "NotoSans-Bold.ttf"
 
-# ──────────────────────────────────────────────────────────
-# 2️⃣  Little helpers
-# ──────────────────────────────────────────────────────────
-class PDF(FPDF, HTMLMixin):
-    """FPDF subclass with HTML support and sensible defaults."""
-    pass
+# ──────────────────────────────────────────────────────────────────────────────
+# 2️⃣  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()
 
-CHART_TAG_REGEX = r'<generate_chart:\s*"([^"]+)"\s*>'
+# Quick Gemini client check
+from google import genai
+try:
+    genai.Client(api_key=api_key)
+except Exception as e:
+    st.exception(e)
+    st.stop()
 
+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            (unchanged)
+# ──���───────────────────────────────────────────────────────────────────────────
+async def run_with_runner_async(root_agent, context_data):
+    from google.adk.sessions import InMemorySessionService
+    from google.adk.runners  import Runner
+    from google.genai import types
 
-async def run_with_runner_async(root_agent, context):
     svc        = InMemorySessionService()
     app_name   = "business_report_app"
     user_id    = "user1"
@@ -66,56 +67,22 @@ async def run_with_runner_async(root_agent, context):
                      session_service=svc)
 
     content = types.Content(role="user",
-                            parts=[types.Part(text=json.dumps(context))])
+                            parts=[types.Part(text=json.dumps(context_data))])
 
     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
-
-
-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)
-
-
-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)
+    return None        # should not occur
 
-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
-# ──────────────────────────────────────────────────────────
+# ──────────────────────────────────────────────────────────────────────────────
+# 4️⃣  Main execution block
+# ──────────────────────────────────────────────────────────────────────────────
 if run_button:
-    # ›› 4.1 read file
+    # 4.1  Load dataset
     if not uploaded:
-        st.warning("⚠️ Please upload a dataset first.")
+        st.warning("⚠️  Please upload a dataset first.")
         st.stop()
 
     try:
@@ -127,24 +94,25 @@ if run_button:
         st.error(f"Failed to read file: {e}")
         st.stop()
 
-    # ›› 4.2 build agents
+    # 4.2  Build the report-writing agent            (unchanged)
+    from google.adk.agents import LlmAgent, SequentialAgent
     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.
-
-When you need a visual, insert a placeholder exactly like:
-<generate_chart: "bar chart of total_sales by region">
+Insert placeholder tags for visualisations 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("ReportPipeline", [report_agent])
+    root_agent = SequentialAgent(name="ReportPipeline",
+                                 sub_agents=[report_agent])
 
-    # ›› 4.3 context
-    ctx = {
+    # 4.3  Context passed to Gemini
+    context_data = {
         "dataset_info": {
             "shape": df.shape,
             "columns": list(df.columns),
@@ -152,69 +120,100 @@ When you need a visual, insert a placeholder exactly like:
             "missing": df.isna().sum().to_dict(),
         },
         "user_context": user_ctx,
-        "preview": df.head().to_dict(),
+        "preview": df.head().to_dict()
     }
 
-    with st.spinner("🔍 Running AI analysis …"):
-        report_md = asyncio.run(run_with_runner_async(root_agent, ctx))
+    # 4.4  Run the AI pipeline
+    with st.spinner("🔍  Running AI analysis…"):
+        report_md = asyncio.run(
+            run_with_runner_async(root_agent, context_data)
+        )
 
     if not report_md:
-        st.error("⚠️ No response from AI.")
+        st.error("⚠️  No response from AI.")
         st.stop()
 
-    # ›› 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,
-    )
+    # 4.5  Parse chart placeholders & generate charts         ★ REWORKED ★
+    chart_tags    = re.findall(r'<generate_chart:\s*"([^"]+)"\s*>', report_md)
+    tag_to_path   = {}                    # desc → temp PNG path
+    chart_html_snippets = {}              # desc → <img …>
+
+    if chart_tags:
+        from langchain_experimental.agents import create_pandas_dataframe_agent
+        from langchain_google_genai import ChatGoogleGenerativeAI
+
+        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
+        )
+        st.subheader("📈 Generated Charts")
+
+        for desc in chart_tags:
+            # LLM draws the chart with matplotlib
+            _ = pandas_agent.run(f"Create a {desc} using matplotlib")
+            fig = plt.gcf()
+
+            tmp_path = Path(tempfile.gettempdir()) / f"{uuid.uuid4()}.png"
+            fig.savefig(tmp_path, dpi=300, bbox_inches="tight")
+            plt.close()
+
+            tag_to_path[desc] = str(tmp_path)
+            chart_html_snippets[desc] = (
+                f'<img src="{tmp_path}" width="100%" style="margin:10px 0;" />'
+            )
+
+            # Immediate visual feedback in Streamlit
+            st.image(tmp_path, caption=desc, use_column_width=True)
+
+    # 4.6  Replace tags with <img> HTML
+    for desc, html_img in chart_html_snippets.items():
+        report_md = re.sub(
+            rf'<generate_chart:\s*"{re.escape(desc)}"\s*>',
+            html_img,
+            report_md
+        )
+
+    # 4.7  Show nicely-rendered preview in Streamlit
+    st.subheader("🔖 Draft Report")
+    st.markdown(report_md, unsafe_allow_html=True)
+
+    # 4.8  Build PDF with HTML rendering (Markdown → HTML)   ★ NEW SECTION ★
+    class PDF(FPDF, HTMLMixin):
+        pass
 
     pdf = PDF()
-    pdf.set_auto_page_break(True, margin=15)
-    pdf.add_font(FONT_NAME, "",  str(FONT_REGULAR_TTF), uni=True)
+    pdf.set_auto_page_break(auto=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.set_font(FONT_NAME, "B", 18)
+    pdf.cell(0, 12, "AI-Generated Business Report", ln=True)
     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",
-    )
+    # Markdown → HTML
+    # (use markdown-it-py if available – better table support)
+    if 'MarkdownIt' in globals():
+        md2html = MarkdownIt("commonmark", {"breaks": True}).enable("table")
+        html_body = md2html.render(report_md)
+    else:
+        html_body = markdown(report_md, extensions=["tables"])
+
+    # For safety: replace bullets “•” with unicode bullet char if markdown() escaped it
+    html_body = html_body.replace("&bull;", "•")
+
+    pdf.set_font(FONT_NAME, "", 11)
+    pdf.write_html(html_body)
+
+    # 4.9  Produce bytes without re-encoding to latin-1
+    pdf_bytes = pdf.output(dest="S")
+    if not isinstance(pdf_bytes, (bytes, bytearray)):
+        pdf_bytes = bytes(pdf_bytes, "utf-8")      # never latin-1!
+
+    st.download_button("⬇️  Download PDF",
+                       pdf_bytes,
+                       file_name="business_report.pdf",
+                       mime="application/pdf")