1515from __future__ import annotations
1616
1717import inspect
18+ import warnings
1819from typing import Any
1920from typing import AsyncGenerator
2021from typing import Awaitable
@@ -186,19 +187,21 @@ def _load_agent_state(
186187 def _create_agent_state_event (
187188 self ,
188189 ctx : InvocationContext ,
190+ * ,
191+ agent_state : Optional [BaseAgentState ] = None ,
192+ end_of_agent : bool = False ,
189193 ) -> Event :
190- """Returns an event with current agent state set in the invocation context .
194+ """Returns an event with agent state.
191195
192196 Args:
193197 ctx: The invocation context.
194-
195- Returns:
196- An event with the current agent state set in the invocation context.
198+ agent_state: The agent state to checkpoint.
199+ end_of_agent: Whether the agent is finished running.
197200 """
198201 event_actions = EventActions ()
199- if ( agent_state := ctx . agent_states . get ( self . name )) is not None :
200- event_actions .agent_state = agent_state
201- if ctx . end_of_agents . get ( self . name ) :
202+ if agent_state :
203+ event_actions .agent_state = agent_state . model_dump ( mode = 'json' )
204+ if end_of_agent :
202205 event_actions .end_of_agent = True
203206 return Event (
204207 invocation_id = ctx .invocation_id ,
@@ -284,22 +287,27 @@ async def run_async(
284287 Event: the events generated by the agent.
285288 """
286289
287- with tracer .start_as_current_span (f'invoke_agent { self .name } ' ) as span :
288- ctx = self ._create_invocation_context (parent_context )
289- tracing .trace_agent_invocation (span , self , ctx )
290- if event := await self ._handle_before_agent_callback (ctx ):
291- yield event
292- if ctx .end_invocation :
293- return
294-
295- async with Aclosing (self ._run_async_impl (ctx )) as agen :
296- async for event in agen :
290+ async def _run_with_trace () -> AsyncGenerator [Event , None ]:
291+ with tracer .start_as_current_span (f'invoke_agent { self .name } ' ) as span :
292+ ctx = self ._create_invocation_context (parent_context )
293+ tracing .trace_agent_invocation (span , self , ctx )
294+ if event := await self .__handle_before_agent_callback (ctx ):
297295 yield event
296+ if ctx .end_invocation :
297+ return
298+
299+ async with Aclosing (self ._run_async_impl (ctx )) as agen :
300+ async for event in agen :
301+ yield event
302+
303+ if ctx .end_invocation :
304+ return
298305
299- if ctx . end_invocation :
300- return
306+ if event := await self . __handle_after_agent_callback ( ctx ) :
307+ yield event
301308
302- if event := await self ._handle_after_agent_callback (ctx ):
309+ async with Aclosing (_run_with_trace ()) as agen :
310+ async for event in agen :
303311 yield event
304312
305313 @final
@@ -317,19 +325,24 @@ async def run_live(
317325 Event: the events generated by the agent.
318326 """
319327
320- with tracer .start_as_current_span (f'invoke_agent { self .name } ' ) as span :
321- ctx = self ._create_invocation_context (parent_context )
322- tracing .trace_agent_invocation (span , self , ctx )
323- if event := await self ._handle_before_agent_callback (ctx ):
324- yield event
325- if ctx .end_invocation :
326- return
328+ async def _run_with_trace () -> AsyncGenerator [Event , None ]:
329+ with tracer .start_as_current_span (f'invoke_agent { self .name } ' ) as span :
330+ ctx = self ._create_invocation_context (parent_context )
331+ tracing .trace_agent_invocation (span , self , ctx )
332+ if event := await self .__handle_before_agent_callback (ctx ):
333+ yield event
334+ if ctx .end_invocation :
335+ return
327336
328- async with Aclosing (self ._run_live_impl (ctx )) as agen :
329- async for event in agen :
337+ async with Aclosing (self ._run_live_impl (ctx )) as agen :
338+ async for event in agen :
339+ yield event
340+
341+ if event := await self .__handle_after_agent_callback (ctx ):
330342 yield event
331343
332- if event := await self ._handle_after_agent_callback (ctx ):
344+ async with Aclosing (_run_with_trace ()) as agen :
345+ async for event in agen :
333346 yield event
334347
335348 async def _run_async_impl (
@@ -406,7 +419,43 @@ def _create_invocation_context(
406419 invocation_context = parent_context .model_copy (update = {'agent' : self })
407420 return invocation_context
408421
409- async def _handle_before_agent_callback (
422+ @property
423+ def canonical_before_agent_callbacks (self ) -> list [_SingleAgentCallback ]:
424+ """Deprecated: Use normalize_callbacks(self.before_agent_callback).
425+
426+ This property is deprecated and will be removed in a future version.
427+ Use normalize_callbacks() from callback_pipeline module instead.
428+
429+ Returns:
430+ List of before_agent callbacks.
431+ """
432+ warnings .warn (
433+ 'canonical_before_agent_callbacks is deprecated. '
434+ 'Use normalize_callbacks(self.before_agent_callback) instead.' ,
435+ DeprecationWarning ,
436+ stacklevel = 2 ,
437+ )
438+ return normalize_callbacks (self .before_agent_callback )
439+
440+ @property
441+ def canonical_after_agent_callbacks (self ) -> list [_SingleAgentCallback ]:
442+ """Deprecated: Use normalize_callbacks(self.after_agent_callback).
443+
444+ This property is deprecated and will be removed in a future version.
445+ Use normalize_callbacks() from callback_pipeline module instead.
446+
447+ Returns:
448+ List of after_agent callbacks.
449+ """
450+ warnings .warn (
451+ 'canonical_after_agent_callbacks is deprecated. '
452+ 'Use normalize_callbacks(self.after_agent_callback) instead.' ,
453+ DeprecationWarning ,
454+ stacklevel = 2 ,
455+ )
456+ return normalize_callbacks (self .after_agent_callback )
457+
458+ async def __handle_before_agent_callback (
410459 self , ctx : InvocationContext
411460 ) -> Optional [Event ]:
412461 """Runs the before_agent_callback if it exists.
@@ -458,7 +507,7 @@ async def _handle_before_agent_callback(
458507
459508 return None
460509
461- async def _handle_after_agent_callback (
510+ async def __handle_after_agent_callback (
462511 self , invocation_context : InvocationContext
463512 ) -> Optional [Event ]:
464513 """Runs the after_agent_callback if it exists.
0 commit comments