import os
import requests
from typing import List
from anthropic import Anthropic
from pydantic import BaseModel, Field
from tensorlake import Agent, Assistant, Function, FunctionTool, User, configure_logging, get_logger, Image, function
# Define the runtime environment
image = Image(name="weather-agent").run("pip install anthropic requests tensorlake pydantic")
class GetCurrentWeatherInput(BaseModel):
location: str = Field(..., description="The city and country, e.g., 'London, UK' or 'New York, USA'.")
@function(image=image, secrets=["OPENWEATHER_API_KEY"])
def get_current_weather(location: str) -> str:
"""
Fetches current weather data for a specified location using the OpenWeatherMap API.
"""
api_key = os.getenv("OPENWEATHER_API_KEY")
base_url = "http://api.openweathermap.org/data/2.5/weather"
params = {
"q": location,
"appid": api_key,
"units": "metric",
}
response = requests.get(base_url, params=params)
if response.status_code == 200:
data = response.json()
desc = data['weather'][0]['description']
temp = data['main']['temp']
return f"Weather in {location}: {desc}, {temp}°C"
else:
return f"Could not fetch weather for {location}."
class WeatherAgent(Agent):
def __init__(self, name: str, system_prompt: str, tools: List[FunctionTool]):
super().__init__(name, system_prompt, tools)
self.anthropic_client = Anthropic()
def call(self, user_message: User, files: List) -> Assistant:
messages = [{"role": "user", "content": user_message.content}]
response = self.anthropic_client.messages.create(
model="claude-3-opus-20240229",
max_tokens=1024,
system=self.system_prompt,
messages=messages,
tools=self.tools,
)
if response.stop_reason == "tool_use":
tool_use = next(block for block in response.content if block.type == "tool_use")
if tool_use.name == "get_current_weather":
location = tool_use.input["location"]
weather_info = get_current_weather(location)
messages.append({"role": "assistant", "content": response.content})
messages.append({
"role": "user",
"content": [
{
"type": "tool_result",
"tool_use_id": tool_use.id,
"content": weather_info
}
]
})
final_response = self.anthropic_client.messages.create(
model="claude-3-opus-20240229",
max_tokens=1024,
system=self.system_prompt,
messages=messages,
tools=self.tools,
)
return Assistant(content=final_response.content)
return Assistant(content=response.content)
# Initialize the agent
weather_agent = WeatherAgent(
name="weather_agent",
system_prompt="You are a helpful weather assistant.",
tools=[
FunctionTool(
name="get_current_weather",
description="Get current weather for a location.",
function=Function(
name="get_current_weather",
description="Get current weather for a location.",
input_schema=GetCurrentWeatherInput.model_json_schema(),
),
)
]
)
