Update app.py

app.py

@@ -1,5 +1,5 @@
 ###############################################################################
-# app.py – AI Report & Slides Demo
+# app.py – AI Report & Slides Demo (stable 2025-06-23)
 ###############################################################################
 import os, re, json, uuid, tempfile, asyncio, base64, io, hashlib
 from pathlib import Path
@@ -20,51 +20,55 @@ except ImportError:
 from pptx import Presentation
 from pptx.util import Inches, Pt
 
-# ------------- Google APIs --------------------------------------------------
+# ── Google APIs ──────────────────────────────────────────────────────────────
 from google.oauth2.service_account import Credentials
 from googleapiclient.discovery import build
 
-# ------------- Gemini / ADK -------------------------------------------------
+# ── Gemini & ADK ─────────────────────────────────────────────────────────────
 from google import genai
 from google.genai import types
-from google.adk.agents import LlmAgent, SequentialAgent
-from google.adk.runners import Runner
+from google.adk.agents   import LlmAgent, SequentialAgent
+from google.adk.runners  import Runner
 from google.adk.sessions import InMemorySessionService
 from langchain_experimental.agents import create_pandas_dataframe_agent
-from langchain_google_genai import ChatGoogleGenerativeAI
+from langchain_google_genai       import ChatGoogleGenerativeAI
 
 ###############################################################################
-# 1. Global config & constants
+# 1. Constants & env
 ###############################################################################
-FONT_DIR          = Path(__file__).parent
-FONT_NAME         = "NotoSans"
-FONT_REGULAR_TTF  = FONT_DIR / "NotoSans-Regular.ttf"
-FONT_BOLD_TTF     = FONT_DIR / "NotoSans-Bold.ttf"
+FONT_DIR   = Path(__file__).parent
+FONT_NAME  = "NotoSans"
+FONT_REG   = FONT_DIR / "NotoSans-Regular.ttf"
+FONT_BOLD  = FONT_DIR / "NotoSans-Bold.ttf"
 
-SLIDE_COUNT       = 7
-TTS_MODEL         = "gemini-2.5-flash-preview-tts"
-GOOGLE_CREDS_JSON = os.getenv("GOOGLE_SERVICE_ACCOUNT_JSON")  # optional
+SLIDES      = 7
+TTS_MODEL   = "gemini-2.5-flash-preview-tts"
+CREDS_JSON  = os.getenv("GOOGLE_SERVICE_ACCOUNT_JSON")          # optional
 
 os.environ["STREAMLIT_CONFIG_DIR"] = tempfile.gettempdir()
 os.environ["MPLCONFIGDIR"]         = tempfile.gettempdir()
 
 ###############################################################################
-# 2. Utility helpers
+# 2. Utility helpers (all cached where heavy)
 ###############################################################################
 def _hash_df(df: pd.DataFrame) -> str:
     return hashlib.sha1(pd.util.hash_pandas_object(df, index=True).values).hexdigest()
 
 @st.cache_resource(show_spinner=False)
-def _get_pandas_agent(df_hash, api_key):
-    df = st.session_state["dataframe"]
+def get_pandas_agent(df_hash: str, api_key: str):
     return create_pandas_dataframe_agent(
         llm=ChatGoogleGenerativeAI(model="gemini-2.0-flash",
                                    google_api_key=api_key),
-        df=df, verbose=False, allow_dangerous_code=True)
+        df=st.session_state.df,
+        verbose=False,
+        allow_dangerous_code=True)
+
+def fix_bullets(text: str) -> str:
+    """Convert stray Win-1252 bullet (0x95) → proper Unicode bullet."""
+    return text.replace("\x95", "•")
 
 def generate_tts_audio(client, text):
