Update app.py

app.py

@@ -1,10 +1,11 @@
 ##############################################################################
-# Sozo Business Studio · 09-Jul-2025                                        #
-#  • Clean narrator text (no scene labels / chart talk)                      #
-#  • Enforce chart-tag-driven visuals (bar, pie, line, scatter, hist)        #
-#  • Fix image generation (Gemini Flash preview) & placeholder fallback      #
-#  • Animation starts blank; artists returned for blit=True                  #
-#  • Silent-audio fallback keeps mux lengths equal                           #
+# Sozo Business Studio · 10-Jul-2025 (full drop-in)                          #
+#  • Restores PDF branch alongside fixed Video branch                         #
+#  • Shared chart-tag grammar across both paths                               #
+#  • Narrator text cleans scene labels + chart talk                           #
+#  • Matplotlib animation starts from blank; artists returned (blit=True)     #
+#  • Gemini Flash-preview image gen with placeholder fallback                 #
+#  • Silent-audio fallback keeps mux lengths equal                            #
 ##############################################################################
 
 import os, re, json, hashlib, uuid, base64, io, tempfile, requests, subprocess
@@ -26,16 +27,14 @@ import cv2
 from langchain_experimental.agents import create_pandas_dataframe_agent
 from langchain_google_genai import ChatGoogleGenerativeAI
 from google import genai
-from google.genai import types   # GenerateContentConfig for image calls
+from google.genai import types   # for GenerateContentConfig
 
-# ────────────────────────────────────────────────────────────────────────────
-# CONFIG
-# ────────────────────────────────────────────────────────────────────────────
+# ─── CONFIG ────────────────────────────────────────────────────────────────
 st.set_page_config(page_title="Sozo Business Studio", layout="wide")
 st.title("📊 Sozo Business Studio")
 st.caption("AI transforms business data into compelling narratives.")
 
-FPS, WIDTH, HEIGHT = 24, 1280, 720
+FPS, WIDTH, HEIGHT  = 24, 1280, 720
 MAX_CHARTS, VIDEO_SCENES = 5, 5
 
 API_KEY = os.getenv("GEMINI_API_KEY")
@@ -43,14 +42,13 @@ if not API_KEY:
     st.error("⚠️ GEMINI_API_KEY is not set."); st.stop()
 GEM = genai.Client(api_key=API_KEY)
 
-DG_KEY = os.getenv("DEEPGRAM_API_KEY")  # optional narration
+DG_KEY = os.getenv("DEEPGRAM_API_KEY")  # optional for narration
 st.session_state.setdefault("bundle", None)
 sha1_bytes = lambda b: hashlib.sha1(b).hexdigest()
 
-# ────────────────────────────────────────────────────────────────────────────
-# HELPERS
-# ────────────────────────────────────────────────────────────────────────────
+# ─── HELPERS ───────────────────────────────────────────────────────────────
 def load_dataframe_safely(buf: bytes, name: str) -> Tuple[pd.DataFrame, str]:
+    """Load CSV/Excel, return (df, err)."""
     try:
         ext = Path(name).suffix.lower()
         df = (pd.read_excel if ext in (".xlsx", ".xls") else pd.read_csv)(io.BytesIO(buf))
@@ -64,6 +62,7 @@ def load_dataframe_safely(buf: bytes, name: str) -> Tuple[pd.DataFrame, str]:
 
 
 def arrow_df(df: pd.DataFrame) -> pd.DataFrame:
+    """Convert for Streamlit Arrow renderer."""
     safe = df.copy()
     for c in safe.columns:
         if safe[c].dtype.name in ("Int64", "Float64", "Boolean"):
@@ -73,9 +72,10 @@ def arrow_df(df: pd.DataFrame) -> pd.DataFrame:
 
 @st.cache_data(show_spinner=False)
 def deepgram_tts(txt: str) -> Tuple[bytes, str]:
+    """Optional audio narration."""
     if not DG_KEY or not txt:
         return None, None
