Skip to content

feat(agents): Add on_llm_start and on_llm_end Lifecycle Hooks #987

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.

Sign up for GitHub

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

Jump to bottom
Open
wants to merge 6 commits into
base: main
Choose a base branch
from
Open

feat(agents): Add on_llm_start and on_llm_end Lifecycle Hooks #987

wants to merge 6 commits into from

Conversation

uzair330
Copy link

@uzair330 uzair330 commented Jul 1, 2025 •
edited
Loading

Motivation

Currently, the AgentHooks provide valuable lifecycle events for the start/end of an agent run and for tool execution (on_tool_start/on_tool_end). However, developers lack the ability to observe the agent's execution at the language model level.

This PR introduces two new hooks, on_llm_start and on_llm_end, to provide this deeper level of observability. This change enables several key use cases:

  • Performance Monitoring: Precisely measure the latency of LLM calls.
  • Debugging & Logging: Log the exact prompts sent to and raw responses received from the model.
  • Implementing Custom Logic: Trigger actions (e.g., updating a UI, saving state) immediately before or after the agent "thinks."

Summary of Changes

  • src/agents/lifecycle.py
    Added two new async methods, on_llm_start and on_llm_end, to the AgentHooks base class, matching the existing on_*_start/on_*_end pattern.

  • src/agents/run.py
    Wrapped the call to model.get_response(...) in _get_new_response with invocations of the new hooks so that they fire immediately before and after each LLM call.

  • tests/test_agent_llm_hooks.py
    Added unit tests (using a mock model and spy hooks) to validate:

    1. The correct sequence of on_start → on_llm_start → on_llm_end → on_end in a chat‑only run.
    2. The correct wrapping of tool execution in a tool‑using run:
      on_start → on_llm_start → on_llm_end → on_tool_start → on_tool_end → on_llm_start → on_llm_end → on_end.
    3. That the agent still runs without error when agent.hooks is None.

Usage Examples

1. Async Example (awaitable via run)

import asyncio
from typing import Any, Optional

from dotenv import load_dotenv

from agents.agent import Agent
from agents.items import ModelResponse, TResponseInputItem
from agents.lifecycle import AgentHooks, RunContextWrapper
from agents.run import Runner

# Load any OPENAI_API_KEY or other env vars
load_dotenv()


# --- 1. Define a custom hooks class to track LLM calls ---
class LLMTrackerHooks(AgentHooks[Any]):
    async def on_llm_start(
        self,
        context: RunContextWrapper,
        agent: Agent,
        system_prompt: Optional[str],
        input_items: list[TResponseInputItem],
    ) -> None:
        print(
            f">>> [HOOK] Agent '{agent.name}' is calling the LLM with system prompt: {system_prompt or '[none]'}"
        )

    async def on_llm_end(
        self,
        context: RunContextWrapper,
        agent: Agent,
        response: ModelResponse,
    ) -> None:
        if response.usage:
            print(f">>> [HOOK] LLM call finished. Tokens used: {response.usage.total_tokens}")


# --- 2. Create your agent with these hooks ---
my_agent = Agent(
    name="MyMonitoredAgent",
    instructions="Tell me a joke.",
    hooks=LLMTrackerHooks(),
)


# --- 3. Drive it via an async main() ---
async def main():
    result = await Runner.run(my_agent, "Tell me a joke.")
    print(f"\nAgent output:\n{result.final_output}")


if __name__ == "__main__":
    asyncio.run(main())

2. Sync Example (blocking via run_sync)

from typing import Any, Optional

from dotenv import load_dotenv

from agents.agent import Agent
from agents.items import ModelResponse, TResponseInputItem
from agents.lifecycle import AgentHooks, RunContextWrapper
from agents.run import Runner

# Load any OPENAI_API_KEY or other env vars
load_dotenv()


# --- 1. Define a custom hooks class to track LLM calls ---
class LLMTrackerHooks(AgentHooks[Any]):
    async def on_llm_start(
        self,
        context: RunContextWrapper,
        agent: Agent,
        system_prompt: Optional[str],
        input_items: list[TResponseInputItem],
    ) -> None:
        print(
            f">>> [HOOK] Agent '{agent.name}' is calling the LLM with system prompt: {system_prompt or '[none]'}"
        )

    async def on_llm_end(
        self,
        context: RunContextWrapper,
        agent: Agent,
        response: ModelResponse,
    ) -> None:
        if response.usage:
            print(f">>> [HOOK] LLM call finished. Tokens used: {response.usage.total_tokens}")


# --- 2. Create your agent with these hooks ---
my_agent = Agent(
    name="MyMonitoredAgent",
    instructions="Tell me a joke.",
    hooks=LLMTrackerHooks(),
)


# --- 3. Drive it via an async main() ---
def main():
    result = Runner.run_sync(my_agent, "Tell me a joke.")
    print(f"\nAgent output:\n{result.final_output}")


if __name__ == "__main__":
    main()

Note

Streaming support for on_llm_start and on_llm_end is not yet implemented. These hooks currently fire only on non‑streamed (batch) LLM calls. Support for streaming invocations will be added in a future release.

Checklist

  • My code follows the style guidelines of this project (checked with ruff).
  • I have added tests that prove my feature works.
  • All new and existing tests passed locally with my changes.

@seratch seratch added enhancement New feature or request feature:core labels Jul 8, 2025
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers
No reviews
Assignees
No one assigned
Labels
enhancement New feature or request feature:core
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
2 participants
@uzair330 @seratch