-    if len(text) > 500:
-        text = text[:500] + "..."
+    text = (text[:500] + "...") if len(text) > 500 else text
     try:
         resp = client.models.generate_content(
             model=TTS_MODEL,
@@ -81,66 +85,87 @@ def generate_tts_audio(client, text):
         st.error(f"TTS error: {e}")
         return None, None
 
-def _convert_pcm_to_wav(raw_bytes: bytes, sample_rate=24000, width=2):
+def pcm_to_wav(raw: bytes, rate=24000, width=2):
     import wave, contextlib
     buf = io.BytesIO()
     with contextlib.closing(wave.open(buf, "wb")) as wf:
-        wf.setnchannels(1)
-        wf.setsampwidth(width)
-        wf.setframerate(sample_rate)
-        wf.writeframes(raw_bytes)
+        wf.setnchannels(1); wf.setsampwidth(width); wf.setframerate(rate); wf.writeframes(raw)
     return buf.getvalue()
 
-def build_pdf(markdown_src: str) -> bytes:
-    """Return final PDF bytes with charts & tables formatted."""
-    # ─ Markdown → HTML
-    if MarkdownIt:
-        html = MarkdownIt("commonmark", {"breaks": True}).enable("table").render(markdown_src)
-    else:
-        html = markdown(markdown_src, extensions=["tables"])
-
-    # ─ Dummy pdf for width calcs
-    _tmp = FPDF(); page_w_mm = _tmp.w - 2 * _tmp.l_margin; px = int(page_w_mm * _tmp.k)
-
-    # ─ Fix tables
-    html = re.sub(r'(<table[^>]*>)', r'\1<font size="8">', html)
-    html = html.replace('</table>', '</font></table>')
-    html = re.sub(r'<table(?![^>]*\bwidth=)', lambda m: f'<table width="{px}"', html)
+@st.cache_resource(show_spinner=False)
+def build_charts(chart_descs: tuple, df_hash: str, api_key: str):
+    """Return {desc:str → path:str} for unique chart descriptions."""
+    agent = get_pandas_agent(df_hash, api_key)
+    desc2path = {}
+    for desc in chart_descs:
+        if desc in desc2path: continue
+        agent.run(f"Create a {desc} using matplotlib")
+        fig = plt.gcf()
+        p = Path(tempfile.gettempdir()) / f"{uuid.uuid4()}.png"
+        fig.savefig(p, dpi=300, bbox_inches="tight"); plt.close()
+        desc2path[desc] = str(p)
+    return desc2path
 
-    # ─ Scale images & clean stray bullets
+@st.cache_resource(show_spinner=False)
+def build_pdf(markdown_src: str) -> bytes:
+    """Return final PDF bytes with fonts, charts, tables fixed."""
+    markdown_src = fix_bullets(markdown_src)
+    # 1 Markdown → HTML
+    html = (MarkdownIt("commonmark", {"breaks": True}).enable("table").render(markdown_src)
+            if MarkdownIt else
+            markdown(markdown_src, extensions=["tables"]))
+    # 2 Dummy pdf for metrics
+    tmp = FPDF(); w_mm, px = tmp.w - 2*tmp.l_margin, lambda mm: int(mm*tmp.k)
+    # 3 Tables
+    html = re.sub(r'(<table[^>]*>)', r'\1<font size="8">', html)\
+           .replace('</table>', '</font></table>')
+    html = re.sub(r'<table(?![^>]*\bwidth=)', lambda m: f'<table width="{px(w_mm)}"', html)
+    # 4 Scale imgs
     html = re.sub(r'<img\s+src="([^"]+\.png)"\s*/?>',
-                  lambda m: f'<img src="{m.group(1)}" width="{int(page_w_mm)}" />', html)
-    html = html.replace("\x95", "•")
-
-    # ─ Build PDF
+                  lambda m: f'<img src="{m.group(1)}" width="{int(w_mm)}" />', html)
+    # 5 PDF build
     class PDF(FPDF, HTMLMixin): pass
     pdf = PDF(); 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.add_font(FONT_NAME, "I", str(FONT_REGULAR_TTF), uni=True)
