Skip to content

Commit 325fae5

Browse files
committed
feat: update run.py to support new lifecycle hooks
1 parent 42422fb commit 325fae5

File tree

1 file changed

+15
-0
lines changed

1 file changed

+15
-0
lines changed

src/agents/run.py

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1252,6 +1252,14 @@ async def _get_new_response(
12521252
model = cls._get_model(agent, run_config)
12531253
model_settings = agent.model_settings.resolve(run_config.model_settings)
12541254
model_settings = RunImpl.maybe_reset_tool_choice(agent, tool_use_tracker, model_settings)
1255+
# If the agent has hooks, we need to call them before and after the LLM call
1256+
if agent.hooks:
1257+
await agent.hooks.on_llm_start(
1258+
context_wrapper,
1259+
agent,
1260+
system_prompt,
1261+
input
1262+
)
12551263

12561264
new_response = await model.get_response(
12571265
system_instructions=filtered.instructions,
@@ -1266,6 +1274,13 @@ async def _get_new_response(
12661274
previous_response_id=previous_response_id,
12671275
prompt=prompt_config,
12681276
)
1277+
# If the agent has hooks, we need to call them after the LLM call
1278+
if agent.hooks:
1279+
await agent.hooks.on_llm_end(
1280+
context_wrapper,
1281+
agent,
1282+
new_response
1283+
)
12691284

12701285
context_wrapper.usage.add(new_response.usage)
12711286

0 commit comments

Comments
 (0)