import json
from typing import Any, Callable, Dict
from agno.agent import Agent
from agno.run import RunContext
from agno.tools import Toolkit
# ---------------------------------------------------------------------------
# Create Agent
# ---------------------------------------------------------------------------
class CustomerDBTools(Toolkit):
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self.register(self.retrieve_customer_profile)
def retrieve_customer_profile(self, customer: str):
"""
Retrieves a customer profile from the database.
Args:
customer: The ID of the customer to retrieve.
Returns:
A string containing the customer profile.
"""
return customer
# When used as a tool hook, this function will receive the contextual Agent, function_name, etc as parameters
def grab_customer_profile_hook(
run_context: RunContext,
function_call: Callable,
arguments: Dict[str, Any],
):
cust_id = arguments.get("customer")
if cust_id not in run_context.session_state["customer_profiles"]: # type: ignore
raise ValueError(f"Customer profile for {cust_id} not found")
customer_profile = run_context.session_state["customer_profiles"][cust_id] # type: ignore
# Replace the customer with the customer_profile
arguments["customer"] = json.dumps(customer_profile)
# Call the function with the updated arguments
result = function_call(**arguments)
return result
agent = Agent(
tools=[CustomerDBTools()],
tool_hooks=[grab_customer_profile_hook],
session_state={
"customer_profiles": {
"123": {"name": "Jane Doe", "email": "jane.doe@example.com"},
"456": {"name": "John Doe", "email": "john.doe@example.com"},
}
},
)
# This should work
# ---------------------------------------------------------------------------
# Run Agent
# ---------------------------------------------------------------------------
if __name__ == "__main__":
agent.print_response("I am customer 456, please retrieve my profile.")
# This should fail
# agent.print_response("I am customer 789, please retrieve my profile.")
Run the Example
# Clone and setup repo
git clone https://github.com/agno-agi/agno.git
cd agno/cookbook/91_tools/tool_hooks
# Create and activate virtual environment
./scripts/demo_setup.sh
source .venvs/demo/bin/activate
python tool_hook_in_toolkit_with_state.py