From bb2c1f3cc5fea759fa905758f74546420412974c Mon Sep 17 00:00:00 2001 From: Your Name Date: Tue, 1 Jul 2025 15:08:13 +0500 Subject: [PATCH 1/7] feat: add lifecycle hooks to agent --- src/agents/lifecycle.py | 23 +++++++++++++++++++++++ 1 file changed, 23 insertions(+) diff --git a/src/agents/lifecycle.py b/src/agents/lifecycle.py index 8643248b1..96d3feb8b 100644 --- a/src/agents/lifecycle.py +++ b/src/agents/lifecycle.py @@ -10,6 +10,29 @@ class RunHooks(Generic[TContext]): override the methods you need. """ + #Two new hook methods added to the RunHooks class to handle LLM start and end events. + #These methods allow you to perform actions just before and after the LLM call for an agent. + #This is useful for logging, monitoring, or modifying the context before and after the LLM call. + async def on_llm_start( + self, + context: RunContextWrapper[TContext], + agent: Agent[TContext], + system_prompt: str | None, + input_items: List[TResponseInputItem] + ) -> None: + """Called just before invoking the LLM for this agent.""" + pass + + async def on_llm_end( + self, + context: RunContextWrapper[TContext], + agent: Agent[TContext], + response: ModelResponse + ) -> None: + """Called immediately after the LLM call returns for this agent.""" + pass + + async def on_agent_start( self, context: RunContextWrapper[TContext], agent: Agent[TContext] ) -> None: From 9302c49fda21a13aa9a56ddb3e0791c386ad4262 Mon Sep 17 00:00:00 2001 From: Your Name Date: Tue, 1 Jul 2025 15:23:38 +0500 Subject: [PATCH 2/7] feat: update run.py to support new lifecycle hooks --- src/agents/run.py | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/src/agents/run.py b/src/agents/run.py index e5f9378ec..cb37f84a2 100644 --- a/src/agents/run.py +++ b/src/agents/run.py @@ -1067,6 +1067,14 @@ async def _get_new_response( model = cls._get_model(agent, run_config) model_settings = agent.model_settings.resolve(run_config.model_settings) model_settings = RunImpl.maybe_reset_tool_choice(agent, tool_use_tracker, model_settings) + # If the agent has hooks, we need to call them before and after the LLM call + if agent.hooks: + await agent.hooks.on_llm_start( + context_wrapper, + agent, + system_prompt, + input + ) new_response = await model.get_response( system_instructions=system_prompt, @@ -1081,6 +1089,13 @@ async def _get_new_response( previous_response_id=previous_response_id, prompt=prompt_config, ) + # 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 + ) context_wrapper.usage.add(new_response.usage) From a775d559973386327660de106dfa37ba2d57326d Mon Sep 17 00:00:00 2001 From: Your Name Date: Tue, 1 Jul 2025 16:30:49 +0500 Subject: [PATCH 3/7] add type imports to lifecycle.py --- src/agents/lifecycle.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/agents/lifecycle.py b/src/agents/lifecycle.py index 96d3feb8b..d1b392260 100644 --- a/src/agents/lifecycle.py +++ b/src/agents/lifecycle.py @@ -1,5 +1,5 @@ -from typing import Any, Generic - +from typing import Any, Generic,List +from .items import ModelResponse, TResponseInputItem from .agent import Agent from .run_context import RunContextWrapper, TContext from .tool import Tool From 853b59595dfc76bf971c2c540933e9d73abd1950 Mon Sep 17 00:00:00 2001 From: Your Name Date: Tue, 1 Jul 2025 19:10:11 +0500 Subject: [PATCH 4/7] test: add unit tests for new agent lifecycle hooks --- tests/test_agent_llm_hooks.py | 183 ++++++++++++++++++++++++++++++++++ 1 file changed, 183 insertions(+) create mode 100644 tests/test_agent_llm_hooks.py diff --git a/tests/test_agent_llm_hooks.py b/tests/test_agent_llm_hooks.py new file mode 100644 index 000000000..143b4787e --- /dev/null +++ b/tests/test_agent_llm_hooks.py @@ -0,0 +1,183 @@ + +from typing import Any, List + +import pytest + +# Core SDK Imports +from agents.agent import Agent +from agents.run import Runner +from agents.lifecycle import AgentHooks +from agents.tool import Tool, function_tool, FunctionTool +from agents.items import ModelResponse +from agents.usage import Usage, InputTokensDetails, OutputTokensDetails +from agents.models.interface import Model + +# Types from the openai library used by the SDK +from openai.types.responses import ResponseFunctionToolCall, ResponseOutputMessage + +# --- 1. Spy Hook Implementation --- +class LoggingAgentHooks(AgentHooks[Any]): + def __init__(self): + super().__init__() + self.called_hooks: List[str] = [] + + # Spy on the NEW hooks + async def on_llm_start(self, *args, **kwargs): + self.called_hooks.append("on_llm_start") + + async def on_llm_end(self, *args, **kwargs): + self.called_hooks.append("on_llm_end") + + # Spy on EXISTING hooks to serve as landmarks for sequence verification + async def on_start(self, *args, **kwargs): + self.called_hooks.append("on_start") + + async def on_end(self, *args, **kwargs): + self.called_hooks.append("on_end") + + async def on_tool_start(self, *args, **kwargs): + self.called_hooks.append("on_tool_start") + + async def on_tool_end(self, *args, **kwargs): + self.called_hooks.append("on_tool_end") + +# --- 2. Mock Model and Tools --- +class MockModel(Model): + """A mock model that can be configured to either return a chat message or a tool call.""" + def __init__(self): + self._call_count = 0 + self._should_call_tool = False + self._tool_to_call: Tool | None = None + + def configure_for_tool_call(self, tool: Tool): + self._should_call_tool = True + self._tool_to_call = tool + + def configure_for_chat(self): + self._should_call_tool = False + self._tool_to_call = None + + async def get_response(self, *args, **kwargs) -> ModelResponse: + self._call_count += 1 + response_items: List[Any] = [] + + if self._should_call_tool and self._call_count == 1: + response_items.append( + ResponseFunctionToolCall(name=self._tool_to_call.name, arguments='{}', call_id="call123", type="function_call") + ) + else: + response_items.append( + ResponseOutputMessage(id="msg1", content=[{"type":"output_text", "text":"Mock response", "annotations":[]}], role="assistant", status="completed", type="message") + ) + + mock_usage = Usage( + requests=1, input_tokens=10, output_tokens=10, total_tokens=20, + input_tokens_details=InputTokensDetails(cached_tokens=0), + output_tokens_details=OutputTokensDetails(reasoning_tokens=0) + ) + return ModelResponse(output=response_items, usage=mock_usage, response_id="resp123") + + async def stream_response(self, *args, **kwargs): + final_response = await self.get_response(*args, **kwargs) + from openai.types.responses import ResponseCompletedEvent + class MockSDKResponse: + def __init__(self, id, output, usage): self.id, self.output, self.usage = id, output, usage + yield ResponseCompletedEvent(response=MockSDKResponse(final_response.response_id, final_response.output, final_response.usage), type="response_completed") + +@function_tool +def mock_tool(a: int, b: int) -> int: + """A mock tool for testing tool call hooks.""" + return a + b + +# --- 3. Pytest Fixtures for Test Setup --- +@pytest.fixture +def logging_hooks() -> LoggingAgentHooks: + """Provides a fresh instance of LoggingAgentHooks for each test.""" + return LoggingAgentHooks() + +@pytest.fixture +def chat_agent(logging_hooks: LoggingAgentHooks) -> Agent: + """Provides an agent configured for a simple chat interaction.""" + mock_model = MockModel() + mock_model.configure_for_chat() + return Agent( + name="ChatAgent", + instructions="Test agent for chat.", + model=mock_model, + hooks=logging_hooks + ) + +@pytest.fixture +def tool_agent(logging_hooks: LoggingAgentHooks) -> Agent: + """Provides an agent configured to use a tool.""" + mock_model = MockModel() + mock_model.configure_for_tool_call(mock_tool) + return Agent( + name="ToolAgent", + instructions="Test agent for tools.", + model=mock_model, + hooks=logging_hooks, + tools=[mock_tool] + ) + +# --- 4. Test Cases Focused on New Hooks --- +@pytest.mark.asyncio +async def test_llm_hooks_fire_in_chat_scenario( + chat_agent: Agent, logging_hooks: LoggingAgentHooks +): + """ + Tests that on_llm_start and on_llm_end fire correctly for a chat-only turn. + """ + await Runner.run(chat_agent, "Hello") + + sequence = logging_hooks.called_hooks + + expected_sequence = [ + "on_start", + "on_llm_start", + "on_llm_end", + "on_end", + ] + assert sequence == expected_sequence + +@pytest.mark.asyncio +async def test_llm_hooks_wrap_tool_hooks_in_tool_scenario( + tool_agent: Agent, logging_hooks: LoggingAgentHooks +): + """ + Tests that on_llm_start and on_llm_end wrap the tool execution cycle. + """ + await Runner.run(tool_agent, "Use your tool") + + sequence = logging_hooks.called_hooks + + expected_sequence = [ + "on_start", + "on_llm_start", + "on_llm_end", + "on_tool_start", + "on_tool_end", + "on_llm_start", + "on_llm_end", + "on_end" + ] + assert sequence == expected_sequence + +@pytest.mark.asyncio +async def test_no_hooks_run_if_hooks_is_none(): + """ + Ensures that the agent runs without error when agent.hooks is None. + """ + mock_model = MockModel() + mock_model.configure_for_chat() + agent_no_hooks = Agent( + name="NoHooksAgent", + instructions="Test agent without hooks.", + model=mock_model, + hooks=None + ) + + try: + await Runner.run(agent_no_hooks, "Hello") + except Exception as e: + pytest.fail(f"Runner.run failed when agent.hooks was None: {e}") \ No newline at end of file From 34e097d537ab09ab154a5c04546d805330cbc06c Mon Sep 17 00:00:00 2001 From: Your Name Date: Tue, 1 Jul 2025 20:39:35 +0500 Subject: [PATCH 5/7] style: Apply ruff formatting across project --- src/agents/lifecycle.py | 19 +++----- src/agents/run.py | 13 +---- tests/test_agent_llm_hooks.py | 91 ++++++++++++++++++++++------------- 3 files changed, 67 insertions(+), 56 deletions(-) diff --git a/src/agents/lifecycle.py b/src/agents/lifecycle.py index d1b392260..022dde3fe 100644 --- a/src/agents/lifecycle.py +++ b/src/agents/lifecycle.py @@ -1,6 +1,7 @@ -from typing import Any, Generic,List -from .items import ModelResponse, TResponseInputItem +from typing import Any, Generic + from .agent import Agent +from .items import ModelResponse, TResponseInputItem from .run_context import RunContextWrapper, TContext from .tool import Tool @@ -10,28 +11,24 @@ class RunHooks(Generic[TContext]): override the methods you need. """ - #Two new hook methods added to the RunHooks class to handle LLM start and end events. - #These methods allow you to perform actions just before and after the LLM call for an agent. - #This is useful for logging, monitoring, or modifying the context before and after the LLM call. + # Two new hook methods added to the RunHooks class to handle LLM start and end events. + # These methods allow you to perform actions just before and after the LLM call for an agent. + # This is useful for logging, monitoring, or modifying the context before and after the LLM call async def on_llm_start( self, context: RunContextWrapper[TContext], agent: Agent[TContext], system_prompt: str | None, - input_items: List[TResponseInputItem] + input_items: list[TResponseInputItem], ) -> None: """Called just before invoking the LLM for this agent.""" pass async def on_llm_end( - self, - context: RunContextWrapper[TContext], - agent: Agent[TContext], - response: ModelResponse + self, context: RunContextWrapper[TContext], agent: Agent[TContext], response: ModelResponse ) -> None: """Called immediately after the LLM call returns for this agent.""" pass - async def on_agent_start( self, context: RunContextWrapper[TContext], agent: Agent[TContext] diff --git a/src/agents/run.py b/src/agents/run.py index cb37f84a2..841813bd5 100644 --- a/src/agents/run.py +++ b/src/agents/run.py @@ -1069,12 +1069,7 @@ async def _get_new_response( model_settings = RunImpl.maybe_reset_tool_choice(agent, tool_use_tracker, model_settings) # If the agent has hooks, we need to call them before and after the LLM call if agent.hooks: - await agent.hooks.on_llm_start( - context_wrapper, - agent, - system_prompt, - input - ) + await agent.hooks.on_llm_start(context_wrapper, agent, system_prompt, input) new_response = await model.get_response( system_instructions=system_prompt, @@ -1091,11 +1086,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) diff --git a/tests/test_agent_llm_hooks.py b/tests/test_agent_llm_hooks.py index 143b4787e..3ba9fa3b6 100644 --- a/tests/test_agent_llm_hooks.py +++ b/tests/test_agent_llm_hooks.py @@ -1,25 +1,25 @@ - -from typing import Any, List +from typing import Any import pytest +# Types from the openai library used by the SDK +from openai.types.responses import ResponseFunctionToolCall, ResponseOutputMessage + # Core SDK Imports from agents.agent import Agent -from agents.run import Runner -from agents.lifecycle import AgentHooks -from agents.tool import Tool, function_tool, FunctionTool from agents.items import ModelResponse -from agents.usage import Usage, InputTokensDetails, OutputTokensDetails +from agents.lifecycle import AgentHooks from agents.models.interface import Model +from agents.run import Runner +from agents.tool import Tool, function_tool +from agents.usage import InputTokensDetails, OutputTokensDetails, Usage -# Types from the openai library used by the SDK -from openai.types.responses import ResponseFunctionToolCall, ResponseOutputMessage # --- 1. Spy Hook Implementation --- class LoggingAgentHooks(AgentHooks[Any]): def __init__(self): super().__init__() - self.called_hooks: List[str] = [] + self.called_hooks: list[str] = [] # Spy on the NEW hooks async def on_llm_start(self, *args, **kwargs): @@ -41,9 +41,11 @@ async def on_tool_start(self, *args, **kwargs): async def on_tool_end(self, *args, **kwargs): self.called_hooks.append("on_tool_end") + # --- 2. Mock Model and Tools --- class MockModel(Model): """A mock model that can be configured to either return a chat message or a tool call.""" + def __init__(self): self._call_count = 0 self._should_call_tool = False @@ -59,54 +61,77 @@ def configure_for_chat(self): async def get_response(self, *args, **kwargs) -> ModelResponse: self._call_count += 1 - response_items: List[Any] = [] + response_items: list[Any] = [] if self._should_call_tool and self._call_count == 1: response_items.append( - ResponseFunctionToolCall(name=self._tool_to_call.name, arguments='{}', call_id="call123", type="function_call") + ResponseFunctionToolCall( + name=self._tool_to_call.name, + arguments="{}", + call_id="call123", + type="function_call", + ) ) else: response_items.append( - ResponseOutputMessage(id="msg1", content=[{"type":"output_text", "text":"Mock response", "annotations":[]}], role="assistant", status="completed", type="message") + ResponseOutputMessage( + id="msg1", + content=[{"type": "output_text", "text": "Mock response", "annotations": []}], + role="assistant", + status="completed", + type="message", + ) ) - + mock_usage = Usage( - requests=1, input_tokens=10, output_tokens=10, total_tokens=20, + requests=1, + input_tokens=10, + output_tokens=10, + total_tokens=20, input_tokens_details=InputTokensDetails(cached_tokens=0), - output_tokens_details=OutputTokensDetails(reasoning_tokens=0) + output_tokens_details=OutputTokensDetails(reasoning_tokens=0), ) return ModelResponse(output=response_items, usage=mock_usage, response_id="resp123") async def stream_response(self, *args, **kwargs): final_response = await self.get_response(*args, **kwargs) from openai.types.responses import ResponseCompletedEvent + class MockSDKResponse: - def __init__(self, id, output, usage): self.id, self.output, self.usage = id, output, usage - yield ResponseCompletedEvent(response=MockSDKResponse(final_response.response_id, final_response.output, final_response.usage), type="response_completed") + def __init__(self, id, output, usage): + self.id, self.output, self.usage = id, output, usage + + yield ResponseCompletedEvent( + response=MockSDKResponse( + final_response.response_id, final_response.output, final_response.usage + ), + type="response_completed", + ) + @function_tool def mock_tool(a: int, b: int) -> int: """A mock tool for testing tool call hooks.""" return a + b + # --- 3. Pytest Fixtures for Test Setup --- @pytest.fixture def logging_hooks() -> LoggingAgentHooks: """Provides a fresh instance of LoggingAgentHooks for each test.""" return LoggingAgentHooks() + @pytest.fixture def chat_agent(logging_hooks: LoggingAgentHooks) -> Agent: """Provides an agent configured for a simple chat interaction.""" mock_model = MockModel() mock_model.configure_for_chat() return Agent( - name="ChatAgent", - instructions="Test agent for chat.", - model=mock_model, - hooks=logging_hooks + name="ChatAgent", instructions="Test agent for chat.", model=mock_model, hooks=logging_hooks ) + @pytest.fixture def tool_agent(logging_hooks: LoggingAgentHooks) -> Agent: """Provides an agent configured to use a tool.""" @@ -117,21 +142,20 @@ def tool_agent(logging_hooks: LoggingAgentHooks) -> Agent: instructions="Test agent for tools.", model=mock_model, hooks=logging_hooks, - tools=[mock_tool] + tools=[mock_tool], ) + # --- 4. Test Cases Focused on New Hooks --- @pytest.mark.asyncio -async def test_llm_hooks_fire_in_chat_scenario( - chat_agent: Agent, logging_hooks: LoggingAgentHooks -): +async def test_llm_hooks_fire_in_chat_scenario(chat_agent: Agent, logging_hooks: LoggingAgentHooks): """ Tests that on_llm_start and on_llm_end fire correctly for a chat-only turn. """ await Runner.run(chat_agent, "Hello") - + sequence = logging_hooks.called_hooks - + expected_sequence = [ "on_start", "on_llm_start", @@ -140,6 +164,7 @@ async def test_llm_hooks_fire_in_chat_scenario( ] assert sequence == expected_sequence + @pytest.mark.asyncio async def test_llm_hooks_wrap_tool_hooks_in_tool_scenario( tool_agent: Agent, logging_hooks: LoggingAgentHooks @@ -159,10 +184,11 @@ async def test_llm_hooks_wrap_tool_hooks_in_tool_scenario( "on_tool_end", "on_llm_start", "on_llm_end", - "on_end" + "on_end", ] assert sequence == expected_sequence + @pytest.mark.asyncio async def test_no_hooks_run_if_hooks_is_none(): """ @@ -171,13 +197,10 @@ async def test_no_hooks_run_if_hooks_is_none(): mock_model = MockModel() mock_model.configure_for_chat() agent_no_hooks = Agent( - name="NoHooksAgent", - instructions="Test agent without hooks.", - model=mock_model, - hooks=None + name="NoHooksAgent", instructions="Test agent without hooks.", model=mock_model, hooks=None ) - + try: await Runner.run(agent_no_hooks, "Hello") except Exception as e: - pytest.fail(f"Runner.run failed when agent.hooks was None: {e}") \ No newline at end of file + pytest.fail(f"Runner.run failed when agent.hooks was None: {e}") From fc50e934a0d92bf88ac2ea37c935d5d159a0e8f6 Mon Sep 17 00:00:00 2001 From: Your Name Date: Wed, 9 Jul 2025 01:04:29 +0500 Subject: [PATCH 6/7] fix: Resolve all CI failures for LLM hooks feature --- src/agents/lifecycle.py | 23 ++- tests/test_agent_llm_hooks.py | 267 ++++++++++------------------------ 2 files changed, 94 insertions(+), 196 deletions(-) diff --git a/src/agents/lifecycle.py b/src/agents/lifecycle.py index 022dde3fe..3e5903af8 100644 --- a/src/agents/lifecycle.py +++ b/src/agents/lifecycle.py @@ -1,4 +1,4 @@ -from typing import Any, Generic +from typing import Any, Generic, Optional from .agent import Agent from .items import ModelResponse, TResponseInputItem @@ -18,7 +18,7 @@ async def on_llm_start( self, context: RunContextWrapper[TContext], agent: Agent[TContext], - system_prompt: str | None, + system_prompt: Optional[str], input_items: list[TResponseInputItem], ) -> None: """Called just before invoking the LLM for this agent.""" @@ -123,3 +123,22 @@ async def on_tool_end( ) -> None: """Called after a tool is invoked.""" pass + + async def on_llm_start( + self, + context: RunContextWrapper[TContext], + agent: Agent[TContext], + system_prompt: Optional[str], + input_items: list[TResponseInputItem], + ) -> None: + """Called immediately before the agent issues an LLM call.""" + pass + + async def on_llm_end( + self, + context: RunContextWrapper[TContext], + agent: Agent[TContext], + response: ModelResponse, + ) -> None: + """Called immediately after the agent receives the LLM response.""" + pass diff --git a/tests/test_agent_llm_hooks.py b/tests/test_agent_llm_hooks.py index 3ba9fa3b6..ded7e18de 100644 --- a/tests/test_agent_llm_hooks.py +++ b/tests/test_agent_llm_hooks.py @@ -1,206 +1,85 @@ -from typing import Any +from collections import defaultdict +from typing import Any, Optional import pytest -# Types from the openai library used by the SDK -from openai.types.responses import ResponseFunctionToolCall, ResponseOutputMessage - -# Core SDK Imports from agents.agent import Agent -from agents.items import ModelResponse +from agents.items import ModelResponse, TResponseInputItem from agents.lifecycle import AgentHooks -from agents.models.interface import Model from agents.run import Runner -from agents.tool import Tool, function_tool -from agents.usage import InputTokensDetails, OutputTokensDetails, Usage - - -# --- 1. Spy Hook Implementation --- -class LoggingAgentHooks(AgentHooks[Any]): - def __init__(self): - super().__init__() - self.called_hooks: list[str] = [] - - # Spy on the NEW hooks - async def on_llm_start(self, *args, **kwargs): - self.called_hooks.append("on_llm_start") - - async def on_llm_end(self, *args, **kwargs): - self.called_hooks.append("on_llm_end") +from agents.run_context import RunContextWrapper, TContext +from agents.tool import Tool - # Spy on EXISTING hooks to serve as landmarks for sequence verification - async def on_start(self, *args, **kwargs): - self.called_hooks.append("on_start") +from .fake_model import FakeModel +from .test_responses import ( + get_function_tool, + get_text_message, +) - async def on_end(self, *args, **kwargs): - self.called_hooks.append("on_end") - - async def on_tool_start(self, *args, **kwargs): - self.called_hooks.append("on_tool_start") - - async def on_tool_end(self, *args, **kwargs): - self.called_hooks.append("on_tool_end") - - -# --- 2. Mock Model and Tools --- -class MockModel(Model): - """A mock model that can be configured to either return a chat message or a tool call.""" +class AgentHooksForTests(AgentHooks): def __init__(self): - self._call_count = 0 - self._should_call_tool = False - self._tool_to_call: Tool | None = None - - def configure_for_tool_call(self, tool: Tool): - self._should_call_tool = True - self._tool_to_call = tool - - def configure_for_chat(self): - self._should_call_tool = False - self._tool_to_call = None - - async def get_response(self, *args, **kwargs) -> ModelResponse: - self._call_count += 1 - response_items: list[Any] = [] - - if self._should_call_tool and self._call_count == 1: - response_items.append( - ResponseFunctionToolCall( - name=self._tool_to_call.name, - arguments="{}", - call_id="call123", - type="function_call", - ) - ) - else: - response_items.append( - ResponseOutputMessage( - id="msg1", - content=[{"type": "output_text", "text": "Mock response", "annotations": []}], - role="assistant", - status="completed", - type="message", - ) - ) - - mock_usage = Usage( - requests=1, - input_tokens=10, - output_tokens=10, - total_tokens=20, - input_tokens_details=InputTokensDetails(cached_tokens=0), - output_tokens_details=OutputTokensDetails(reasoning_tokens=0), - ) - return ModelResponse(output=response_items, usage=mock_usage, response_id="resp123") - - async def stream_response(self, *args, **kwargs): - final_response = await self.get_response(*args, **kwargs) - from openai.types.responses import ResponseCompletedEvent - - class MockSDKResponse: - def __init__(self, id, output, usage): - self.id, self.output, self.usage = id, output, usage - - yield ResponseCompletedEvent( - response=MockSDKResponse( - final_response.response_id, final_response.output, final_response.usage - ), - type="response_completed", - ) - - -@function_tool -def mock_tool(a: int, b: int) -> int: - """A mock tool for testing tool call hooks.""" - return a + b - - -# --- 3. Pytest Fixtures for Test Setup --- -@pytest.fixture -def logging_hooks() -> LoggingAgentHooks: - """Provides a fresh instance of LoggingAgentHooks for each test.""" - return LoggingAgentHooks() - - -@pytest.fixture -def chat_agent(logging_hooks: LoggingAgentHooks) -> Agent: - """Provides an agent configured for a simple chat interaction.""" - mock_model = MockModel() - mock_model.configure_for_chat() - return Agent( - name="ChatAgent", instructions="Test agent for chat.", model=mock_model, hooks=logging_hooks - ) - - -@pytest.fixture -def tool_agent(logging_hooks: LoggingAgentHooks) -> Agent: - """Provides an agent configured to use a tool.""" - mock_model = MockModel() - mock_model.configure_for_tool_call(mock_tool) - return Agent( - name="ToolAgent", - instructions="Test agent for tools.", - model=mock_model, - hooks=logging_hooks, - tools=[mock_tool], - ) - - -# --- 4. Test Cases Focused on New Hooks --- -@pytest.mark.asyncio -async def test_llm_hooks_fire_in_chat_scenario(chat_agent: Agent, logging_hooks: LoggingAgentHooks): - """ - Tests that on_llm_start and on_llm_end fire correctly for a chat-only turn. - """ - await Runner.run(chat_agent, "Hello") - - sequence = logging_hooks.called_hooks - - expected_sequence = [ - "on_start", - "on_llm_start", - "on_llm_end", - "on_end", - ] - assert sequence == expected_sequence - - -@pytest.mark.asyncio -async def test_llm_hooks_wrap_tool_hooks_in_tool_scenario( - tool_agent: Agent, logging_hooks: LoggingAgentHooks -): - """ - Tests that on_llm_start and on_llm_end wrap the tool execution cycle. - """ - await Runner.run(tool_agent, "Use your tool") - - sequence = logging_hooks.called_hooks - - expected_sequence = [ - "on_start", - "on_llm_start", - "on_llm_end", - "on_tool_start", - "on_tool_end", - "on_llm_start", - "on_llm_end", - "on_end", - ] - assert sequence == expected_sequence - - + self.events: dict[str, int] = defaultdict(int) + + def reset(self): + self.events.clear() + + async def on_start(self, context: RunContextWrapper[TContext], agent: Agent[TContext]) -> None: + self.events["on_start"] += 1 + + async def on_end( + self, context: RunContextWrapper[TContext], agent: Agent[TContext], output: Any + ) -> None: + self.events["on_end"] += 1 + + async def on_handoff( + self, context: RunContextWrapper[TContext], agent: Agent[TContext], source: Agent[TContext] + ) -> None: + self.events["on_handoff"] += 1 + + async def on_tool_start( + self, context: RunContextWrapper[TContext], agent: Agent[TContext], tool: Tool + ) -> None: + self.events["on_tool_start"] += 1 + + async def on_tool_end( + self, + context: RunContextWrapper[TContext], + agent: Agent[TContext], + tool: Tool, + result: str, + ) -> None: + self.events["on_tool_end"] += 1 + + # NEW: LLM hooks + async def on_llm_start( + self, + context: RunContextWrapper[TContext], + agent: Agent[TContext], + system_prompt: Optional[str], + input_items: list[TResponseInputItem], + ) -> None: + self.events["on_llm_start"] += 1 + + async def on_llm_end( + self, + ccontext: RunContextWrapper[TContext], + agent: Agent[TContext], + response: ModelResponse, + ) -> None: + self.events["on_llm_end"] += 1 + + +# Example test using the above hooks: @pytest.mark.asyncio -async def test_no_hooks_run_if_hooks_is_none(): - """ - Ensures that the agent runs without error when agent.hooks is None. - """ - mock_model = MockModel() - mock_model.configure_for_chat() - agent_no_hooks = Agent( - name="NoHooksAgent", instructions="Test agent without hooks.", model=mock_model, hooks=None +async def test_non_streamed_agent_hooks_with_llm(): + hooks = AgentHooksForTests() + model = FakeModel() + agent = Agent( + name="A", model=model, tools=[get_function_tool("f", "res")], handoffs=[], hooks=hooks ) - - try: - await Runner.run(agent_no_hooks, "Hello") - except Exception as e: - pytest.fail(f"Runner.run failed when agent.hooks was None: {e}") + # 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} From 525a3658dd9f2f4169a4f8841d137254f8d8679a Mon Sep 17 00:00:00 2001 From: Your Name Date: Wed, 16 Jul 2025 16:35:33 +0500 Subject: [PATCH 7/7] feat: add streaming and tests to new hooks --- src/agents/run.py | 8 +++++- tests/test_agent_llm_hooks.py | 51 ++++++++++++++++++++++++++++++++--- 2 files changed, 55 insertions(+), 4 deletions(-) diff --git a/src/agents/run.py b/src/agents/run.py index 841813bd5..406e1b549 100644 --- a/src/agents/run.py +++ b/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!") diff --git a/tests/test_agent_llm_hooks.py b/tests/test_agent_llm_hooks.py index ded7e18de..2eb2cfb03 100644 --- a/tests/test_agent_llm_hooks.py +++ b/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}