From d98b73253b2fd0d99d07b8f23eb978da3d639842 Mon Sep 17 00:00:00 2001 From: Hemang Date: Mon, 10 Feb 2025 05:43:06 +0100 Subject: [PATCH] Skip Anthropic test if API key not set. --- tests/anthropic/test_claude_weather_agent.py | 103 +++++++++---------- 1 file changed, 51 insertions(+), 52 deletions(-) diff --git a/tests/anthropic/test_claude_weather_agent.py b/tests/anthropic/test_claude_weather_agent.py index bb61c3c..9315225 100644 --- a/tests/anthropic/test_claude_weather_agent.py +++ b/tests/anthropic/test_claude_weather_agent.py @@ -1,26 +1,25 @@ -import anthropic -from typing import Dict -import os -from tavily import TavilyClient -import anthropic -from httpx import Client -import os -# from invariant import testing -tavily_client = TavilyClient(api_key=os.getenv("TAVILY_API_KEY")) import datetime +import os +from typing import Dict + +import anthropic +import pytest +from httpx import Client +from tavily import TavilyClient + class WeatherAgent: def __init__(self, api_key: str): - dataset_name = "claude_weather_agent_test" + str(datetime.datetime.now().strftime("%Y%m%d%H%M%S")) - invariant_api_key = os.environ.get("INVARIANT_API_KEY") + self.tavily_client = TavilyClient(api_key=os.getenv("TAVILY_API_KEY")) + dataset_name = "claude_weather_agent_test" + str( + datetime.datetime.now().strftime("%Y%m%d%H%M%S") + ) self.client = anthropic.Anthropic( http_client=Client( - headers={ - "Invariant-Authorization": f"Bearer {invariant_api_key}" - }, - ), - base_url=f"http://localhost/api/v1/proxy/{dataset_name}/anthropic", - ) + headers={"Invariant-Authorization": "Bearer "}, + ), + base_url=f"http://localhost/api/v1/proxy/{dataset_name}/anthropic", + ) self.get_weather_function = { "name": "get_weather", "description": "Get the current weather in a given location", @@ -29,39 +28,34 @@ class WeatherAgent: "properties": { "location": { "type": "string", - "description": "The city and state, e.g. San Francisco, CA" + "description": "The city and state, e.g. San Francisco, CA", }, "unit": { "type": "string", "enum": ["celsius", "fahrenheit"], - "description": "The unit of temperature, either \"celsius\" or \"fahrenheit\"" - } + "description": 'The unit of temperature, either "celsius" or "fahrenheit"', + }, }, - "required": ["location"] - } + "required": ["location"], + }, } - # self.system_prompt = """You are an assistant that can perform weather searches using function calls. - # When a user asks for weather information, respond with a JSON object specifying + # self.system_prompt = """You are an assistant that can perform weather searches using function calls. + # When a user asks for weather information, respond with a JSON object specifying # the function name `get_weather` and the arguments latitude and longitude are needed.""" def get_response(self, user_query: str) -> Dict: """ Get the response from the agent for a given user query for weather. """ - messages = [ - { - "role": "user", - "content": user_query - } - ] + messages = [{"role": "user", "content": user_query}] while True: response = self.client.messages.create( # system=self.system_prompt, - tools = [self.get_weather_function], + tools=[self.get_weather_function], model="claude-3-5-sonnet-20241022", max_tokens=1024, - messages=messages + messages=messages, ) # If there's tool call, Extract the tool call parameters from the response @@ -69,44 +63,49 @@ class WeatherAgent: tool_call_params = response.content[1].input tool_call_result = self.get_weather(tool_call_params["location"]) tool_call_id = response.content[1].id - messages.append({ - "role": response.role, - "content": response.content + messages.append({"role": response.role, "content": response.content}) + messages.append( + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": tool_call_id, + "content": tool_call_result, + } + ], } ) - messages.append({ - "role": "user", - "content": [{ - "type": "tool_result", - "tool_use_id": tool_call_id, - "content": tool_call_result - }] - }) else: return response.content[0].text def get_weather(self, location: str): """Get the current weather in a given location using latitude and longitude.""" query = f"What is the weather in {location}?" - response = tavily_client.search(query) + response = self.tavily_client.search(query) response_content = response["results"][0]["content"] return response["results"][0]["title"] + ":\n" + response_content -# Initialize agent with your Anthropic API key -anthropic_api_key = os.getenv("ANTHROPIC_API_KEY") -weather_agent = WeatherAgent(anthropic_api_key) - -def test_proxy_response(): +@pytest.mark.skipif( + not os.getenv("ANTHROPIC_API_KEY") or not os.getenv("TAVILY_API_KEY"), + reason="API keys not set", +) +def test_proxy_response(): + """Test the proxy response for the weather agent.""" + # Initialize agent with Anthropic API key + anthropic_api_key = os.getenv("ANTHROPIC_API_KEY") + weather_agent = WeatherAgent(anthropic_api_key) + # Example queries queries = [ "What's the weather like in Zurich city?", "Tell me the forecast for New York", - "How's the weather in London next week?" + "How's the weather in London next week?", ] cities = ["Zurich", "New York", "London"] # Process each query - for index,query in enumerate(queries): + for index, query in enumerate(queries): response = weather_agent.get_response(query) assert response is not None - assert cities[index] in response \ No newline at end of file + assert cities[index] in response