diff --git a/.github/workflows/tests_ci.yml b/.github/workflows/tests_ci.yml index fa97fb5..f5c135a 100644 --- a/.github/workflows/tests_ci.yml +++ b/.github/workflows/tests_ci.yml @@ -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 diff --git a/proxy/routes/anthropic.py b/proxy/routes/anthropic.py index b407c6b..fd0f0fa 100644 --- a/proxy/routes/anthropic.py +++ b/proxy/routes/anthropic.py @@ -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 diff --git a/proxy/utils/explorer.py b/proxy/utils/explorer.py index 71108d7..1d6a18f 100644 --- a/proxy/utils/explorer.py +++ b/proxy/utils/explorer.py @@ -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, diff --git a/resources/streaming_chunk_text/anthropic.txt b/resources/streaming_chunk_text/anthropic.txt new file mode 100644 index 0000000..c8c4480 --- /dev/null +++ b/resources/streaming_chunk_text/anthropic.txt @@ -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'] diff --git a/run.sh b/run.sh index 226f04c..01f1bb4 100755 --- a/run.sh +++ b/run.sh @@ -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 $@ } diff --git a/tests/anthropic/test_anthropic_with_tool_call.py b/tests/anthropic/test_anthropic_with_tool_call.py new file mode 100644 index 0000000..56be242 --- /dev/null +++ b/tests/anthropic/test_anthropic_with_tool_call.py @@ -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() + diff --git a/tests/anthropic/test_anthropic_without_tool_call.py b/tests/anthropic/test_anthropic_without_tool_call.py new file mode 100644 index 0000000..91995f6 --- /dev/null +++ b/tests/anthropic/test_anthropic_without_tool_call.py @@ -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] + } + ] \ No newline at end of file diff --git a/tests/anthropic/test_claude_weather_agent.py b/tests/anthropic/test_claude_weather_agent.py deleted file mode 100644 index 9315225..0000000 --- a/tests/anthropic/test_claude_weather_agent.py +++ /dev/null @@ -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 "}, - ), - 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 diff --git a/tests/open_ai/test_chat_with_tool_call.py b/tests/open_ai/test_chat_with_tool_call.py index 422e5ff..72f7f8f 100644 --- a/tests/open_ai/test_chat_with_tool_call.py +++ b/tests/open_ai/test_chat_with_tool_call.py @@ -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 diff --git a/tests/open_ai/test_chat_without_tool_calls.py b/tests/open_ai/test_chat_without_tool_calls.py index 130b597..b96ae43 100644 --- a/tests/open_ai/test_chat_without_tool_calls.py +++ b/tests/open_ai/test_chat_without_tool_calls.py @@ -6,7 +6,6 @@ import uuid import pytest from httpx import Client - # add tests folder (parent) to sys.path from openai import OpenAI