feat: add bubbling of events

Author: marcusschiesser

.env

@@ -5,7 +5,7 @@
 MODEL_PROVIDER=openai
 
 # The name of LLM model to use.
-MODEL=gpt-4o
+MODEL=gpt-4o-mini
 
 # Name of the embedding model to use.
 EMBEDDING_MODEL=text-embedding-3-large

.vscode/launch.json

@@ -0,0 +1,16 @@
+{
+    // Use IntelliSense to learn about possible attributes.
+    // Hover to view descriptions of existing attributes.
+    // For more information, visit: https://go.microsoft.com/fwlink/?linkid=830387
+    "version": "0.2.0",
+    "configurations": [
+        {
+            "name": "Python Debugger: Current File",
+            "type": "debugpy",
+            "request": "launch",
+            "program": "main.py",
+            "console": "integratedTerminal",
+            "justMyCode": false
+        }
+    ]
+}

app/core/agent_call.py

@@ -1,13 +1,54 @@
 import asyncio
 from typing import Any, List
 
-from llama_index.core.tools import FunctionTool
+from llama_index.core.tools.types import ToolMetadata, ToolOutput
+from llama_index.core.tools.utils import create_schema_from_function
+from llama_index.core.workflow import Context, Workflow
 
+from app.core.function_call import (
+    AgentRunResult,
+    ContextAwareTool,
+    FunctionCallingAgent,
+)
 from app.core.planner_agent import StructuredPlannerAgent
-from app.core.prefix import PrintPrefix
-from app.core.function_call import AgentRunResult, FunctionCallingAgent
 
-import textwrap
+
+class AgentCallTool(ContextAwareTool):
+    def __init__(self, agent: Workflow) -> None:
+        self.agent = agent
+        # create the schema without the context
+        name = f"call_{agent.name}"
+
+        async def schema_call(input: str) -> str:
+            pass
+
+        # create the schema without the Context
+        fn_schema = create_schema_from_function(name, schema_call)
+        self._metadata = ToolMetadata(
+            name=name,
+            description=(
+                f"Use this tool to delegate a sub task to the {agent.name} agent."
+                + (f" The agent is an {agent.role}." if agent.role else "")
+            ),
+            fn_schema=fn_schema,
+        )
+
+    # overload the acall function with the ctx argument as it's needed for bubbling the events
+    async def acall(self, ctx: Context, input: str) -> ToolOutput:
+        # FIXME: reset contexts, not needed after https://github.com/run-llama/llama_index/pull/15776
+        self.agent._contexts = set()
+        task = asyncio.create_task(self.agent.run(input=input))
+        # bubble all events while running the agent to the calling agent
+        async for ev in self.agent.stream_events():
+            ctx.write_event_to_stream(ev)
+        ret: AgentRunResult = await task
+        response = ret.response.message.content
+        return ToolOutput(
+            content=str(response),
+            tool_name=self.metadata.name,
+            raw_input={"args": input, "kwargs": {}},
+            raw_output=response,
+        )
 
 
 class AgentCallingAgent(FunctionCallingAgent):
@@ -19,7 +60,7 @@ def __init__(
         **kwargs: Any,
     ) -> None:
         agents = agents or []
-        tools = [_create_call_workflow_fn(self, agent) for agent in agents]
+        tools = [AgentCallTool(agent=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})
@@ -29,12 +70,12 @@ class AgentOrchestrator(StructuredPlannerAgent):
     def __init__(
         self,
         *args: Any,
-        agents: List[FunctionCallingAgent] | None = None,
         name: str = "orchestrator",
+        agents: List[FunctionCallingAgent] | None = None,
         **kwargs: Any,
     ) -> None:
         agents = agents or []
-        tools = [_create_call_workflow_fn(self, agent) for agent in agents]
+        tools = [AgentCallTool(agent=agent) for agent in agents]
         super().__init__(
             *args,
             name=name,
@@ -43,35 +84,3 @@ def __init__(
         )
         # call add_workflows so agents will get detected by llama agents automatically
         self.add_workflows(**{agent.name: agent for agent in agents})
-
-
-def _create_call_workflow_fn(
-    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)
-        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(
-        async_fn=acall_workflow_fn,
-        name=f"call_{agent.name}",
-        description=(
-            f"Use this tool to delegate a sub task to the {agent.name} agent."
-            + (f" The agent is an {agent.role}." if agent.role else "")
-        ),
-    )

app/core/function_call.py

@@ -1,3 +1,4 @@
+from abc import abstractmethod
 from typing import Any, List, Optional
 
 from llama_index.core.llms import ChatMessage, ChatResponse
@@ -6,6 +7,8 @@
 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.tools import FunctionTool
+
 from llama_index.core.workflow import (
     Context,
     Event,
@@ -25,11 +28,30 @@ class ToolCallEvent(Event):
     tool_calls: list[ToolSelection]
 
 
+class AgentRunEvent(Event):
+    name: str
+    _msg: str
+
+    @property
+    def msg(self):
+        return self._msg
+
+    @msg.setter
+    def msg(self, value):
+        self._msg = value
+
+
 class AgentRunResult(BaseModel):
     response: ChatResponse
     sources: list[ToolOutput]
 
 
+class ContextAwareTool(FunctionTool):
+    @abstractmethod
+    async def acall(self, ctx: Context, input: Any) -> ToolOutput:
+        pass
+
+
 class FunctionCallingAgent(Workflow):
     def __init__(
         self,
@@ -40,13 +62,15 @@ def __init__(
         verbose: bool = False,
         timeout: float = 360.0,
         name: str,
+        write_events: bool = True,
         role: Optional[str] = None,
         **kwargs: Any,
     ) -> None:
         super().__init__(*args, verbose=verbose, timeout=timeout, **kwargs)
         self.tools = tools or []
         self.name = name
         self.role = role
+        self.write_events = write_events
 
         if llm is None:
             llm = Settings.llm
@@ -72,9 +96,10 @@ async def prepare_chat_history(self, ctx: Context, 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}")
-        )
+        if self.write_events:
+            ctx.write_event_to_stream(
+                AgentRunEvent(name=self.name, msg=f"Start to work on: {user_input}")
+            )
 
         # get chat history
         chat_history = self.memory.get()
@@ -96,7 +121,10 @@ async def handle_llm_input(
         )
 
         if not tool_calls:
-            ctx.session.write_event_to_stream(Event(msg=f"[{self.name}] Finished task"))
+            if self.write_events:
+                ctx.write_event_to_stream(
+                    AgentRunEvent(name=self.name, msg="Finished task")
+                )
             return StopEvent(
                 result=AgentRunResult(response=response, sources=[*self.sources])
             )
@@ -128,7 +156,11 @@ async def handle_tool_calls(self, ctx: Context, ev: ToolCallEvent) -> InputEvent
                 continue
 
             try:
-                tool_output = await tool.acall(**tool_call.tool_kwargs)
+                if isinstance(tool, ContextAwareTool):
+                    # inject context for calling an context aware tool
+                    tool_output = await tool.acall(ctx=ctx, **tool_call.tool_kwargs)
+                else:
+                    tool_output = await tool.acall(**tool_call.tool_kwargs)
                 self.sources.append(tool_output)
                 tool_msgs.append(
                     ChatMessage(

app/core/planner_agent.py

@@ -1,3 +1,4 @@
+import asyncio
 from typing import Any, List
 
 from llama_index.core.llms.function_calling import FunctionCallingLLM
@@ -10,7 +11,7 @@
     step,
 )
 
-from app.core.function_call import AgentRunResult, FunctionCallingAgent
+from app.core.function_call import AgentRunEvent, AgentRunResult, FunctionCallingAgent
 from app.core.planner import Planner, SubTask, Plan
 from llama_index.core.tools import BaseTool
 
@@ -35,8 +36,7 @@ class PlanEventType(Enum):
     REFINED = "refined"
 
 
-class PlanEvent(Event):
-
+class PlanEvent(AgentRunEvent):
     event_type: PlanEventType
     plan: Plan
 
@@ -68,9 +68,11 @@ def __init__(
             name="executor",
             llm=llm,
             tools=self.tools,
+            write_events=False,
             # it's important to instruct to just return the tool call, otherwise the executor will interpret and change the result
             system_prompt="You are an expert in completing given tasks by calling the right tool for the task. Just return the result of the tool call. Don't add any information yourself",
         )
+        self.add_workflows(executor=self.executor)
 
     @step()
     async def create_plan(
@@ -80,8 +82,8 @@ async def create_plan(
         ctx.data["task"] = ev.input
         ctx.data["act_plan_id"] = plan_id
         # inform about the new plan
-        ctx.session.write_event_to_stream(
-            PlanEvent(event_type=PlanEventType.CREATED, plan=plan)
+        ctx.write_event_to_stream(
+            PlanEvent(name=self.name, event_type=PlanEventType.CREATED, plan=plan)
         )
         if self._verbose:
             print("=== Executing plan ===\n")
@@ -97,7 +99,7 @@ async def execute_plan(self, ctx: Context, ev: ExecutePlanEvent) -> SubTaskEvent
         # send an event per sub task
         events = [SubTaskEvent(sub_task=sub_task) for sub_task in upcoming_sub_tasks]
         for event in events:
-            ctx.session.send_event(event)
+            ctx.send_event(event)
 
         return None
 
@@ -107,7 +109,13 @@ async def execute_sub_task(
     ) -> SubTaskResultEvent:
         if self._verbose:
             print(f"=== Executing sub task: {ev.sub_task.name} ===")
-        result: AgentRunResult = await self.executor.run(input=ev.sub_task.input)
+        # FIXME: reset contexts, not needed after https://github.com/run-llama/llama_index/pull/15776
+        self.executor._contexts = set()
+        task = asyncio.create_task(self.executor.run(input=ev.sub_task.input))
+        # bubble all events while running the executor to the planner
+        async for event in self.executor.stream_events():
+            ctx.write_event_to_stream(event)
+        result: AgentRunResult = await task
         if self._verbose:
             print("=== Done executing sub task ===\n")
         self.planner.state.add_completed_sub_task(ctx.data["act_plan_id"], ev.sub_task)
@@ -141,8 +149,10 @@ async def gather_results(
             )
             # inform about the new plan
             if new_plan is not None:
-                ctx.session.write_event_to_stream(
-                    PlanEvent(event_type=PlanEventType.REFINED, plan=new_plan)
+                ctx.write_event_to_stream(
+                    PlanEvent(
+                        name=self.name, event_type=PlanEventType.REFINED, plan=new_plan
+                    )
                 )
 
         # continue executing plan

main.py

@@ -1,9 +1,10 @@
 # flake8: noqa: E402
 import asyncio
 import os
+import textwrap
 from dotenv import load_dotenv
 from app.core.agent_call import AgentCallingAgent, AgentOrchestrator
-from app.core.function_call import AgentRunResult, FunctionCallingAgent
+from app.core.function_call import AgentRunEvent, AgentRunResult, FunctionCallingAgent
 from app.engine.index import get_index
 from app.settings import init_settings
 from llama_index.core.tools import QueryEngineTool, ToolMetadata
@@ -80,6 +81,11 @@ def create_orchestrator():
     )
 
 
+def info(prefix: str, text: str) -> None:
+    truncated = textwrap.shorten(text, width=255, placeholder="...")
+    print(f"[{prefix}] {truncated}")
+
+
 async def main():
     # agent = create_choreography()
     agent = create_orchestrator()
@@ -88,10 +94,11 @@ async def main():
     )
 
     async for ev in agent.stream_events():
-        print(ev.msg)
+        if isinstance(ev, AgentRunEvent):
+            info(ev.name, ev.msg)
 
     ret: AgentRunResult = await task
-    print(ret.response.message.content)
+    print(f"\n\nResult:\n\n{ret.response.message.content}")
 
 
 if __name__ == "__main__":