feat: Add lifecycle hooks for the agent LLM request/response cycle

Author: uzair330

src/agents/lifecycle.py

@@ -34,9 +34,7 @@ async def on_llm_end(
         """Called immediately after the LLM call returns for this agent."""
         pass
 
-    async def on_agent_start(
-        self, context: RunContextWrapper[TContext], agent: TAgent
-    ) -> None:
+    async def on_agent_start(self, context: RunContextWrapper[TContext], agent: TAgent) -> None:
         """Called before the agent is invoked. Called each time the current agent changes."""
         pass
 
@@ -152,4 +150,4 @@ async def on_llm_end(
 """Run hooks when using `Agent`."""
 
 AgentHooks = AgentHooksBase[TContext, Agent]
-"""Agent hooks for `Agent`s."""
+"""Agent hooks for `Agent`s."""

src/agents/run.py

@@ -943,11 +943,13 @@ async def _run_single_turn_streamed(
             input_items=input,
             system_instructions=system_prompt,
         )
-        
+
         # Call hook just before the model is invoked, with the correct system_prompt.
         if agent.hooks:
-            await agent.hooks.on_llm_start(context_wrapper, agent, filtered.instructions, filtered.input)
-        
+            await agent.hooks.on_llm_start(
+                context_wrapper, agent, filtered.instructions, filtered.input
+            )
+
         # 1. Stream the output events
         async for event in model.stream_response(
             filtered.instructions,
@@ -985,7 +987,7 @@ async def _run_single_turn_streamed(
             streamed_result._event_queue.put_nowait(RawResponsesStreamEvent(data=event))
 
         # Call hook just after the model response is finalized.
-        if agent.hooks:
+        if agent.hooks and final_response is not None:
             await agent.hooks.on_llm_end(context_wrapper, agent, final_response)
 
         # 2. At this point, the streaming is complete for this turn of the agent loop.
@@ -1266,8 +1268,8 @@ async def _get_new_response(
             await agent.hooks.on_llm_start(
                 context_wrapper,
                 agent,
-                filtered.instructions, # Use filtered instructions
-                filtered.input # Use filtered input
+                filtered.instructions,  # Use filtered instructions
+                filtered.input,  # Use filtered input
             )
 
         new_response = await model.get_response(
@@ -1285,11 +1287,7 @@ async def _get_new_response(
         )
         # If the agent has hooks, we need to call them after the LLM call
         if agent.hooks:
-            await agent.hooks.on_llm_end(
-                context_wrapper,
-                agent,
-                new_response
-            )
+            await agent.hooks.on_llm_end(context_wrapper, agent, new_response)
 
         context_wrapper.usage.add(new_response.usage)
 
@@ -1404,4 +1402,4 @@ async def _save_result_to_session(
 def _copy_str_or_list(input: str | list[TResponseInputItem]) -> str | list[TResponseInputItem]:
     if isinstance(input, str):
         return input
-    return input.copy()
+    return input.copy()

tests/test_agent_llm_hooks.py

@@ -127,4 +127,4 @@ async def test_streamed_agent_hooks_with_llm():
                 print(f"[EVENT] message_output_item → {text}")
 
     # Expect one on_start, one on_llm_start, one on_llm_end, and one on_end
-    assert hooks.events == {"on_start": 1, "on_llm_start": 1, "on_llm_end": 1, "on_end": 1}
+    assert hooks.events == {"on_start": 1, "on_llm_start": 1, "on_llm_end": 1, "on_end": 1}