Integrate Your Custom Python Functions with the Tool Use Agent¶
Overview¶
The Tool Use Agent is a utility agent designed to perform function calls using provided tools. It enables dynamic execution of functions based on user queries, allowing for a flexible and extensible system. By integrating both built-in and custom tools, the agent can process a wide range of tasks—from simple calculations to complex data processing.
Configuration¶
To utilize the Tool Use Agent, you need to define its configuration in a YAML file. This configuration specifies the agent's settings and the tools it can access.
Configuration Parameters¶
agent_class: The class name of the agent. For the Tool Use Agent, this should beToolUseAgent.config: Configuration settings for the agent.wait_time: The maximum time (in seconds) to wait for an external function call to complete.enable_interpreter: A boolean indicating whether to use the Interpreter Agent to process function outputs before returning them to the user.builtin_tools: A list of built-in tool names that the agent can use.custom_tools: A list of custom tool definitions (as JSON strings) that the agent can use.
Here's an example configuration (example.yaml):
utility_agents:
- agent_class: ToolUseAgent
agent_name: "Tool Use Agent"
agent_description: "An agent that performs function calling using provided tools."
config:
wait_time: 120
enable_interpreter: true
builtin_tools:
- "calculate_expression"
custom_tools:
- |
{
"type": "function",
"function": {
"name": "generate_password",
"description": "Generate a random password of specified length.",
"parameters": {
"type": "object",
"properties": {
"length": {
"type": "integer",
"description": "Length of the password to generate. Default is 12.",
"default": 12
}
},
"required": []
}
}
}
orchestrator:
agent_list:
- agent_name: "Tool Use Agent"
Built-in Tools¶
The builtin_tools parameter allows you to specify which predefined tools the agent can access. In the example above, the agent has access to the calculate_expression tool, which evaluates mathematical expressions. We plan to add more built-in tools in the future to extend the agent's capabilities, providing even greater functionality.
Custom Tools¶
You can provide custom Python functions to the Tool Use Agent using the custom_tools parameter. To do this, you need to:
- Add your custom functions to the
executor_dict: This is a dictionary that maps function names to their implementations. - Provide the JSON schema of each function in the configuration: This lets the agent understand how to call your custom functions and what parameters they expect.
This ensures that the agent can correctly interpret user queries and map them to the appropriate custom functions.
Example Usage¶
Here's how to use the Tool Use Agent with your custom Python functions:
Define Your Custom Python Function¶
First, define your custom function. For example, a function to generate a random password:
import asyncio
import os
import random
import string
from air import DistillerClient
from dotenv import load_dotenv
load_dotenv() # loads your API_KEY from your local '.env' file
api_key=str(os.getenv("API_KEY"))
async def generate_password(length: int = 12) -> str:
"""
Generate a random password of specified length.
Args:
length (int, optional): The total length of the password to generate.
Must be at least 4 to include one of each required character type.
Defaults to 12.
Returns:
str: A randomly generated password string.
"""
if length < 4:
raise ValueError("Password length should be at least 4 characters.")
password_chars = [
random.choice(string.ascii_uppercase),
random.choice(string.ascii_lowercase),
random.choice(string.digits),
random.choice(string.punctuation),
]
if length > 4:
all_chars = string.ascii_letters + string.digits + string.punctuation
password_chars.extend(random.choice(all_chars) for _ in range(length - 4))
random.shuffle(password_chars)
password = "".join(password_chars)
return password
Initialize the Distiller Client and Add Your Function¶
Next, initialize the DistillerClient and add your function to the executor_dict:
async def tool_use_demo():
# Initialize the DistillerClient
distiller_client = DistillerClient(api_key=api_key)
# Register a new project with your configuration
distiller_client.create_project(config_path="example.yaml", project="example")
# Map custom agent names to their corresponding functions
executor_dict = {
"generate_password": generate_password,
}
async with distiller_client(
project="example",
uuid="test_user",
executor_dict=executor_dict,
) as dc:
# List of queries to process
queries = [
"Generate a safe password with 23 chars.",
]
for query in queries:
# Send the query and print responses
responses = await dc.query(query=query)
print(f"----\nQuery: {query}")
async for response in responses:
print(f"Response: {response['content']}")
if __name__ == "__main__":
asyncio.run(tool_use_demo())
In this example:
DistillerClient: Connects to the distiller server.create_project: Sets up a new project using your configuration file (example.yaml).executor_dict: Links custom agent names to their corresponding Python functions.queries: A list of user queries for the agent to process.Processing Queries: The agent handles each query, invokes the appropriate functions as needed, and returns the responses.
Expected Output¶
After running the code above, you can expect the following outputs:
With Interpreter Agent¶
If enable_interpreter is set to true, the Interpreter Agent processes the function's output to make it more user-friendly. The response might look like:
----
Query: Generate a safe password with 23 chars.
Response: Sure! Here's a randomly generated safe password with 23 characters:
`g@5Yq^12Bz&Mn8$!j0Rc)w#`
This password includes uppercase letters, lowercase letters, digits, and special characters to enhance security.
Without Interpreter Agent¶
If enable_interpreter is set to false, the raw output from the function is returned. The response might look like:
In the raw output, the password is provided without any additional explanation or formatting.
Advanced Guidelines for wait_time - Cancellable Custom Tool Functions¶
Timeout Behavior¶
If a tool runs longer than wait_time, the system stops waiting and returns an error:
"I am unable to handle this query at this moment."
However, the tool may still be running in the background unless it was written to stop cleanly when cancelled. Timeout does not automatically guarantee termination. Proper cancellation depends on how the tool is implemented.
For instance, async def tool functions that spend their time in await calls (like asyncio.sleep() or async I/O) can stop quickly, while blocking calls (like input() and time.sleep()) usually keep running until they finish.
Cancellable Examples (Recommended)¶
import asyncio
import sys
async def collect_terminal_feedback() -> str:
"""
Collect terminal input asynchronously so ToolUseAgent timeout can cancel it.
"""
query = "Share feedback: "
# Use the current event loop so stdin can be handled without blocking.
loop = asyncio.get_running_loop()
print(query, end="", flush=True)
# Future is completed by the stdin-ready callback below.
fut = loop.create_future()
def on_ready() -> None:
if fut.done():
return
# Read one line and resolve the future when stdin becomes readable.
fut.set_result(sys.stdin.readline().rstrip("\n"))
# Register a non-blocking stdin callback (Unix/macOS/Linux).
loop.add_reader(sys.stdin, on_ready)
try:
# Await is cancellation-friendly, so timeout can cancel this wait.
return await fut
finally:
# Always remove the reader to avoid leaking callbacks across calls.
loop.remove_reader(sys.stdin)
async def slow_operation() -> str:
"""Sleep longer than ToolUseAgent.wait_time to force timeout."""
delay = 15
print("Start sleeping")
# asyncio.sleep is cancellation-aware and can be interrupted by timeout.
await asyncio.sleep(delay)
print("Finish sleeping")
return f"Finished sleeping for {delay} seconds"
Non-Cancellable Blocking Examples (Avoid)¶
import time
def collect_terminal_feedback() -> str:
"""Blocking input call; timeout cancellation will not interrupt promptly."""
prompt = "Share feedback: "
print(prompt, end="", flush=True)
# input() blocks the thread; timeout cancellation cannot interrupt it promptly.
response = input()
return response or "(empty feedback)"
def slow_operation() -> str:
"""Blocking sleep call; timeout cancellation will not interrupt promptly."""
delay = 15
print("Start sleeping")
# time.sleep() blocks the thread; timeout cancellation cannot interrupt it promptly.
time.sleep(delay)
print("Finish sleeping")
return f"Finished sleeping for {delay} seconds"