-    pdf.add_font(FONT_NAME, "BI", str(FONT_BOLD_TTF),   uni=True)
+    for style, path in [("",FONT_REG), ("B",FONT_BOLD), ("I",FONT_REG), ("BI",FONT_BOLD)]:
+        pdf.add_font(FONT_NAME, style, str(path), uni=True)
     pdf.set_fallback_fonts([FONT_NAME])
-
     pdf.add_page(); pdf.set_font(FONT_NAME, "B", 18)
     pdf.cell(0, 12, "AI-Generated Business Report", ln=True); pdf.ln(4)
     pdf.set_font(FONT_NAME, "", 11); pdf.write_html(html)
     return bytes(pdf.output(dest="S"))
 
-def upload_pptx_to_slides(path: Path) -> str | None:
-    if not GOOGLE_CREDS_JSON:
-        st.warning("Service-account JSON env var not set; skipping Slides upload.")
+@st.cache_resource(show_spinner=False)
+def build_pptx(slide_texts: tuple, desc2path: dict) -> Path:
+    prs = Presentation(); blank = prs.slide_layouts[6]
+    for i, txt in enumerate(slide_texts):
+        txt = fix_bullets(txt)
+        slide = prs.slides.add_slide(blank)
+        slide.shapes.add_textbox(Inches(0.6), Inches(0.4), Inches(9), Inches(0.8))\
+             .text_frame.text = f"Slide {i+1}"
+        body = slide.shapes.add_textbox(Inches(0.6), Inches(1.4), Inches(9), Inches(3.2))
+        body.text_frame.text = re.sub(r'<.*?>', '', txt)
+        m = re.search(r'<generate_chart:\s*"([^"]+)"\s*>', txt)
+        if m and m.group(1) in desc2path:
+            slide.shapes.add_picture(desc2path[m.group(1)],
+                                     Inches(1), Inches(3.7), width=Inches(8))
+    out = Path(tempfile.gettempdir()) / f"slides_{uuid.uuid4().hex}.pptx"
+    prs.save(out); return out
+
+def upload_to_slides(pptx_path: Path) -> str | None:
+    if not CREDS_JSON:
+        st.warning("Service-account JSON not configured; skipping Slides upload.")
         return None
     creds = Credentials.from_service_account_file(
-        GOOGLE_CREDS_JSON,
+        CREDS_JSON,
         scopes=["https://www.googleapis.com/auth/drive.file",
                 "https://www.googleapis.com/auth/presentations"])
-    drive = build("drive", "v3", credentials=creds)
-    meta = {"name": f"AI-Slides-{uuid.uuid4().hex[:6]}",
-            "mimeType": "application/vnd.google-apps.presentation"}
+    drive = build("drive",  "v3", credentials=creds)
+    meta  = {"name": f"AI-Slides-{uuid.uuid4().hex[:6]}",
+             "mimeType": "application/vnd.google-apps.presentation"}
     media = {"mimeType": "application/vnd.openxmlformats-officedocument.presentationml.presentation",
-             "body": open(path, "rb")}
-    file = drive.files().create(body=meta, media_body=media).execute()
-    return f"https://docs.google.com/presentation/d/{file['id']}"
+             "body": open(pptx_path, "rb")}
+    f = drive.files().create(body=meta, media_body=media).execute()
+    return f"https://docs.google.com/presentation/d/{f['id']}"
 
 ###############################################################################
 # 3. Streamlit UI
@@ -148,49 +173,46 @@ def upload_pptx_to_slides(path: Path) -> str | None:
 st.set_page_config(page_title="AI Report & Slides", layout="wide")
 st.title("📊 AI-Generated Report & Presentation")
 
