Update app.py

app.py

@@ -303,6 +303,8 @@ def run_query_sync(query: str):
         web_agent.run(query, ctx=ctx)
     )
 
+stream_queue = asyncio.Queue()
+
 async def run_query(query: str):
     trace_id = f"agent-run-{uuid.uuid4().hex}"
     try:
@@ -311,29 +313,50 @@ async def run_query(query: str):
             session_id="web-agent-session",
             user_id=ANON_USER_ID,
         ):
-            # We'll use this to accumulate the response for the final trace
-            full_response = ""
+            # Clear the queue before starting
+            while not stream_queue.empty():
+                try:
+                    stream_queue.get_nowait()
+                except:
+                    pass
+            
+            # Add thinking message to the queue
+            await stream_queue.put("🤔 Thinking about your question...\n\n")
+            
+            # Patch the agent's methods to capture tool usage
+            original_call_function = web_agent._call_function
+            
+            async def patched_call_function(function_call):
+                tool_name = function_call.get("name", "unknown tool")
+                await stream_queue.put(f"🔧 Using tool: {tool_name}...\n")
+                result = await original_call_function(function_call)
+                await stream_queue.put(f"📊 Got result from {tool_name}\n")
+                return result
+            
+            # Apply the patch
+            web_agent._call_function = patched_call_function
             
-            # Return a generator that we can iterate through in gradio_query
-            async def response_generator():
-                nonlocal full_response
-                # Use the agent normally, but we'll handle streaming separately
-                result = await web_agent.run(query, ctx=ctx)
-                # Get the final response
-                response_text = result.response if isinstance(result.response, str) else str(result.response)
-                
-                # Simulate streaming by yielding words
-                words = response_text.split()
-                partial = ""
-                for word in words:
-                    partial += word + " "
-                    yield partial
-                    await asyncio.sleep(0.05)  # Simulate typing speed
-                
-                # Store the full response for tracing
-                full_response = response_text
+            # Start the agent run in a separate task
+            run_task = asyncio.create_task(web_agent.run(query, ctx=ctx))
             
-            return response_generator()
+            # Stream from the queue while the agent is running
+            while not run_task.done():
+                try:
+                    chunk = await asyncio.wait_for(stream_queue.get(), timeout=0.5)
+                    yield chunk
+                except asyncio.TimeoutError:
+                    # No new items, just wait
+                    await asyncio.sleep(0.1)
+            
+            # Get the final result
+            result = await run_task
+            final_response = result.response if isinstance(result.response, str) else str(result.response)
+            
+            # Restore the original method
+            web_agent._call_function = original_call_function
+            
+            # Yield the final answer
+            yield f"\n\n✅ Final answer: {final_response}"
     finally:
         instrumentor.flush()
 
@@ -346,13 +369,13 @@ async def gradio_query(user_input, chat_history=None):
     history.append({"role": "assistant", "content": "Thinking..."})
     yield history, history
     
-    # Get response generator
-    response_gen = await run_query(user_input)
-    
-    # Update the response as chunks come in
-    async for chunk in response_gen:
-        history[-1]["content"] = chunk
-        yield history, history
+    # Get streaming response
+    full_response = ""
+    async for chunk in run_query(user_input):
+        if chunk:
+            full_response += chunk
+            history[-1]["content"] = full_response
+            yield history, history
 
 # Build and launch Gradio app
 grb = gr.Blocks()