bubble events

Author: marcusschiesser

app/core/agent_call.py

@@ -1,3 +1,4 @@
+import asyncio
 from typing import Any, List
 
 from llama_index.core.tools import FunctionTool
@@ -18,7 +19,7 @@ def __init__(
         **kwargs: Any,
     ) -> None:
         agents = agents or []
-        tools = [_create_call_workflow_fn(name, agent) for agent in agents]
+        tools = [_create_call_workflow_fn(self, agent) for agent in agents]
         super().__init__(*args, name=name, tools=tools, **kwargs)
         # call add_workflows so agents will get detected by llama agents automatically
         self.add_workflows(**{agent.name: agent for agent in agents})
@@ -33,7 +34,7 @@ def __init__(
         **kwargs: Any,
     ) -> None:
         agents = agents or []
-        tools = [_create_call_workflow_fn(name, agent) for agent in agents]
+        tools = [_create_call_workflow_fn(self, agent) for agent in agents]
         super().__init__(
             *args,
             name=name,
@@ -45,18 +46,25 @@ def __init__(
 
 
 def _create_call_workflow_fn(
-    caller_name: str, agent: FunctionCallingAgent
+    caller: FunctionCallingAgent, agent: FunctionCallingAgent
 ) -> FunctionTool:
     def info(prefix: str, text: str) -> None:
         truncated = textwrap.shorten(text, width=255, placeholder="...")
         print(f"{prefix}: '{truncated}'")
 
     async def acall_workflow_fn(input: str) -> str:
-        info(f"[{caller_name}->{agent.name}]", input)
-        with PrintPrefix(f"[{agent.name}]"):
-            ret: AgentRunResult = await agent.run(input=input)
-            response = ret.response.message.content
-        info(f"[{caller_name}<-{agent.name}]", response)
+        # info(f"[{caller_name}->{agent.name}]", input)
+        task = asyncio.create_task(agent.run(input=input))
+        # bubble all events while running the agent to the calling agent
+        if len(caller._sessions) > 1:
+            print("XXX: Bubbling events only works with single-session agents")
+        else:
+            session = next(iter(caller._sessions))
+            async for ev in agent.stream_events():
+                session.write_event_to_stream(ev)
+        ret: AgentRunResult = await task
+        response = ret.response.message.content
+        # info(f"[{caller_name}<-{agent.name}]", response)
         return response
 
     return FunctionTool.from_defaults(

app/core/function_call.py

@@ -1,14 +1,19 @@
 from typing import Any, List, Optional
 
+from llama_index.core.llms import ChatMessage, ChatResponse
 from llama_index.core.llms.function_calling import FunctionCallingLLM
 from llama_index.core.memory import ChatMemoryBuffer
-from llama_index.core.tools.types import BaseTool
-from llama_index.core.workflow import Workflow, StartEvent, StopEvent, step
-
-from llama_index.core.llms import ChatMessage, ChatResponse
-from llama_index.core.tools import ToolSelection, ToolOutput
-from llama_index.core.workflow import Event
 from llama_index.core.settings import Settings
+from llama_index.core.tools import ToolOutput, ToolSelection
+from llama_index.core.tools.types import BaseTool
+from llama_index.core.workflow import (
+    Context,
+    Event,
+    StartEvent,
+    StopEvent,
+    Workflow,
+    step,
+)
 from pydantic import BaseModel
 
 
@@ -54,7 +59,7 @@ def __init__(
         self.sources = []
 
     @step()
-    async def prepare_chat_history(self, ev: StartEvent) -> InputEvent:
+    async def prepare_chat_history(self, ctx: Context, ev: StartEvent) -> InputEvent:
         # clear sources
         self.sources = []
 
@@ -67,13 +72,18 @@ async def prepare_chat_history(self, ev: StartEvent) -> InputEvent:
         user_input = ev.input
         user_msg = ChatMessage(role="user", content=user_input)
         self.memory.put(user_msg)
+        ctx.session.write_event_to_stream(
+            Event(msg=f"[{self.name}] Start to work on: {user_input}")
+        )
 
         # get chat history
         chat_history = self.memory.get()
         return InputEvent(input=chat_history)
 
     @step()
-    async def handle_llm_input(self, ev: InputEvent) -> ToolCallEvent | StopEvent:
+    async def handle_llm_input(
+        self, ctx: Context, ev: InputEvent
+    ) -> ToolCallEvent | StopEvent:
         chat_history = ev.input
 
         response = await self.llm.achat_with_tools(
@@ -86,14 +96,15 @@ async def handle_llm_input(self, ev: InputEvent) -> ToolCallEvent | StopEvent:
         )
 
         if not tool_calls:
+            ctx.session.write_event_to_stream(Event(msg=f"[{self.name}] Finished task"))
             return StopEvent(
                 result=AgentRunResult(response=response, sources=[*self.sources])
             )
         else:
             return ToolCallEvent(tool_calls=tool_calls)
 
     @step()
-    async def handle_tool_calls(self, ev: ToolCallEvent) -> InputEvent:
+    async def handle_tool_calls(self, ctx: Context, ev: ToolCallEvent) -> InputEvent:
         tool_calls = ev.tool_calls
         tools_by_name = {tool.metadata.get_name(): tool for tool in self.tools}