-api_key = os.getenv("GEMINI_API_KEY")
-if not api_key:
-    st.error("Set `GEMINI_API_KEY` env var."); st.stop()
-client = genai.Client(api_key=api_key)
+API_KEY = os.getenv("GEMINI_API_KEY")
+if not API_KEY:
+    st.error("Set `GEMINI_API_KEY` in env."); st.stop()
+client = genai.Client(api_key=API_KEY)
 
-MODE = st.radio("Choose output:", ["Report", "Presentation", "Both"], horizontal=True)
+mode = st.radio("Output:", ["Report", "Presentation", "Both"], horizontal=True)
 uploaded = st.file_uploader("Upload CSV / XLSX", ["csv", "xlsx"])
-user_ctx = st.text_area("Optional business context")
-run_btn  = st.button("🚀 Generate")
+user_ctx = st.text_area("Optional business context / objectives")
+run_btn = st.button("🚀 Generate")
 
 ###############################################################################
-# 4. Pipeline
+# 4. Main pipeline
 ###############################################################################
 if run_btn:
-    if not uploaded:
-        st.warning("Upload a dataset."); st.stop()
-
-    # 4.1 Load data
+    if not uploaded: st.warning("Please upload a dataset."); st.stop()
+    # 4.1 Load
     try:
         df = pd.read_excel(uploaded) if uploaded.name.lower().endswith(".xlsx") else pd.read_csv(uploaded)
     except Exception as e:
         st.error(f"Read error: {e}"); st.stop()
-    st.session_state["dataframe"] = df
+    st.session_state.df = df
     df_hash = _hash_df(df)
 
-    # 4.2 Build agents
+    # 4.2 Agents
     report_agent = LlmAgent(
         name="ReportAgent", model="gemini-2.5-flash",
-        instruction="""You are a senior business analyst. Write an executive-level Markdown report with charts tags like <generate_chart: "histogram of sales">.""",
-        description="Writes report")
+        instruction=("You are a senior business analyst. Write an executive-level Markdown "
+                     "report with <generate_chart: \"...\"> placeholders."))
 
-    presentation_agent = LlmAgent(
+    pres_agent = LlmAgent(
         name="PresentationAgent", model="gemini-2.5-flash",
-        instruction=f"""Create exactly {SLIDE_COUNT} presentation slides. Each slide starts with 'Slide X:' then bullet/paragraph. Use <generate_chart: "..."> tags where helpful.""",
-        description="Writes deck")
+        instruction=(f"Create exactly {SLIDES} slides. Each slide starts with 'Slide X:' "
+                     "and may include <generate_chart: \"...\"> tags."))
 
-    sub_agents = []
-    if MODE in ("Report", "Both"): sub_agents.append(report_agent)
-    if MODE in ("Presentation", "Both"): sub_agents.append(presentation_agent)
+    subs = []
+    if mode in ("Report", "Both"): subs.append(report_agent)
+    if mode in ("Presentation", "Both"): subs.append(pres_agent)
+    pipeline = SequentialAgent(name="Pipeline", sub_agents=subs)
 
