feat: add streaming and tests to new hooks

Author: uzair330

src/agents/run.py

@@ -812,7 +812,9 @@ async def _run_single_turn_streamed(
 
         input = ItemHelpers.input_to_new_input_list(streamed_result.input)
         input.extend([item.to_input_item() for item in streamed_result.new_items])
-
+        # 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, system_prompt, input)
         # 1. Stream the output events
         async for event in model.stream_response(
             system_prompt,
@@ -849,6 +851,10 @@ 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:
+            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.
         if not final_response:
             raise ModelBehaviorError("Model did not produce a final response!")

tests/test_agent_llm_hooks.py

@@ -4,7 +4,7 @@
 import pytest
 
 from agents.agent import Agent
-from agents.items import ModelResponse, TResponseInputItem
+from agents.items import ItemHelpers, ModelResponse, TResponseInputItem
 from agents.lifecycle import AgentHooks
 from agents.run import Runner
 from agents.run_context import RunContextWrapper, TContext
@@ -63,7 +63,7 @@ async def on_llm_start(
 
     async def on_llm_end(
         self,
-        ccontext: RunContextWrapper[TContext],
+        context: RunContextWrapper[TContext],
         agent: Agent[TContext],
         response: ModelResponse,
     ) -> None:
@@ -72,7 +72,7 @@ async def on_llm_end(
 
 # Example test using the above hooks:
 @pytest.mark.asyncio
-async def test_non_streamed_agent_hooks_with_llm():
+async def test_async_agent_hooks_with_llm():
     hooks = AgentHooksForTests()
     model = FakeModel()
     agent = Agent(
@@ -83,3 +83,48 @@ async def test_non_streamed_agent_hooks_with_llm():
     await Runner.run(agent, input="hello")
     # 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}
+
+
+# test_sync_agent_hook_with_llm()
+def test_sync_agent_hook_with_llm():
+    hooks = AgentHooksForTests()
+    model = FakeModel()
+    agent = Agent(
+        name="A", model=model, tools=[get_function_tool("f", "res")], handoffs=[], hooks=hooks
+    )
+    # Simulate a single LLM call producing an output:
+    model.set_next_output([get_text_message("hello")])
+    Runner.run_sync(agent, input="hello")
+    # 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}
+
+
+# test_streamed_agent_hooks_with_llm():
+@pytest.mark.asyncio
+async def test_streamed_agent_hooks_with_llm():
+    hooks = AgentHooksForTests()
+    model = FakeModel()
+    agent = Agent(
+        name="A", model=model, tools=[get_function_tool("f", "res")], handoffs=[], hooks=hooks
+    )
+    # Simulate a single LLM call producing an output:
+    model.set_next_output([get_text_message("hello")])
+    stream = Runner.run_streamed(agent, input="hello")
+
+    async for event in stream.stream_events():
+        if event.type == "raw_response_event":
+            continue
+        if event.type == "agent_updated_stream_event":
+            print(f"[EVENT] agent_updated → {event.new_agent.name}")
+        elif event.type == "run_item_stream_event":
+            item = event.item
+            if item.type == "tool_call_item":
+                print("[EVENT] tool_call_item")
+            elif item.type == "tool_call_output_item":
+                print(f"[EVENT] tool_call_output_item → {item.output}")
+            elif item.type == "message_output_item":
+                text = ItemHelpers.text_message_output(item)
+                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}