feat: add streaming and tests to new hooks

Author: uzair330

src/agents/run.py

@@ -935,14 +935,19 @@ 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])
 
+        # THIS IS THE RESOLVED CONFLICT BLOCK
         filtered = await cls._maybe_filter_model_input(
             agent=agent,
             run_config=run_config,
             context_wrapper=context_wrapper,
             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)
+        
         # 1. Stream the output events
         async for event in model.stream_response(
             filtered.instructions,
@@ -979,6 +984,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!")
@@ -1257,8 +1266,8 @@ async def _get_new_response(
             await agent.hooks.on_llm_start(
                 context_wrapper,
                 agent,
-                system_prompt,
-                input
+                filtered.instructions, # Use filtered instructions
+                filtered.input # Use filtered input
             )
 
         new_response = await model.get_response(
@@ -1395,4 +1404,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

@@ -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(
@@ -81,5 +81,50 @@ async def test_non_streamed_agent_hooks_with_llm():
     # Simulate a single LLM call producing an output:
     model.set_next_output([get_text_message("hello")])
     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}