-    root_agent = SequentialAgent(name="Pipeline", sub_agents=sub_agents)
-
-    # 4.3 Context
+    # 4.3 Context → Gemini
     ctx = {"dataset_info": {"shape": df.shape,
                             "columns": list(df.columns),
                             "dtypes": df.dtypes.astype(str).to_dict(),
@@ -198,117 +220,96 @@ if run_btn:
            "user_context": user_ctx,
            "preview": df.head().to_dict()}
 
-    # 4.4 Run ADK pipeline
-    svc = InMemorySessionService()
-    sid = str(uuid.uuid4())
-    asyncio.run(svc.create_session(app_name="app", user_id="user1", session_id=sid))
-    runner = Runner(agent=root_agent, app_name="app", session_service=svc)
+    # 4.4 Run ADK
+    svc, sid = InMemorySessionService(), str(uuid.uuid4())
+    asyncio.run(svc.create_session(app_name="app", user_id="user", session_id=sid))
+    runner = Runner(agent=pipeline, app_name="app", session_service=svc)
 
-    final_texts = {}
-    async def _run():
+    final = {}
+    async def _collect():
         async for ev in runner.run_async(
-                user_id="user1",
-                session_id=sid,
+                user_id="user", session_id=sid,
                 new_message=types.Content(role="user",
                                           parts=[types.Part(text=json.dumps(ctx))])):
             if ev.is_final_response():
-                final_texts[ev.author] = ev.content.parts[0].text
-    asyncio.run(_run())
+                final[ev.author] = ev.content.parts[0].text
+    asyncio.run(_collect())
 
-    # 4.5 Build charts
-    chart_tags = re.findall(r'<generate_chart:\s*"([^"]+)"\s*>', "\n".join(final_texts.values()))
-    desc2path = {}
-    if chart_tags:
-        pandas_agent = _get_pandas_agent(df_hash, api_key)
-        for desc in chart_tags:
-            if desc in desc2path: continue
-            pandas_agent.run(f"Create a {desc} using matplotlib")
-            fig = plt.gcf(); p = Path(tempfile.gettempdir()) / f"{uuid.uuid4()}.png"
-            fig.savefig(p, dpi=300, bbox_inches="tight"); plt.close()
-            desc2path[desc] = str(p)
+    # 4.5 Charts
+    chart_tags = re.findall(r'<generate_chart:\s*"([^"]+)"\s*>', "\n".join(final.values()))
+    desc2path = build_charts(tuple(chart_tags), df_hash, API_KEY)
 
     # 4.6 Tabs
-    if MODE == "Both":
+    if mode == "Both":
         tab_report, tab_pres = st.tabs(["📄 Report", "📑 Presentation"])
     else:
         tab_report = tab_pres = st
 
     #######################################################################
-    # Report tab
+    # Report
     #######################################################################
-    if MODE in ("Report", "Both"):
-        report_md = final_texts["ReportAgent"]
-        for d, pth in desc2path.items():
-            report_md = re.sub(rf'<generate_chart:\s*"{re.escape(d)}"\s*>',
-                               f'<img src="{pth}" />', report_md)
+    if mode in ("Report", "Both"):
+        md = fix_bullets(final["ReportAgent"])
+        # Replace tags: base-64 for preview + file path for PDF
+        preview_md = md; pdf_md = md
+        for d,pth in desc2path.items():
+            with open(pth, "rb") as f:
+                b64 = base64.b64encode(f.read()).decode()
+            embed = f'<img src="data:image/png;base64,{b64}" style="max-width:100%;margin:8px 0;" />'
+            preview_md = re.sub(rf'<generate_chart:\s*"{re.escape(d)}"\s*>', embed, preview_md)
+            pdf_md     = re.sub(rf'<generate_chart:\s*"{re.escape(d)}"\s*>',
+                                f'<img src="{pth}" />', pdf_md)
 
         with tab_report:
             st.markdown("### Report Preview")
-            st.markdown(report_md, unsafe_allow_html=True)
+            st.markdown(preview_md, unsafe_allow_html=True)
             st.download_button("⬇️ Download PDF",
-                               build_pdf(report_md),
-                               file_name="business_report.pdf",
+                               build_pdf(pdf_md),
+                               "business_report.pdf",
                                mime="application/pdf")
 
     #######################################################################
-    # Presentation tab
+    # Presentation
     #######################################################################
-    if MODE in ("Presentation", "Both"):
-        raw = final_texts["PresentationAgent"]
-        slides_text = [s.strip() for s in re.split(r'\bSlide\s+\d+:', raw) if s.strip()][:SLIDE_COUNT]
-
-        st.session_state.app_state = {"steps": slides_text, "current": 0}
-
-        def render_step(idx, text):
-            tot = len(slides_text)
-            st.markdown(f"#### Slide {idx+1} / {tot}")
-            st.write(text)
-            if st.button(f"🔊 Narrate", key=f"tts_{idx}"):
+    if mode in ("Presentation", "Both"):
+        raw = fix_bullets(final["PresentationAgent"])
+        slides_txt = [s.strip() for s in re.split(r'\bSlide\s+\d+:', raw) if s.strip()][:SLIDES]
+        st.session_state.slides = slides_txt
+        st.session_state.curr   = 0
+
+        def render(idx: int):
+            txt = st.session_state.slides[idx]
+            st.markdown(f"#### Slide {idx+1} / {len(slides_txt)}")
+            st.write(txt)
+            if st.button("🔊 Narrate", key=f"tts_{idx}"):
                 with st.spinner("TTS…"):
-                    aud, mime = generate_tts_audio(client, text)
+                    aud,mime = generate_tts_audio(client, txt)
                     if aud:
                         if 'pcm' in mime or 'L16' in mime:
-                            st.audio(_convert_pcm_to_wav(aud), format="audio/wav")
+                            st.audio(pcm_to_wav(aud), format="audio/wav")
                         else:
                             st.audio(aud, format=mime)
 
         with tab_pres:
-            col_l, col_r = st.columns([1,1])
-            if col_l.button("⬅️ Back", disabled=st.session_state.app_state['current']==0):
-                st.session_state.app_state['current'] -= 1
-            if col_r.button("Next ➡️",
-                            disabled=st.session_state.app_state['current']==len(slides_text)-1):
-                st.session_state.app_state['current'] += 1
-            render_step(st.session_state.app_state['current'],
-                        slides_text[st.session_state.app_state['current']])
-
-            # -- PPTX build (cached) ----------------------------------------
-            @st.cache_resource(show_spinner=False)
-            def _build_pptx(txts, d2p):
-                prs = Presentation(); blank = prs.slide_layouts[6]
-                for i, txt in enumerate(txts):
-                    sl = prs.slides.add_slide(blank)
-                    title = sl.shapes.add_textbox(Inches(0.5), Inches(0.4),
-                                                  Inches(9), Inches(1))
-                    title.text_frame.text = f"Slide {i+1}"
-                    body = sl.shapes.add_textbox(Inches(0.5), Inches(1.4),
-                                                 Inches(9), Inches(3.5))
-                    body.text_frame.text = re.sub(r'<.*?>', '', txt)
-                    m = re.search(r'<generate_chart:\s*"([^"]+)"\s*>', txt)
-                    if m and m.group(1) in d2p:
-                        sl.shapes.add_picture(d2p[m.group(1)],
-                                              Inches(1), Inches(3.5), width=Inches(8))
-                out = Path(tempfile.gettempdir()) / f"slides_{uuid.uuid4().hex}.pptx"
-                prs.save(out); return out
-            pptx_path = _build_pptx(slides_text, desc2path)
+            col1,col2 = st.columns(2)
+            if col1.button("⬅️ Back", disabled=st.session_state.curr==0):
+                st.session_state.curr -= 1
+            if col2.button("Next ➡️",
+                           disabled=st.session_state.curr==len(slides_txt)-1):
+                st.session_state.curr += 1
+            render(st.session_state.curr)
+
+            # PPTX build (cached)
+            pptx = build_pptx(tuple(slides_txt), desc2path)
             st.download_button("⬇️ Download PPTX",
-                               pptx_path.read_bytes(),
-                               file_name="presentation.pptx",
+                               pptx.read_bytes(),
+                               "presentation.pptx",
                                mime=("application/vnd.openxmlformats-officedocument."
                                      "presentationml.presentation"))
             if st.button("⬆️ Upload to Google Slides"):
-                url = upload_pptx_to_slides(pptx_path)
-                if url: st.success(f"[Open Slides]({url})")
+                url = upload_to_slides(pptx)
+                if url:
+                    st.success(f"[Open Slides]({url})")
 ###############################################################################
-# End of file
+# End
 ###############################################################################