-    txt = re.sub(r"[^\w\s.,!?;:-]", "", txt)[:1000]  # Deepgram text hygiene
+    txt = re.sub(r"[^\w\s.,!?;:-]", "", txt)[:1000]
     try:
         r = requests.post(
             "https://api.deepgram.com/v1/speak",
@@ -118,31 +118,11 @@ re_scene = re.compile(r"^\s*scene\s*\d+[:.\- ]*", re.I)
 def clean_narration(txt: str) -> str:
     txt = re_scene.sub("", txt)
     txt = TAG_RE.sub("", txt)
-    txt = re.sub(r"\s*\([^)]*\)", "", txt)          # remove parentheticals
+    txt = re.sub(r"\s*\([^)]*\)", "", txt)
     txt = re.sub(r"\s{2,}", " ", txt).strip()
     return txt
 
 
-# ─── PDF GENERATION (unchanged logic) ───────────────────────────────────────
-class PDF(FPDF, HTMLMixin):
-    pass
-
-
-def build_pdf(md: str, charts: Dict[str, str]) -> bytes:
-    html = MarkdownIt("commonmark", {"breaks": True}).enable("table").render(
-        TAG_RE.sub(lambda m: f'<img src="{charts.get(m.group("d").strip(), "")}">', md)
-    )
-    pdf = PDF()
-    pdf.set_auto_page_break(True, margin=15)
-    pdf.add_page()
-    pdf.set_font("Arial", "B", 18)
-    pdf.cell(0, 12, "AI-Generated Business Report", ln=True)
-    pdf.ln(3)
-    pdf.set_font("Arial", "", 11)
-    pdf.write_html(html)
-    return bytes(pdf.output(dest="S"))
-
-
 # ─── IMAGE GENERATION & PLACEHOLDER ────────────────────────────────────────
 def placeholder_img() -> Image.Image:
     return Image.new("RGB", (WIDTH, HEIGHT), (230, 230, 230))
@@ -151,7 +131,7 @@ def placeholder_img() -> Image.Image:
 def generate_image_from_prompt(prompt: str) -> Image.Image:
     model_main = "gemini-2.0-flash-exp-image-generation"
     model_fallback = "gemini-2.0-flash-preview-image-generation"
-    full_prompt = ("A clean business-presentation illustration: " + prompt)
+    full_prompt = "A clean business-presentation illustration: " + prompt
 
     def fetch(model_name):
         res = GEM.models.generate_content(
@@ -171,6 +151,86 @@ def generate_image_from_prompt(prompt: str) -> Image.Image:
         return placeholder_img()
 
 
+# ─── PDF GENERATION ────────────────────────────────────────────────────────
+class PDF(FPDF, HTMLMixin):
+    pass
+
+
+def build_pdf(md: str, charts: Dict[str, str]) -> bytes:
+    html = MarkdownIt("commonmark", {"breaks": True}).enable("table").render(
+        TAG_RE.sub(lambda m: f'<img src="{charts.get(m.group("d").strip(), "")}">', md)
+    )
+    pdf = PDF()
+    pdf.set_auto_page_break(True, margin=15)
+    pdf.add_page()
+    pdf.set_font("Arial", "B", 18)
+    pdf.cell(0, 12, "AI-Generated Business Report", ln=True)
+    pdf.ln(3)
+    pdf.set_font("Arial", "", 11)
+    pdf.write_html(html)
+    return bytes(pdf.output(dest="S"))
+
+
+def generate_report(buf: bytes, name: str, ctx: str, key: str):
+    df, err = load_dataframe_safely(buf, name)
+    if err:
+        st.error(err); return None
+
+    llm = ChatGoogleGenerativeAI(model="gemini-2.0-flash",
+                                 google_api_key=API_KEY, temperature=0.1)
+
+    ctx_dict = {
+        "shape": df.shape,
+        "columns": list(df.columns),
+        "user_ctx": ctx or "General business analysis",
+    }
+    cols = ", ".join(ctx_dict["columns"][:6])
+    report_prompt = (
+        "You are a senior business analyst. Write an executive-level Markdown report "
+        "with insights & recommendations.\n"
+        'When a visual is helpful, insert a tag like <generate_chart: "pie | sales by region"> '
+        "(chart_type first, then a description). Valid chart types: bar, pie, line, scatter, hist.\n"
+        f"Base every chart on columns ({cols}) from the dataset.\n"
+        f"Data context:\n{json.dumps(ctx_dict, indent=2)}"
+    )
+    md = llm.invoke(report_prompt).content
+
+    chart_descs = extract_chart_tags(md)[:MAX_CHARTS]
+    charts: Dict[str, str] = {}
+    if chart_descs:
+        agent = create_pandas_dataframe_agent(
+            llm=llm, df=df, verbose=False, allow_dangerous_code=True
+        )
+        for d in chart_descs:
+            with st.spinner(f"Generating chart: {d}"):
+                with plt.ioff():
+                    try:
+                        agent.run(f"Create a {d} with Matplotlib and save.")
+                        fig = plt.gcf()
+                        if fig.axes:
+                            p = Path(tempfile.gettempdir()) / f"{uuid.uuid4()}.png"
+                            fig.savefig(p, dpi=300, bbox_inches="tight", facecolor="white")
+                            charts[d] = str(p)
+                        plt.close("all")
+                    except Exception:
+                        plt.close("all")
+
+    preview = TAG_RE.sub(
+        lambda m: f'<img src="data:image/png;base64,{base64.b64encode(Path(charts[m.group("d").strip()]).read_bytes()).decode()}">' 
+        if m.group("d").strip() in charts else m.group(0),
+        md
+    )
+    pdf_bytes = build_pdf(md, charts)
+
+    return {
+        "type": "report",
+        "preview": preview,
+        "pdf": pdf_bytes,
+        "report_md": md,
+        "key": key,
+    }
+
+
 # ─── ANIMATION HELPERS ─────────────────────────────────────────────────────
 def animate_image_fade(img_cv2: np.ndarray, dur: float, out: Path, fps: int = FPS) -> str:
     frames = max(int(dur * fps), fps)
@@ -188,7 +248,6 @@ def animate_chart(desc: str, df: pd.DataFrame, dur: float, out: Path, fps: int =
     ctype = ctype or "bar"
     title = rest[0] if rest else desc
 
-    # aggregate or prepare data
     if ctype == "pie":
         cat = df.select_dtypes(exclude="number").columns[0]
         num = df.select_dtypes(include="number").columns[0]
@@ -196,7 +255,7 @@ def animate_chart(desc: str, df: pd.DataFrame, dur: float, out: Path, fps: int =
     elif ctype in ("bar", "hist"):
         num = df.select_dtypes(include="number").columns[0]
         pdf = df[num]
-    else:  # line/scatter
+    else:
         cols = df.select_dtypes(include="number").columns[:2]
         pdf = df[list(cols)].sort_index()
 
@@ -207,54 +266,29 @@ def animate_chart(desc: str, df: pd.DataFrame, dur: float, out: Path, fps: int =
         wedges, _ = ax.pie(pdf, labels=pdf.index, startangle=90)
         ax.set_title(title)
 
-        def init():
-            for w in wedges: w.set_alpha(0)
-            return wedges
-
-        def update(i):
-            a = i / frames
-            for w in wedges: w.set_alpha(a)
-            return wedges
+        def init(): [w.set_alpha(0) for w in wedges]; return wedges
+        def update(i): a=i/frames; [w.set_alpha(a) for w in wedges]; return wedges
 
     elif ctype == "bar":
         bars = ax.bar(pdf.index, np.zeros_like(pdf.values), color="#1f77b4")
-        ax.set_ylim(0, pdf.max() * 1.1)
-        ax.set_title(title)
-
-        def init():
-            return bars
+        ax.set_ylim(0, pdf.max() * 1.1); ax.set_title(title)
 
-        def update(i):
-            f = i / frames
-            for b, h in zip(bars, pdf.values):
-                b.set_height(h * f)
-            return bars
+        def init(): return bars
+        def update(i): f=i/frames; [b.set_height(h*f) for b,h in zip(bars,pdf.values)]; return bars
 
     elif ctype == "hist":
         _, _, patches = ax.hist(pdf, bins=20, color="#1f77b4", alpha=0)
         ax.set_title(title)
 
-        def init():
-            for p in patches: p.set_alpha(0)
-            return patches
-
-        def update(i):
-            a = i / frames
-            for p in patches: p.set_alpha(a)
-            return patches
+        def init(): [p.set_alpha(0) for p in patches]; return patches
+        def update(i): a=i/frames; [p.set_alpha(a) for p in patches]; return patches
 
     elif ctype == "scatter":
         pts = ax.scatter(pdf.iloc[:, 0], pdf.iloc[:, 1], s=10, alpha=0)
-        ax.set_title(title)
-        ax.grid(alpha=0.3)
+        ax.set_title(title); ax.grid(alpha=.3)
 
-        def init():
-            pts.set_alpha(0)
-            return [pts]
-
-        def update(i):
-            pts.set_alpha(i / frames)
-            return [pts]
+        def init(): pts.set_alpha(0); return [pts]
+        def update(i): pts.set_alpha(i/frames); return [pts]
 
     else:  # line
         line, = ax.plot([], [], lw=2)
@@ -262,21 +296,13 @@ def animate_chart(desc: str, df: pd.DataFrame, dur: float, out: Path, fps: int =
         y_full = pdf.iloc[:, 1] if pdf.shape[1] > 1 else pdf.iloc[:, 0]
         ax.set_xlim(x_full.min(), x_full.max())
         ax.set_ylim(y_full.min(), y_full.max())
-        ax.set_title(title)
-        ax.grid(alpha=0.3)
+        ax.set_title(title); ax.grid(alpha=.3)
 
-        def init():
-            line.set_data([], [])
-            return [line]
+        def init(): line.set_data([], []); return [line]
+        def update(i): k=max(2,int(len(x_full)*i/frames)); line.set_data(x_full[:k],y_full.iloc[:k]); return [line]
 
-        def update(i):
-            k = max(2, int(len(x_full) * i / frames))
-            line.set_data(x_full[:k], y_full.iloc[:k])
-            return [line]
-
-    anim = FuncAnimation(
-        fig, update, init_func=init, frames=frames,
-        blit=True, interval=1000 / fps)
+    anim = FuncAnimation(fig, update, init_func=init, frames=frames,
+                         blit=True, interval=1000 / fps)
     anim.save(str(out), writer=FFMpegWriter(fps=fps, metadata={'artist': 'Sozo'}), dpi=144)
     plt.close(fig)
     return str(out)
@@ -288,10 +314,9 @@ def safe_chart(desc, df, dur, out):
     except Exception:
         with plt.ioff():
             df.plot(ax=plt.gca())
-            tmp_png = Path(tempfile.gettempdir()) / f"{uuid.uuid4()}.png"
-            plt.savefig(tmp_png, bbox_inches="tight")
-            plt.close()
-        img = cv2.resize(cv2.imread(str(tmp_png)), (WIDTH, HEIGHT))
+            p = Path(tempfile.gettempdir()) / f"{uuid.uuid4()}.png"
+            plt.savefig(p, bbox_inches="tight"); plt.close()
+        img = cv2.resize(cv2.imread(str(p)), (WIDTH, HEIGHT))
         return animate_image_fade(img, dur, out)
 
 
@@ -310,9 +335,7 @@ def concat_media(paths: List[str], out: Path, kind="video"):
     lst.unlink(missing_ok=True)
 
 
-# ────────────────────────────────────────────────────────────────────────────
-# PROMPT HELPERS
-# ────────────────��───────────────────────────────────────────────────────────
+# ─── VIDEO GENERATION ──────────────────────────────────────────────────────
 def build_story_prompt(ctx_dict):
     cols = ", ".join(ctx_dict["columns"][:6])
     return (
@@ -326,9 +349,6 @@ def build_story_prompt(ctx_dict):
     )
 
 
-# ────────────────────────────────────────────────────────────────────────────
-# VIDEO GENERATION
-# ────────────────────────────────────────────────────────────────────────────
 def generate_video(buf: bytes, name: str, ctx: str, key: str):
     try:
         subprocess.run(["ffmpeg", "-version"], check=True, capture_output=True)
@@ -358,28 +378,28 @@ def generate_video(buf: bytes, name: str, ctx: str, key: str):
         descs = extract_chart_tags(sc)
         narrative = clean_narration(sc)
 
-        # ----- audio ---------------------------------------------------------
+        # --- audio ---
         audio_bytes, _ = deepgram_tts(narrative)
-        mp3_path = Path(tempfile.gettempdir()) / f"{uuid.uuid4()}.mp3"
+        mp3 = Path(tempfile.gettempdir()) / f"{uuid.uuid4()}.mp3"
         if audio_bytes:
-            mp3_path.write_bytes(audio_bytes)
-            dur = audio_duration(str(mp3_path))
+            mp3.write_bytes(audio_bytes)
+            dur = audio_duration(str(mp3))
         else:
             dur = 5.0
-            generate_silence_mp3(dur, mp3_path)
-        audio_parts.append(str(mp3_path)); temps.append(mp3_path)
+            generate_silence_mp3(dur, mp3)
+        audio_parts.append(str(mp3)); temps.append(mp3)
 
-        # ----- visual --------------------------------------------------------
-        mp4_path = Path(tempfile.gettempdir()) / f"{uuid.uuid4()}.mp4"
+        # --- visual ---
+        mp4 = Path(tempfile.gettempdir()) / f"{uuid.uuid4()}.mp4"
         if descs:
-            safe_chart(descs[0], df, dur, mp4_path)
+            safe_chart(descs[0], df, dur, mp4)
         else:
             img = generate_image_from_prompt(narrative)
             img_cv = cv2.cvtColor(np.array(img.resize((WIDTH, HEIGHT))), cv2.COLOR_RGB2BGR)
-            animate_image_fade(img_cv, dur, mp4_path)
-        video_parts.append(str(mp4_path)); temps.append(mp4_path)
+            animate_image_fade(img_cv, dur, mp4)
+        video_parts.append(str(mp4)); temps.append(mp4)
 
-    # ----- concatenate -------------------------------------------------------
+    # concat
     silent_vid = Path(tempfile.gettempdir()) / f"{uuid.uuid4()}.mp4"
     concat_media(video_parts, silent_vid, "video")
     audio_mix = Path(tempfile.gettempdir()) / f"{uuid.uuid4()}.mp3"
@@ -397,37 +417,56 @@ def generate_video(buf: bytes, name: str, ctx: str, key: str):
     return str(final_vid)
 
 
-# ────────────────────────────────────────────────────────────────────────────
-# UI
-# ────────────────────────────────────────────────────────────────────────────
+# ─── UI ─────────────────────────────────────────────────────────────────────
+mode = st.radio("Select Output Format:", ["Report (PDF)", "Video Narrative"], horizontal=True)
+
 upl = st.file_uploader("Upload CSV or Excel", type=["csv", "xlsx", "xls"])
 if upl:
-    df_preview, _ = load_dataframe_safely(upl.getvalue(), upl.name)
+    df_prev, _ = load_dataframe_safely(upl.getvalue(), upl.name)
     with st.expander("📊 Data Preview"):
-        st.dataframe(arrow_df(df_preview.head()))
+        st.dataframe(arrow_df(df_prev.head()))
 
 ctx = st.text_area("Business context or specific instructions (optional)")
 
-if st.button("🚀 Generate Video", type="primary", disabled=not upl):
-    key = sha1_bytes(b"".join([upl.getvalue(), ctx.encode()]))
-    st.session_state.bundle = None
+if st.button("🚀 Generate", type="primary", disabled=not upl):
+    key = sha1_bytes(b"".join([upl.getvalue(), mode.encode(), ctx.encode()]))
+
     with st.spinner("Generating…"):
-        path = generate_video(upl.getvalue(), upl.name, ctx, key)
-        if path:
-            st.session_state.bundle = {"video_path": path, "key": key}
+        if mode == "Report (PDF)":
+            st.session_state.bundle = generate_report(upl.getvalue(), upl.name, ctx, key)
+        else:
+            st.session_state.bundle = None
+            path = generate_video(upl.getvalue(), upl.name, ctx, key)
+            if path:
+                st.session_state.bundle = {"type": "video", "video_path": path, "key": key}
     st.rerun()
 
-# ────────────────────────────────────────────────────────────────────────────
-# OUTPUT
-# ────────────────────────────────────────────────────────────────────────────
+# ─── OUTPUT ────────────────────────────────────────────────────────────────
 if bundle := st.session_state.get("bundle"):
-    vp = bundle["video_path"]
-    if Path(vp).exists():
-        with open(vp, "rb") as f:
-            st.video(f.read())
-        with open(vp, "rb") as f:
-            st.download_button("Download Video", f,
-                               f"sozo_narrative_{bundle['key'][:8]}.mp4",
-                               "video/mp4")
-    else:
-        st.error("Video file missing – generation failed.")
+    if bundle["type"] == "report":
+        st.subheader("📄 Generated Report")
+        with st.expander("View Report", expanded=True):
+            st.markdown(bundle["preview"], unsafe_allow_html=True)
+
+        c1, c2 = st.columns(2)
+        with c1:
+            st.download_button("Download PDF", bundle["pdf"],
+                               "business_report.pdf", "application/pdf",
+                               use_container_width=True)
+        with c2:
+            if DG_KEY and st.button("🔊 Narrate Summary", use_container_width=True):
+                txt = re.sub(r"<[^>]+>", "", bundle["report_md"])
+                audio, mime = deepgram_tts(txt)
+                st.audio(audio, format=mime) if audio else st.error("Narration failed.")
+
+    else:  # video
+        st.subheader("🎬 Generated Video Narrative")
+        vp = bundle["video_path"]
+        if Path(vp).exists():
+            with open(vp, "rb") as f:
+                st.video(f.read())
+            with open(vp, "rb") as f:
+                st.download_button("Download Video", f,
+                                   f"sozo_narrative_{bundle['key'][:8]}.mp4", "video/mp4")
+        else:
+            st.error("Video file missing – generation failed.")