Update app.py

app.py

@@ -1,9 +1,10 @@
-# app.py (Fixed: Robust Line-by-Line Streaming)
+# app.py (Final Version: Visuals + Financials + Robust Streaming)
 import gradio as gr
 import httpx
 import os
 import json
-from visuals import create_progress_chart, create_calibration_table
+# Import the visualization logic from our new module
+from visuals import create_progress_chart, create_calibration_table, create_cost_summary
 
 # --- Config ---
 BLAXEL_BASE_URL = os.getenv("BLAXEL_BACKEND_URL")
@@ -38,7 +39,7 @@ def call_blaxel_backend(user_problem, google_key, anthropic_key, sambanova_key):
     yield {status_output: f"Connecting to MudabbirAI..."}
     
     try:
-        # Increased timeout to 600s because complex agents can take time
+        # High timeout for complex agentic workflows
         with httpx.stream("POST", full_endpoint_url, json=payload, headers=headers, timeout=600.0) as response:
             if response.status_code != 200:
                 yield {status_output: f"HTTP Error: {response.status_code}"}
@@ -47,24 +48,16 @@ def call_blaxel_backend(user_problem, google_key, anthropic_key, sambanova_key):
             final_json = None
             full_log = ""
 
-            # --- FIX: Use iter_lines() for robust SSE handling ---
+            # Robust SSE Line Parsing
             for line in response.iter_lines():
-                if not line: 
-                    continue # Skip empty keep-alive lines
+                if not line: continue 
+                if hasattr(line, "decode"): line = line.decode("utf-8")
+                if line.startswith(":"): continue # Skip keep-alive comments
 
-                # Decode line (httpx yields strings in newer versions, but safe to check)
-                if hasattr(line, "decode"):
-                    line = line.decode("utf-8")
-
-                # Filter out keep-alive comments (lines starting with :)
-                if line.startswith(":"):
-                    continue
-
-                # Process data lines
                 if line.startswith("data: "):
                     content = line.replace("data: ", "", 1).strip()
                     
-                    # Check for the FINAL payload
+                    # Check for FINAL payload
                     if content.startswith("FINAL:"):
                         try:
                             json_str = content.replace("FINAL:", "", 1)
@@ -72,18 +65,18 @@ def call_blaxel_backend(user_problem, google_key, anthropic_key, sambanova_key):
                         except json.JSONDecodeError as e:
                             full_log += f"\n[Error Parsing Final JSON]: {e}"
                     else:
-                        # Normal log update
-                        # We add a newline to separate log entries cleanly
+                        # Streaming Log Update
                         full_log += content + "\n"
                         yield {status_output: full_log}
 
-            # --- Render Results ---
+            # Final Rendering
             if final_json:
                 log_data = final_json.get("log")
                 
-                # Generate Visuals
+                # Create Visuals & Cost Report
                 chart = create_progress_chart(log_data)
                 calib_table = create_calibration_table(log_data)
+                cost_md = create_cost_summary(log_data)
 
                 yield {
                     status_output: full_log + "\nDone!",
@@ -91,10 +84,10 @@ def call_blaxel_backend(user_problem, google_key, anthropic_key, sambanova_key):
                     final_audio_output: final_json.get("audio"),
                     final_json_log: log_data,
                     progress_plot: chart,
-                    calibration_data: calib_table
+                    calibration_data: calib_table,
+                    cost_display: cost_md
                 }
             else:
-                # If stream ended but no FINAL json, show what we got
                 yield {status_output: full_log + "\n[Stream ended without final payload]"}
 
     except Exception as e:
@@ -125,6 +118,8 @@ with gr.Blocks(theme=gr.themes.Soft(primary_hue="emerald", secondary_hue="slate"
                 progress_plot = gr.Plot(label="Improvement Trajectory (Draft 1 vs. Final)")
             with gr.Row():
                 calibration_data = gr.Dataframe(label="Team Calibration Results (Model Scores)")
+            with gr.Row():
+                cost_display = gr.Markdown(label="Financial Intelligence")
                 
         with gr.TabItem("📝 Final Briefing", id="result_tab"):
             final_text_output = gr.Markdown(label="Strategic Report")
@@ -137,7 +132,7 @@ with gr.Blocks(theme=gr.themes.Soft(primary_hue="emerald", secondary_hue="slate"
     submit_button.click(
         fn=call_blaxel_backend,
         inputs=[problem_input, google_key_input, anthropic_key_input, sambanova_key_input],
-        outputs=[status_output, final_text_output, final_audio_output, final_json_log, progress_plot, calibration_data]
+        outputs=[status_output, final_text_output, final_audio_output, final_json_log, progress_plot, calibration_data, cost_display]
     )
 
 demo.launch()