CalvinBackup/invariant-gateway
1
0
Fork 0
mirror of https://github.com/invariantlabs-ai/invariant-gateway.git synced 2026-07-09 20:28:36 +02:00
Code Issues Packages Projects Releases Wiki Activity

add anthropic policy and related tests

Browse Source 
add anthropic policy and related tests
This commit is contained in:
zishan-wei
2025-02-18 16:55:37 +01:00
committed by GitHub
parent acb451a10e 0adf3729c2
commit d6e551cb3b
10 changed files with 559 additions and 146 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.github/workflows/tests_ci.yml
+1
View File
@@ -19,4 +19,5 @@ jobs:
- name: Run tests
env:
OPENAI_API_KEY: ${{ secrets.INVARIANT_TESTING_OPENAI_KEY }}
ANTHROPIC_API_KEY: ${{ secrets.INVARIANT_TESTING_ANTHROPIC_KEY}}
run: ./run.sh tests -s -vv
proxy/routes/anthropic.py
+149 -32
View File
@@ -7,24 +7,30 @@ import httpx
from fastapi import APIRouter, Depends, Header, HTTPException, Request
from utils.constants import CLIENT_TIMEOUT, IGNORED_HEADERS
from utils.explorer import push_trace
from starlette.responses import StreamingResponse
proxy = APIRouter()
ALLOWED_ANTHROPIC_ENDPOINTS = {"v1/messages"}
MISSING_INVARIANT_AUTH_HEADER = "Missing invariant-authorization header"
MISSING_ANTHROPIC_AUTH_HEADER = "Missing athropic authorization header"
NOT_SUPPORTED_ENDPOINT = "Not supported OpenAI endpoint"
MISSING_INVARIANT_AUTH_API_KEY = "Missing invariant authorization header"
MISSING_ANTHROPIC_AUTH_HEADER = "Missing Anthropic authorization header"
NOT_SUPPORTED_ENDPOINT = "Not supported Anthropic endpoint"
FAILED_TO_PUSH_TRACE = "Failed to push trace to the dataset: "
END_REASONS = ["end_turn", "max_tokens", "stop_sequence"]
MESSAGE_START = "message_start"
MESSGAE_DELTA = "message_delta"
MESSAGE_STOP = "message_stop"
CONTENT_BLOCK_START = "content_block_start"
CONTENT_BLOCK_DELTA = "content_block_delta"
CONTENT_BLOCK_STOP = "content_block_stop"
def validate_headers(
invariant_authorization: str = Header(None), x_api_key: str = Header(None)
x_api_key: str = Header(None)
):
"""Require the invariant-authorization and authorization headers to be present"""
if invariant_authorization is None:
raise HTTPException(status_code=400, detail=MISSING_INVARIANT_AUTH_HEADER)
"""Require the headers to be present"""
if x_api_key is None:
raise HTTPException(status_code=400, detail=MISSING_ANTHROPIC_AUTH_HEADER)
@@ -44,22 +50,44 @@ async def anthropic_proxy(
headers = {
k: v for k, v in request.headers.items() if k.lower() not in IGNORED_HEADERS
}
headers["accept-encoding"] = "identity"
if request.headers.get(
"invariant-authorization"
) is None and "|invariant-auth:" not in request.headers.get("authorization"):
raise HTTPException(status_code=400, detail=MISSING_INVARIANT_AUTH_API_KEY)
if request.headers.get("invariant-authorization"):
invariant_authorization = request.headers.get("invariant-authorization")
else:
authorization = request.headers.get("x-api-key")
api_keys = authorization.split("|invariant-auth: ")
invariant_authorization = f"Bearer {api_keys[1].strip()}"
# Update the authorization header to pass the OpenAI API Key to the OpenAI API
headers["x-api-key"] = f"{api_keys[0].strip()}"
request_body = await request.body()
request_body_json = json.loads(request_body)
anthropic_url = f"https://api.anthropic.com/{endpoint}"
client = httpx.AsyncClient(timeout=httpx.Timeout(CLIENT_TIMEOUT))
anthropic_request = client.build_request(
"POST", anthropic_url, headers=headers, data=request_body
)
invariant_authorization = request.headers.get("invariant-authorization")
async with client:
response = await client.send(anthropic_request)
if request_body_json.get("stream"):
return await handle_streaming_response(client, anthropic_request, dataset_name, invariant_authorization)
else:
try:
response = await client.send(anthropic_request)
except httpx.HTTPStatusError as e:
raise HTTPException(
status_code=response.status_code,
detail=f"Failed to fetch response: {response.text}, got error{e}",
)
await handle_non_streaming_response(
response, dataset_name, request_body_json, invariant_authorization
)
@@ -71,13 +99,15 @@ async def push_to_explorer(
merged_response: dict[str, Any],
request_body: dict[str, Any],
invariant_authorization: str,
reformat: bool = True,
) -> None:
"""Pushes the full trace to the Invariant Explorer"""
# Combine the messages from the request body and Anthropic response
messages = request_body.get("messages", [])
messages += [merged_response]
messages = anthropic_to_invariant_messages(messages)
if reformat:
messages = anthropic_to_invariant_messages(messages)
_ = await push_trace(
dataset_name=dataset_name,
messages=[messages],
@@ -102,6 +132,90 @@ async def handle_non_streaming_response(
invariant_authorization,
)
async def handle_streaming_response(
client: httpx.AsyncClient,
anthropic_request: httpx.Request,
dataset_name: str,
invariant_authorization: str
) -> StreamingResponse:
formatted_invariant_response = []
async def event_generator() -> Any:
async with client.stream(
"POST",
anthropic_request.url,
headers=anthropic_request.headers,
content=anthropic_request.content,
) as response:
if response.status_code != 200:
yield json.dumps(
{"error": f"Failed to fetch response: {response.status_code}"}
).encode()
return
async for chunk in response.aiter_bytes():
yield chunk
process_chunk_text(
chunk,
formatted_invariant_response
)
if formatted_invariant_response and formatted_invariant_response[-1].get("stop_reason") in END_REASONS:
await push_to_explorer(
dataset_name,
formatted_invariant_response[-1],
json.loads(anthropic_request.content),
invariant_authorization,
)
generator = event_generator()
return StreamingResponse(generator, media_type="text/event-stream")
def process_chunk_text(chunk, formatted_invariant_response):
"""
Process the chunk of text and update the formatted_invariant_response
Example of chunk list can be find in:
../../resources/streaming_chunk_text/anthropic.txt
"""
text_decode = chunk.decode().strip()
for text_block in text_decode.split("\n\n"):
# might be empty block
if len(text_block.split("\ndata:"))>1:
text_data = text_block.split("\ndata:")[1]
text_json = json.loads(text_data)
update_formatted_invariant_response(text_json, formatted_invariant_response)
def update_formatted_invariant_response(text_json, formatted_invariant_response):
if text_json.get("type") == MESSAGE_START:
message = text_json.get("message")
formatted_invariant_response.append({
"id": message.get("id"),
"role": message.get("role"),
"content": "",
"model": message.get("model"),
"stop_reason": message.get("stop_reason"),
"stop_sequence": message.get("stop_sequence"),
})
elif text_json.get("type") == CONTENT_BLOCK_START and text_json.get("content_block").get("type")=="tool_use":
content_block = text_json.get("content_block")
formatted_invariant_response.append(
{
"role": "tool",
"tool_id": content_block.get("id"),
"content": "",
}
)
elif text_json.get("type") == CONTENT_BLOCK_DELTA:
if formatted_invariant_response[-1]["role"]=="assistant":
formatted_invariant_response[-1]["content"] += text_json.get("delta").get("text")
elif formatted_invariant_response[-1]["role"]=="tool":
formatted_invariant_response[-1]["content"] += text_json.get("delta").get("partial_json")
elif text_json.get("type") == MESSGAE_DELTA:
formatted_invariant_response[-1]["stop_reason"] = text_json.get("delta").get("stop_reason")
def anthropic_to_invariant_messages(
messages: list[dict], keep_empty_tool_response: bool = False
@@ -155,24 +269,27 @@ def handle_user_message(message, keep_empty_tool_response):
def handle_assistant_message(message):
output = []
for sub_message in message["content"]:
if sub_message["type"] == "text":
output.append({"role": "assistant", "content": sub_message.get("text")})
elif sub_message["type"] == "tool_use":
output.append(
{
"role": "assistant",
"content": None,
"tool_calls": [
{
"tool_id": sub_message.get("id"),
"type": "function",
"function": {
"name": sub_message.get("name"),
"arguments": sub_message.get("input"),
},
}
],
}
)
if isinstance(message["content"], list):
for sub_message in message["content"]:
if sub_message["type"] == "text":
output.append({"role": "assistant", "content": sub_message.get("text")})
elif sub_message["type"] == "tool_use":
output.append(
{
"role": "assistant",
"content": None,
"tool_calls": [
{
"tool_id": sub_message.get("id"),
"type": "function",
"function": {
"name": sub_message.get("name"),
"arguments": sub_message.get("input"),
},
}
],
}
)
else:
output.append({"role": "assistant", "content": message["content"]})
return output
proxy/utils/explorer.py
-1
View File
@@ -8,7 +8,6 @@ from invariant_sdk.types.push_traces import PushTracesRequest, PushTracesRespons
DEFAULT_API_URL = "https://explorer.invariantlabs.ai"
async def push_trace(
messages: List[List[Dict[str, Any]]],
dataset_name: str,
resources/streaming_chunk_text/anthropic.txt
+6
View File
@@ -0,0 +1,6 @@
anthropic_chunk_list = [b'event: message_start\ndata: {"type":"message_start","message":{"id":"msg_012KWB6kiKvzx7r1SKs5nGA1","type":"message","role":"assistant","model":"claude-3-5-sonnet-20241022","content":[],"stop_reason":null,"stop_sequence":null,"usage":{"input_tokens":20,"cache_creation_input_tokens":0,"cache_read_input_tokens":0,"output_tokens":5}} }\n\nevent: content_block_start\ndata: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""} }\n\n'
,b'event: content_block_delta\ndata: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":" making it an attractive destination for both business"} }\n\n'
, b'event: content_block_delta\ndata: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":" and leisure."} }\n\n'
, b'event: content_block_stop\ndata: {"type":"content_block_stop","index":0}\n\n'
, b'event: message_delta\ndata: {"type":"message_delta","delta":{"stop_reason":"end_turn","stop_sequence":null},"usage":{"output_tokens":301} }\n\n'
, b'event: message_stop\ndata: {"type":"message_stop" }\n\n']
run.sh
+1
View File
@@ -68,6 +68,7 @@ tests() {
--mount type=bind,source=./tests,target=/tests \
--network invariant-proxy-web-test \
-e OPENAI_API_KEY="$OPENAI_API_KEY" \
-e ANTHROPIC_API_KEY="$ANTHROPIC_API_KEY"\
--env-file ./tests/.env.test \
explorer-proxy-tests $@
}
tests/anthropic/test_anthropic_with_tool_call.py
+264
View File
@@ -0,0 +1,264 @@
import datetime
import os
from typing import Dict
import json
import anthropic
import pytest
from httpx import Client
import sys
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from util import * # needed for pytest fixtures
pytest_plugins = ("pytest_asyncio",)
class WeatherAgent:
def __init__(self,proxy_url):
self.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","None")
self.client = anthropic.Anthropic(
http_client=Client(
headers={"Invariant-Authorization": f"Bearer {invariant_api_key}"},
),
base_url=f"{proxy_url}/api/v1/proxy/{self.dataset_name}/anthropic",
)
self.get_weather_function = {
"name": "get_weather",
"description": "Get the current weather in a given location",
"input_schema": {
"type": "object",
"properties": {
"location": {
"type": "string",
"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"',
},
},
"required": ["location"],
},
}
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}]
response_list = []
while True:
response = self.client.messages.create(
tools=[self.get_weather_function],
model="claude-3-5-sonnet-20241022",
max_tokens=1024,
messages=messages
)
response_list.append(response)
# If there's tool call, Extract the tool call parameters from the response
if len(response.content) > 1 and response.content[1].type == "tool_use":
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": "user",
"content": [
{
"type": "tool_result",
"tool_use_id": tool_call_id,
"content": tool_call_result,
}
],
}
)
else:
return response_list
def get_streaming_response(self, user_query: str) -> Dict:
messages = [{"role": "user", "content": user_query}]
response_list = []
def clean_quotes(text):
# Convert \' to '
text = text.replace("\'", "'")
# Convert \" to "
text = text.replace('\"', '"')
text = text.replace("\n", " ")
return text
while True:
json_data = ""
content = []
with self.client.messages.stream(
tools=[self.get_weather_function],
model="claude-3-5-sonnet-20241022",
max_tokens=1024,
messages=messages,
) as stream:
for event in stream:
if isinstance(event, anthropic.types.RawContentBlockStartEvent):
# Start a new block
current_block = event.content_block
current_text = ""
elif isinstance(event, anthropic.types.RawContentBlockDeltaEvent):
if hasattr(event.delta, 'text'):
# Accumulate text for TextBlock
current_text += clean_quotes(event.delta.text)
elif hasattr(event.delta, 'partial_json'):
# Accumulate JSON for ToolUseBlock
json_data += clean_quotes(event.delta.partial_json)
current_text += clean_quotes(event.delta.partial_json)
elif isinstance(event, anthropic.types.RawContentBlockStopEvent):
# Block is complete, add it to content
if current_block.type == 'text':
content.append(anthropic.types.TextBlock(citations=None, text=current_text, type="text"))
elif current_block.type == 'tool_use':
content.append(
anthropic.types.ToolUseBlock(id=current_block.id,
input=json.loads(current_text),
name=current_block.name,
type="tool_use")
)
response_list.append(content)
if isinstance(event, anthropic.types.RawMessageStopEvent) and event.message.stop_reason == "tool_use":
tool_call_params = json.loads(json_data)
tool_call_result = self.get_weather(tool_call_params["location"])
messages.append({"role": "assistant", "content": content})
messages.append(
{
"role": "user",
"content": [
{
"type": "tool_result",
"tool_use_id": content[-1].id,
"content": tool_call_result,
}
],
}
)
else:
return response_list
def get_weather(self, location: str):
"""Get the current weather in a given location using latitude and longitude."""
response = f'''Weather in {location}:
Good morning! Expect overcast skies with intermittent showers throughout the day.
Temperatures will range from a cool 15°C in the early hours to around 19°C by mid-afternoon.
Light winds from the northeast at about 10 km/h will keep conditions mild.
It might be a good idea to carry an umbrella if you’re heading out. Stay dry and have a great day!
'''
return response
@pytest.mark.skipif(not os.getenv("ANTHROPIC_API_KEY"), reason="No ANTHROPIC_API_KEY set")
async def test_chat_completion_without_streaming(
context, explorer_api_url, proxy_url
):
"""Test the chat completion without streaming for the weather agent."""
weather_agent = WeatherAgent(proxy_url)
queries = [
"What's the weather like in Zurich city?",
"Tell me the weather for New York",
]
cities = ["zurich", "new york"]
# Process each query
responses = []
for index, query in enumerate(queries):
response = weather_agent.get_response(query)
assert response is not None
assert response[0].role == "assistant"
assert response[0].stop_reason == "tool_use"
assert response[0].content[0].type == "text"
assert response[0].content[1].type == "tool_use"
assert cities[index] in response[0].content[1].input["location"].lower()
assert response[1].role == "assistant"
assert response[1].stop_reason == "end_turn"
assert cities[index] in response[1].content[0].text.lower()
responses.append(response)
traces_response = await context.request.get(
f"{explorer_api_url}/api/v1/dataset/byuser/developer/{weather_agent.dataset_name}/traces"
)
traces = await traces_response.json()
assert len(traces) == len(queries)
for index,trace in enumerate(traces):
trace_id = trace["id"]
# Fetch the trace
trace_response = await context.request.get(
f"{explorer_api_url}/api/v1/trace/{trace_id}"
)
trace = await trace_response.json()
trace_messages = trace["messages"]
assert trace_messages[0]["role"] == "user"
assert trace_messages[0]["content"] == queries[index]
assert trace_messages[1]["role"] == "assistant"
assert cities[index] in trace_messages[1]["content"].lower()
assert trace_messages[2]["role"] == "assistant"
assert trace_messages[2]["tool_calls"][0]["function"]["name"] == "get_weather"
assert cities[index] in trace_messages[2]["tool_calls"][0]["function"]["arguments"]["location"].lower()
assert trace_messages[3]["role"] == "tool"
assert trace_messages[4]["role"] == "assistant"
assert cities[index] in trace_messages[4]["content"].lower()
@pytest.mark.skipif(not os.getenv("ANTHROPIC_API_KEY"), reason="No ANTHROPIC_API_KEY set")
async def test_chat_completion_with_streaming(
context, explorer_api_url, proxy_url
):
"""Test the chat completion with streaming for the weather agent."""
weather_agent = WeatherAgent(proxy_url)
queries = [
"What's the weather like in Zurich city?",
"Tell me the weather for New York",
]
cities = ["zurich", "new york"]
for index, query in enumerate(queries):
response = weather_agent.get_streaming_response(query)
assert response is not None
assert response[0][0].type == "text"
assert response[0][1].type == "tool_use"
assert response[0][1].name == "get_weather"
assert cities[index] in response[0][1].input["location"].lower()
assert response[1][0].type == "text"
assert cities[index] in response[1][0].text.lower()
traces_response = await context.request.get(
f"{explorer_api_url}/api/v1/dataset/byuser/developer/{weather_agent.dataset_name}/traces"
)
traces = await traces_response.json()
assert len(traces) == len(queries)
for index,trace in enumerate(traces):
trace_id = trace["id"]
# Fetch the trace
trace_response = await context.request.get(
f"{explorer_api_url}/api/v1/trace/{trace_id}"
)
trace = await trace_response.json()
trace_messages = trace["messages"]
assert trace_messages[0]["role"] == "user"
assert trace_messages[0]["content"] == queries[index]
assert trace_messages[1]["role"] == "assistant"
assert cities[index] in trace_messages[1]["content"].lower()
assert trace_messages[2]["role"] == "assistant"
assert trace_messages[2]["tool_calls"][0]["function"]["name"] == "get_weather"
assert cities[index] in trace_messages[2]["tool_calls"][0]["function"]["arguments"]["location"].lower()
assert trace_messages[3]["role"] == "tool"
assert trace_messages[4]["role"] == "assistant"
assert cities[index] in trace_messages[4]["content"].lower()
tests/anthropic/test_anthropic_without_tool_call.py
+137
View File
@@ -0,0 +1,137 @@
import anthropic
import os
from httpx import Client
import datetime
import pytest
import sys
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from util import * # needed for pytest fixtures
pytest_plugins = ("pytest_asyncio")
@pytest.mark.skipif(not os.getenv("ANTHROPIC_API_KEY"), reason="No ANTHROPIC_API_KEY set")
async def test_chat_completion_without_streaming(
context, explorer_api_url,proxy_url
):
dataset_name = "claude_streaming_response_without_toolcall_test" + str(datetime.datetime.now().strftime("%Y%m%d%H%M%S"))
invariant_api_key = os.environ.get("INVARIANT_API_KEY","None")
client = anthropic.Anthropic(
http_client=Client(
headers={"Invariant-Authorization": f"Bearer {invariant_api_key}"},
),
base_url=f"{proxy_url}/api/v1/proxy/{dataset_name}/anthropic",
)
cities = ["zurich", "new york", "london"]
queries = [
"Can you introduce Zurich city within 200 words?",
"Tell me the history of New York within 100 words?",
"How's the weather in London next week?"
]
# Process each query
responses = []
for query in queries:
response = client.messages.create(
model="claude-3-5-sonnet-20241022",
max_tokens=1024,
messages=[{"role": "user", "content": query}],
)
response_text = response.content[0].text
responses.append(response_text)
assert response_text is not None
assert cities[queries.index(query)] in response_text.lower()
traces_response = await context.request.get(
f"{explorer_api_url}/api/v1/dataset/byuser/developer/{dataset_name}/traces"
)
traces = await traces_response.json()
assert len(traces) == len(queries)
for index,trace in enumerate(traces):
trace_id = trace["id"]
# Fetch the trace
trace_response = await context.request.get(
f"{explorer_api_url}/api/v1/trace/{trace_id}"
)
trace = await trace_response.json()
assert trace["messages"] == [
{
"role": "user",
"content": queries[index]
},
{
"role": "assistant",
"content": responses[index]
}
]
@pytest.mark.skipif(not os.getenv("ANTHROPIC_API_KEY"), reason="No ANTHROPIC_API_KEY set")
async def test_streaming_response_without_toolcall(
context,
explorer_api_url,
proxy_url):
dataset_name = "claude_streaming_response_without_toolcall_test" + str(datetime.datetime.now().strftime("%Y%m%d%H%M%S"))
invariant_api_key = os.environ.get("INVARIANT_API_KEY","None")
client = anthropic.Anthropic(
http_client=Client(
headers={"Invariant-Authorization": f"Bearer {invariant_api_key}"},
),
base_url=f"{proxy_url}/api/v1/proxy/{dataset_name}/anthropic",
)
cities = ["zurich", "new york", "london"]
queries = [
"Can you introduce Zurich city within 200 words?",
"Tell me the history of New York within 100 words?",
"How's the weather in London next week?"
]
# Process each query
responses = []
for index,query in enumerate(queries):
messages = [
{
"role": "user",
"content": query
}
]
response_text = ""
with client.messages.stream(
model="claude-3-5-sonnet-20241022",
max_tokens=1024,
messages=messages,
) as response:
for reply in response.text_stream:
response_text += reply
assert cities[index] in response_text.lower()
responses.append(response_text)
assert response_text is not None
assert cities[queries.index(query)] in response_text.lower()
traces_response = await context.request.get(
f"{explorer_api_url}/api/v1/dataset/byuser/developer/{dataset_name}/traces"
)
traces = await traces_response.json()
assert len(traces) == len(queries)
for index,trace in enumerate(traces):
trace_id = trace["id"]
# Fetch the trace
trace_response = await context.request.get(
f"{explorer_api_url}/api/v1/trace/{trace_id}"
)
trace = await trace_response.json()
assert trace["messages"] == [
{
"role": "user",
"content": queries[index]
},
{
"role": "assistant",
"content": responses[index]
}
]
tests/anthropic/test_claude_weather_agent.py
-111
View File
@@ -1,111 +0,0 @@
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):
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": "Bearer <some-api-key>"},
),
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",
"input_schema": {
"type": "object",
"properties": {
"location": {
"type": "string",
"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"',
},
},
"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
# 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}]
while True:
response = self.client.messages.create(
# system=self.system_prompt,
tools=[self.get_weather_function],
model="claude-3-5-sonnet-20241022",
max_tokens=1024,
messages=messages,
)
# If there's tool call, Extract the tool call parameters from the response
if len(response.content) > 1 and response.content[1].type == "tool_use":
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": "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 = self.tavily_client.search(query)
response_content = response["results"][0]["content"]
return response["results"][0]["title"] + ":\n" + response_content
@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?",
]
cities = ["Zurich", "New York", "London"]
# Process each query
for index, query in enumerate(queries):
response = weather_agent.get_response(query)
assert response is not None
assert cities[index] in response
tests/open_ai/test_chat_with_tool_call.py
+1 -1
View File
@@ -7,7 +7,6 @@ import uuid
import pytest
from httpx import Client
# add tests folder (parent) to sys.path
from openai import OpenAI
@@ -205,6 +204,7 @@ async def test_chat_completion_with_tool_call_with_streaming(
model="gpt-4o", messages=history, stream=True
)
final_response = {"role": "assistant", "content": ""}
for chunk in chat_response_final:
if chunk.choices and chunk.choices[0].delta.content:
final_response["content"] += chunk.choices[0].delta.content
tests/open_ai/test_chat_without_tool_calls.py
-1
View File
@@ -6,7 +6,6 @@ import uuid
import pytest
from httpx import Client
# add tests folder (parent) to sys.path
from openai import OpenAI