Update run.py

run.py

@@ -186,6 +186,8 @@ def run_agent_with_streaming(agent, question, stream_callback=None):
 def create_gradio_interface():
     """Create Gradio interface with streaming support"""
     import gradio as gr
+    import time
+    import threading
     
     def process_question(question, model_id, hf_token, serpapi_key, custom_api_endpoint, 
                         custom_api_key, search_provider, search_api_key, custom_search_url):
@@ -194,9 +196,9 @@ def create_gradio_interface():
         agent = create_agent(
             model_id=model_id,
             hf_token=hf_token,
-            openai_api_key=None,  # Add if needed
+            openai_api_key=None,
             serpapi_key=serpapi_key,
-            api_endpoint=None,  # Add if needed
+            api_endpoint=None,
             custom_api_endpoint=custom_api_endpoint,
             custom_api_key=custom_api_key,
             search_provider=search_provider,
@@ -204,22 +206,40 @@ def create_gradio_interface():
             custom_search_url=custom_search_url,
         )
         
-        # Stream output
-        full_output = []
+        # Shared state for streaming
+        output_buffer = []
+        is_complete = False
         
         def stream_callback(text):
-            full_output.append(text)
-            return "".join(full_output)
+            output_buffer.append(text)
         
-        # Generator function for streaming
-        def generate():
+        def run_agent_async():
+            nonlocal is_complete
             try:
                 answer = run_agent_with_streaming(agent, question, stream_callback)
-                yield "".join(full_output) + f"\n\n**FINAL ANSWER:** {answer}"
+                output_buffer.append(f"\n\n**FINAL ANSWER:** {answer}")
             except Exception as e:
-                yield "".join(full_output) + f"\n\n**ERROR:** {str(e)}"
+                output_buffer.append(f"\n\n**ERROR:** {str(e)}")
+            finally:
+                is_complete = True
         
-        return generate()
+        # Start agent in background thread
+        agent_thread = threading.Thread(target=run_agent_async)
+        agent_thread.start()
+        
+        # Generator that yields updates
+        last_length = 0
+        while not is_complete or agent_thread.is_alive():
+            current_output = "".join(output_buffer)
+            if len(current_output) > last_length:
+                yield current_output
+                last_length = len(current_output)
+            time.sleep(0.1)  # Small delay to prevent excessive updates
+        
+        # Final yield to ensure everything is captured
+        final_output = "".join(output_buffer)
+        if len(final_output) > last_length:
+            yield final_output
     
     # Create Gradio interface
     with gr.Blocks(title="Streaming Agent Chat") as demo: