2525from ..telemetry .metrics import Trace
2626from ..telemetry .tracer import get_tracer
2727from ..tools ._validator import validate_and_prepare_tools
28+ from ..types ._events import (
29+ EventLoopStopEvent ,
30+ EventLoopThrottleEvent ,
31+ ForceStopEvent ,
32+ ModelMessageEvent ,
33+ StartEvent ,
34+ StartEventLoopEvent ,
35+ ToolResultMessageEvent ,
36+ )
2837from ..types .content import Message
2938from ..types .exceptions import (
3039 ContextWindowOverflowException ,
@@ -91,8 +100,8 @@ async def event_loop_cycle(agent: "Agent", invocation_state: dict[str, Any]) ->
91100 cycle_start_time , cycle_trace = agent .event_loop_metrics .start_cycle (attributes = attributes )
92101 invocation_state ["event_loop_cycle_trace" ] = cycle_trace
93102
94- yield { "callback" : { "start" : True }}
95- yield { "callback" : { "start_event_loop" : True }}
103+ yield StartEvent ()
104+ yield StartEventLoopEvent ()
96105
97106 # Create tracer span for this event loop cycle
98107 tracer = get_tracer ()
@@ -175,7 +184,7 @@ async def event_loop_cycle(agent: "Agent", invocation_state: dict[str, Any]) ->
175184
176185 if isinstance (e , ModelThrottledException ):
177186 if attempt + 1 == MAX_ATTEMPTS :
178- yield { "callback" : { "force_stop" : True , "force_stop_reason" : str ( e )}}
187+ yield ForceStopEvent ( reason = e )
179188 raise e
180189
181190 logger .debug (
@@ -189,7 +198,7 @@ async def event_loop_cycle(agent: "Agent", invocation_state: dict[str, Any]) ->
189198 time .sleep (current_delay )
190199 current_delay = min (current_delay * 2 , MAX_DELAY )
191200
192- yield { "callback" : { "event_loop_throttled_delay" : current_delay , ** invocation_state }}
201+ yield EventLoopThrottleEvent ( delay = current_delay , invocation_state = invocation_state )
193202 else :
194203 raise e
195204
@@ -201,7 +210,7 @@ async def event_loop_cycle(agent: "Agent", invocation_state: dict[str, Any]) ->
201210 # Add the response message to the conversation
202211 agent .messages .append (message )
203212 agent .hooks .invoke_callbacks (MessageAddedEvent (agent = agent , message = message ))
204- yield { "callback" : { " message" : message }}
213+ yield ModelMessageEvent ( message = message )
205214
206215 # Update metrics
207216 agent .event_loop_metrics .update_usage (usage )
@@ -235,8 +244,8 @@ async def event_loop_cycle(agent: "Agent", invocation_state: dict[str, Any]) ->
235244 cycle_start_time = cycle_start_time ,
236245 invocation_state = invocation_state ,
237246 )
238- async for event in events :
239- yield event
247+ async for typed_event in events :
248+ yield typed_event
240249
241250 return
242251
@@ -264,11 +273,11 @@ async def event_loop_cycle(agent: "Agent", invocation_state: dict[str, Any]) ->
264273 tracer .end_span_with_error (cycle_span , str (e ), e )
265274
266275 # Handle any other exceptions
267- yield { "callback" : { "force_stop" : True , "force_stop_reason" : str ( e )}}
276+ yield ForceStopEvent ( reason = e )
268277 logger .exception ("cycle failed" )
269278 raise EventLoopException (e , invocation_state ["request_state" ]) from e
270279
271- yield { "stop" : (stop_reason , message , agent .event_loop_metrics , invocation_state ["request_state" ])}
280+ yield EventLoopStopEvent (stop_reason , message , agent .event_loop_metrics , invocation_state ["request_state" ])
272281
273282
274283async def recurse_event_loop (agent : "Agent" , invocation_state : dict [str , Any ]) -> AsyncGenerator [dict [str , Any ], None ]:
@@ -295,7 +304,7 @@ async def recurse_event_loop(agent: "Agent", invocation_state: dict[str, Any]) -
295304 recursive_trace = Trace ("Recursive call" , parent_id = cycle_trace .id )
296305 cycle_trace .add_child (recursive_trace )
297306
298- yield { "callback" : { "start" : True }}
307+ yield StartEvent ()
299308
300309 events = event_loop_cycle (agent = agent , invocation_state = invocation_state )
301310 async for event in events :
@@ -339,7 +348,7 @@ async def _handle_tool_execution(
339348 validate_and_prepare_tools (message , tool_uses , tool_results , invalid_tool_use_ids )
340349 tool_uses = [tool_use for tool_use in tool_uses if tool_use .get ("toolUseId" ) not in invalid_tool_use_ids ]
341350 if not tool_uses :
342- yield { "stop" : (stop_reason , message , agent .event_loop_metrics , invocation_state ["request_state" ])}
351+ yield EventLoopStopEvent (stop_reason , message , agent .event_loop_metrics , invocation_state ["request_state" ])
343352 return
344353
345354 tool_events = agent .tool_executor ._execute (
@@ -358,15 +367,15 @@ async def _handle_tool_execution(
358367
359368 agent .messages .append (tool_result_message )
360369 agent .hooks .invoke_callbacks (MessageAddedEvent (agent = agent , message = tool_result_message ))
361- yield { "callback" : { " message" : tool_result_message }}
370+ yield ToolResultMessageEvent ( message = message )
362371
363372 if cycle_span :
364373 tracer = get_tracer ()
365374 tracer .end_event_loop_cycle_span (span = cycle_span , message = message , tool_result_message = tool_result_message )
366375
367376 if invocation_state ["request_state" ].get ("stop_event_loop" , False ):
368377 agent .event_loop_metrics .end_cycle (cycle_start_time , cycle_trace )
369- yield { "stop" : (stop_reason , message , agent .event_loop_metrics , invocation_state ["request_state" ])}
378+ yield EventLoopStopEvent (stop_reason , message , agent .event_loop_metrics , invocation_state ["request_state" ])
370379 return
371380
372381 events = recurse_event_loop (agent = agent , invocation_state = invocation_state )
0 commit comments