Update app.py

app.py

@@ -1,4 +1,4 @@
-# app.py (Fixed SSE Parsing for Large JSON Logs)
+# app.py (Fixed: Robust Line-by-Line Streaming)
 import gradio as gr
 import httpx
 import os
@@ -38,30 +38,33 @@ def call_blaxel_backend(user_problem, google_key, anthropic_key, sambanova_key):
     yield {status_output: f"Connecting to MudabbirAI..."}
     
     try:
-        with httpx.stream("POST", full_endpoint_url, json=payload, headers=headers, timeout=300.0) as response:
+        # Increased timeout to 600s because complex agents can take time
+        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}"}
                 return 
 
             final_json = None
             full_log = ""
-            buffer = ""
 
-            # Iterate over the stream
-            for chunk in response.iter_text():
-                buffer += chunk
-                
-                # Process complete SSE events (separated by double newline)
-                while "\n\n" in buffer:
-                    message, buffer = buffer.split("\n\n", 1)
-                    
-                    # Remove "data: " prefix if present
-                    if message.startswith("data: "):
-                        content = message.replace("data: ", "", 1)
-                    else:
-                        content = message
+            # --- FIX: Use iter_lines() for robust SSE handling ---
+            for line in response.iter_lines():
+                if not line: 
+                    continue # Skip empty keep-alive lines
 
-                    # Check if this is the FINAL payload
+                # 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
                     if content.startswith("FINAL:"):
                         try:
                             json_str = content.replace("FINAL:", "", 1)
@@ -69,11 +72,12 @@ 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:
-                        # Standard log update
+                        # Normal log update
+                        # We add a newline to separate log entries cleanly
                         full_log += content + "\n"
                         yield {status_output: full_log}
 
-            # Once stream ends, check if we got the JSON and render visuals
+            # --- Render Results ---
             if final_json:
                 log_data = final_json.get("log")
                 
@@ -89,9 +93,12 @@ def call_blaxel_backend(user_problem, google_key, anthropic_key, sambanova_key):
                     progress_plot: chart,
                     calibration_data: calib_table
                 }
+            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:
-        yield {status_output: f"Error: {str(e)}"}
+        yield {status_output: f"Connection Error: {str(e)}"}
 
 # --- Gradio UI ---
 with gr.Blocks(theme=gr.themes.Soft(primary_hue="emerald", secondary_hue="slate")) as demo: