Update app.py

app.py

@@ -23,6 +23,9 @@ from llama_index.core.memory import ChatMemoryBuffer
 from llama_index.readers.web import RssReader, SimpleWebPageReader
 from llama_index.core import SummaryIndex
 
+# Import the event types for streaming
+from llama_index.core.agent.workflow import AgentStream, ToolCall, ToolCallResult
+
 import subprocess
 subprocess.run(["playwright", "install"])
 
@@ -296,15 +299,14 @@ web_agent = AgentWorkflow.from_tools_or_functions(
 )
 ctx = Context(web_agent)
 
-# Async helper to run agent queries
+# Async helper to run agent queries (kept for compatibility)
 def run_query_sync(query: str):
     """Helper to run async agent.run in sync context."""
     return asyncio.get_event_loop().run_until_complete(
         web_agent.run(query, ctx=ctx)
     )
 
-stream_queue = asyncio.Queue()
-
+# Updated run_query function to use stream_events
 async def run_query(query: str):
     trace_id = f"agent-run-{uuid.uuid4().hex}"
     try:
@@ -313,90 +315,51 @@ async def run_query(query: str):
             session_id="web-agent-session",
             user_id=ANON_USER_ID,
         ):
-            # Clear the queue before starting
-            while not stream_queue.empty():
-                try:
-                    stream_queue.get_nowait()
-                except:
-                    pass
-            
-            # Add initial messages to the queue
-            await stream_queue.put("🤔 Thinking about your question...\n\n")
-            
-            # The key is to patch each individual tool function to capture its usage
-            original_functions = {}
+            # Start the handler
+            handler = web_agent.run(query, ctx=ctx)
             
-            # Store original functions and patch each tool
-            for tool in tools:
-                tool_name = tool.metadata.name
-                original_fn = tool.fn
-                original_functions[tool_name] = original_fn
-                
-                # Create a wrapper function that will log the tool usage
-                def create_wrapper(orig_fn, tool_name):
-                    async def wrapper(*args, **kwargs):
-                        # Log tool usage
-                        await stream_queue.put(f"🔧 Using tool: {tool_name}...\n")
-                        
-                        # Call original function
-                        if asyncio.iscoroutinefunction(orig_fn):
-                            result = await orig_fn(*args, **kwargs)
-                        else:
-                            result = orig_fn(*args, **kwargs)
-                        
-                        # Log result
-                        await stream_queue.put(f"📊 Got result from {tool_name}\n")
-                        return result
+            # Stream content
+            full_response = ""
+            async for event in handler.stream_events():
+                if isinstance(event, AgentStream):
+                    # This is the text being generated
+                    if event.delta:
+                        yield event.delta
+                elif isinstance(event, ToolCall):
+                    # A tool is being called
+                    tool_name = event.data.get("name", "unknown tool")
+                    yield f"\n\n🔧 Using tool: {tool_name}...\n"
+                elif isinstance(event, ToolCallResult):
+                    # Result from a tool call
+                    # Try to find the tool name in the event data
+                    tool_name = "the tool"
+                    if hasattr(event, "parent_id") and event.parent_id:
+                        # This might help identify the parent tool call
+                        tool_name = f"tool {event.parent_id}"
                     
-                    return wrapper
-                
-                # Replace the function with our wrapped version
-                tool.fn = create_wrapper(original_fn, tool_name)
-            
-            # Start the agent run
-            await stream_queue.put("🧠 Planning approach...\n\n")
-            task = asyncio.create_task(web_agent.run(query, ctx=ctx))
-            
-            # Stream updates while waiting for completion
-            while not task.done():
-                try:
-                    # Check if there's anything in the queue to yield
-                    if not stream_queue.empty():
-                        chunk = await stream_queue.get()
-                        yield chunk
+                    # Get the result from the event data
+                    result = event.data.get("result", "")
+                    
+                    # Truncate long results for display
+                    if isinstance(result, str) and len(result) > 200:
+                        result_preview = result[:200] + "... (truncated)"
                     else:
-                        # Wait a bit and check again
-                        await asyncio.sleep(0.1)
-                except Exception as e:
-                    yield f"\n⚠️ Error during streaming: {str(e)}\n"
+                        result_preview = str(result)
+                    
+                    yield f"\n📊 Got result from {tool_name}\n"
             
-            # Get the final result
-            try:
-                result = await task
-                final_response = result.response if isinstance(result.response, str) else str(result.response)
-                
-                # Yield the final answer
-                yield f"\n\n✅ Final answer: {final_response}"
-            except Exception as e:
-                yield f"\n\n❌ Error getting final result: {str(e)}"
-                
-            # Restore original functions
-            for tool in tools:
-                tool_name = tool.metadata.name
-                if tool_name in original_functions:
-                    tool.fn = original_functions[tool_name]
     except Exception as e:
-        yield f"❌ Error: {str(e)}"
+        yield f"\n\n❌ Error: {str(e)}"
     finally:
         instrumentor.flush()
 
-# Gradio interface function
+# Updated gradio_query function
 async def gradio_query(user_input, chat_history=None):
     history = chat_history or []
     history.append({"role": "user", "content": user_input})
     
     # Add initial assistant message
-    history.append({"role": "assistant", "content": "Thinking..."})
+    history.append({"role": "assistant", "content": "Processing..."})
     yield history, history
     
     # Get streaming response