-
Notifications
You must be signed in to change notification settings - Fork 2.8k
fix(run): fire on_llm_start / on_llm_end in Runner.run() for streaming & non-streaming (aligns with docs) #1619
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Conversation
8faa564 to
cb7bec5
Compare
cb7bec5 to
02792b2
Compare
|
Overall, this looks great. I will look into this early next week. Can you make changes to the lifecycle example too? diff --git a/examples/basic/lifecycle_example.py b/examples/basic/lifecycle_example.py
index f37380b..63ce3f8 100644
--- a/examples/basic/lifecycle_example.py
+++ b/examples/basic/lifecycle_example.py
@@ -1,10 +1,11 @@
import asyncio
import random
-from typing import Any
+from typing import Any, Optional
from pydantic import BaseModel
from agents import Agent, RunContextWrapper, RunHooks, Runner, Tool, Usage, function_tool
+from agents.items import ModelResponse, TResponseInputItem
class ExampleHooks(RunHooks):
@@ -20,6 +21,22 @@ class ExampleHooks(RunHooks):
f"### {self.event_counter}: Agent {agent.name} started. Usage: {self._usage_to_str(context.usage)}"
)
+ async def on_llm_start(
+ self,
+ context: RunContextWrapper,
+ agent: Agent,
+ system_prompt: Optional[str],
+ input_items: list[TResponseInputItem],
+ ) -> None:
+ self.event_counter += 1
+ print(f"### {self.event_counter}: LLM started. Usage: {self._usage_to_str(context.usage)}")
+
+ async def on_llm_end(
+ self, context: RunContextWrapper, agent: Agent, response: ModelResponse
+ ) -> None:
+ self.event_counter += 1
+ print(f"### {self.event_counter}: LLM ended. Usage: {self._usage_to_str(context.usage)}")
+
async def on_agent_end(self, context: RunContextWrapper, agent: Agent, output: Any) -> None:
self.event_counter += 1
print( |
|
Please fix this one too:
|
02792b2 to
784fd29
Compare
|
Thanks for the 👀 @seratch - PR updated.
When updating |
| # 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) | ||
| # If we have run hooks, or if the agent has hooks, we need to call them after the LLM call |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
can you move this right after L1422's usage update? It shouldn't bring any visible overhead in processing time and can provide better insights for the callback
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Updated.
This fixes the issue I was seeing when running the lifecycle test. Thanks!
- added run hooks handling with agent hooks
784fd29 to
4b1a29f
Compare
Summary
Fix a discrepancy between the documented lifecycle and the Python runner:
on_llm_startandon_llm_endnow emit fromRunner.run()consistently in both streaming and non-streaming flows._get_new_responseand_run_single_turn_streamed.Behavior
on_llm_startfires immediately before an LLM request (sync/async/streaming).on_llm_endfires only on successful completion (after final chunk in streaming).on_llm_endandon_agent_endare not emitted; errors propagate. This preserves current semantics.Hook ordering (because run hooks and agent hooks dispatch concurrently):
Compatibility
_get_new_responsesignature updated to support unified flow. No other call sites/overrides in repo; considered safe as underscored/private.Tests added
on_agent_start,on_llm_start,on_llm_end,on_agent_endeach fire once.Runner.run_sync.FakeModel.get_responseto raise; asserton_llm_startfires buton_llm_end/on_agent_enddo not.BoomModel.stream_responseyields once then raises; asserton_llm_startfires buton_llm_end/on_agent_enddo not.Changelog
Fix: Python runner now emits
on_llm_start/on_llm_endas documented in both streaming and non-streaming paths.on_llm_endremains success-only; error paths propagate exceptions without firing end hooks.Fixes #1612