119 lines
3.2 KiB
Python
119 lines
3.2 KiB
Python
"""Example: Custom Tools
|
|
|
|
This example demonstrates how to define and use custom tools with agents.
|
|
"""
|
|
|
|
import asyncio
|
|
import os
|
|
from datetime import datetime
|
|
|
|
from pydantic import BaseModel
|
|
|
|
from agentlite import Agent, OpenAIProvider, tool
|
|
from agentlite.tool import CallableTool2, ToolOk
|
|
|
|
|
|
# Define a tool using the decorator
|
|
@tool()
|
|
async def get_current_time() -> str:
|
|
"""Get the current date and time."""
|
|
return datetime.now().isoformat()
|
|
|
|
|
|
@tool()
|
|
async def calculate(expression: str) -> str:
|
|
"""Evaluate a mathematical expression.
|
|
|
|
Args:
|
|
expression: The mathematical expression to evaluate (e.g., "2 + 2").
|
|
"""
|
|
try:
|
|
# Safe evaluation - only allow basic math operations
|
|
allowed_names = {
|
|
"abs": abs,
|
|
"max": max,
|
|
"min": min,
|
|
"pow": pow,
|
|
"round": round,
|
|
}
|
|
result = eval(expression, {"__builtins__": {}}, allowed_names)
|
|
return str(result)
|
|
except Exception as e:
|
|
return f"Error: {e}"
|
|
|
|
|
|
# Define a tool using CallableTool2 (type-safe)
|
|
class WeatherParams(BaseModel):
|
|
"""Parameters for weather tool."""
|
|
|
|
city: str
|
|
units: str = "celsius"
|
|
|
|
|
|
class GetWeather(CallableTool2[WeatherParams]):
|
|
"""Get weather information for a city."""
|
|
|
|
name = "get_weather"
|
|
description = "Get the current weather for a city."
|
|
params = WeatherParams
|
|
|
|
async def __call__(self, params: WeatherParams) -> ToolOk:
|
|
# This is a mock implementation
|
|
# In a real scenario, you would call a weather API
|
|
weather_data = {
|
|
"Beijing": {"temp": 22, "condition": "Sunny"},
|
|
"Shanghai": {"temp": 25, "condition": "Cloudy"},
|
|
"New York": {"temp": 18, "condition": "Rainy"},
|
|
"London": {"temp": 15, "condition": "Overcast"},
|
|
}
|
|
|
|
city = params.city
|
|
if city in weather_data:
|
|
data = weather_data[city]
|
|
temp = data["temp"]
|
|
if params.units == "fahrenheit":
|
|
temp = temp * 9 // 5 + 32
|
|
return ToolOk(
|
|
output=f"Weather in {city}: {data['condition']}, {temp}°{params.units[0].upper()}"
|
|
)
|
|
return ToolOk(output=f"Weather data not available for {city}")
|
|
|
|
|
|
async def main():
|
|
"""Run the custom tools example."""
|
|
# Create provider
|
|
provider = OpenAIProvider(
|
|
api_key=os.getenv("OPENAI_API_KEY", "your-api-key"),
|
|
model="gpt-4o-mini",
|
|
)
|
|
|
|
# Create agent with tools
|
|
agent = Agent(
|
|
provider=provider,
|
|
system_prompt="You are a helpful assistant with access to tools.",
|
|
tools=[
|
|
get_current_time,
|
|
calculate,
|
|
GetWeather(),
|
|
],
|
|
)
|
|
|
|
# Test tools
|
|
print("=== Testing Tools ===\n")
|
|
|
|
print("User: What time is it?")
|
|
response = await agent.run("What time is it?")
|
|
print(f"Agent: {response}\n")
|
|
|
|
print("User: What is 123 * 456?")
|
|
response = await agent.run("What is 123 * 456?")
|
|
print(f"Agent: {response}\n")
|
|
|
|
print("User: What's the weather in Beijing?")
|
|
response = await agent.run("What's the weather in Beijing?")
|
|
print(f"Agent: {response}\n")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
asyncio.run(main())
|