diff --git a/README.md b/README.md index f50bf8a..d2cd653 100644 --- a/README.md +++ b/README.md @@ -160,15 +160,11 @@ import random import os base_url = "/api/v1/gateway/litellm/{add-your-dataset-name-here}" -EXAMPLE_MODELS = ["openai/gpt-4o", "gemini/gemini-2.0-flash", "anthropic/claude-3-5-haiku-20241022"] +EXAMPLE_MODELS = ["openai/gpt-4o", "anthropic/claude-3-5-haiku-20241022"] model = random.choice(SAMPLE_MODELS) base_url += "/" + model.split("/")[0] # append /gemini /openai or /anthropic. -if model.split("/")[0] == "gemini": - base_url += f"/v1beta/models/{model.split('/')[1]}" # gemini expects the model name in the url. - - chat_response = completion( model=model, messages=[{"role": "user", "content": "What is the capital of France?"}], diff --git a/gateway/mcp/mcp_sessions_manager.py b/gateway/mcp/mcp_sessions_manager.py index ece6a3e..d56b2f0 100644 --- a/gateway/mcp/mcp_sessions_manager.py +++ b/gateway/mcp/mcp_sessions_manager.py @@ -32,6 +32,7 @@ def user_and_host() -> str: return f"{username}@{hostname}" + class McpAttributes(BaseModel): """ A Pydantic model to represent MCP attributes. @@ -118,9 +119,9 @@ class McpAttributes(BaseModel): for arg in extra_args: assert "=" in arg, f"Invalid extra metadata argument: {arg}" key, value = arg.split("=") - assert key.startswith( - "--metadata-" - ), f"Invalid extra metadata argument: {arg}, must start with --metadata-" + assert key.startswith("--metadata-"), ( + f"Invalid extra metadata argument: {arg}, must start with --metadata-" + ) key = key[len("--metadata-") :] metadata[key] = value @@ -185,8 +186,12 @@ class McpSession(BaseModel): self.attributes.explorer_dataset, self._get_invariant_authorization(), # pylint: disable=no-member - self.attributes.metadata.get("mcp_client"), - self.attributes.metadata.get("mcp_server"), + self.attributes.metadata.get( + "client", self.attributes.metadata.get("mcp_client") + ), + self.attributes.metadata.get( + "server", self.attributes.metadata.get("mcp_server") + ), ) @contextlib.asynccontextmanager diff --git a/gateway/routes/anthropic.py b/gateway/routes/anthropic.py index b0e4904..e706ecb 100644 --- a/gateway/routes/anthropic.py +++ b/gateway/routes/anthropic.py @@ -263,7 +263,7 @@ class AnthropicProvider(BaseProvider): b'event: message_start\ndata: {"type":"message_start","message": {"id":"msg_01LkayzAaw7b7QkUAw91psyx","type":"message","role":"assistant" - ,"model":"claude-3-5-sonnet-20241022","content":[],"stop_reason":null, + ,"model":"claude-sonnet-4-5-20250929","content":[],"stop_reason":null, "stop_sequence":null,"usage":{"input_tokens":20,"cache_creation_input_to' and diff --git a/resources/streaming_chunk_text/anthropic.txt b/resources/streaming_chunk_text/anthropic.txt index c8c4480..8252d32 100644 --- a/resources/streaming_chunk_text/anthropic.txt +++ b/resources/streaming_chunk_text/anthropic.txt @@ -1,4 +1,4 @@ -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' +anthropic_chunk_list = [b'event: message_start\ndata: {"type":"message_start","message":{"id":"msg_012KWB6kiKvzx7r1SKs5nGA1","type":"message","role":"assistant","model":"claude-sonnet-4-5-20250929","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' diff --git a/tests/integration/anthropic/test_anthropic_header_with_invariant_key.py b/tests/integration/anthropic/test_anthropic_header_with_invariant_key.py index 1597846..1b43198 100644 --- a/tests/integration/anthropic/test_anthropic_header_with_invariant_key.py +++ b/tests/integration/anthropic/test_anthropic_header_with_invariant_key.py @@ -37,7 +37,7 @@ async def test_gateway_with_invariant_key_in_anthropic_key_header( base_url=f"{gateway_url}/api/v1/gateway/{dataset_name}/anthropic", ) response = client.messages.create( - model="claude-3-5-sonnet-20241022", + model="claude-sonnet-4-5-20250929", max_tokens=1024, messages=[ { diff --git a/tests/integration/anthropic/test_anthropic_with_tool_call.py b/tests/integration/anthropic/test_anthropic_with_tool_call.py index 0af1002..d3eadf7 100644 --- a/tests/integration/anthropic/test_anthropic_with_tool_call.py +++ b/tests/integration/anthropic/test_anthropic_with_tool_call.py @@ -56,7 +56,7 @@ class WeatherAgent: while True: response = self.client.messages.create( tools=[self.get_weather_function], - model="claude-3-5-sonnet-20241022", + model="claude-sonnet-4-5-20250929", max_tokens=1024, messages=messages, ) @@ -100,7 +100,7 @@ class WeatherAgent: event = None with self.client.messages.stream( tools=[self.get_weather_function], - model="claude-3-5-sonnet-20241022", + model="claude-sonnet-4-5-20250929", max_tokens=1024, messages=messages, ) as stream: @@ -188,12 +188,8 @@ async def test_response_with_tool_call(explorer_api_url, gateway_url, push_to_ex 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 city in response[0].content[1].input["location"].lower() - - assert response[1].role == "assistant" - assert response[1].stop_reason == "end_turn" + assert response[0].content[0].type == "tool_use" + assert city in response[0].content[0].input["location"].lower() responses.append(response) if push_to_explorer: @@ -217,18 +213,13 @@ async def test_response_with_tool_call(explorer_api_url, gateway_url, push_to_ex assert trace_messages[0]["role"] == "user" assert trace_messages[0]["content"] == query assert trace_messages[1]["role"] == "assistant" - assert city in trace_messages[1]["content"].lower() - assert trace_messages[2]["role"] == "assistant" - assert trace_messages[2]["tool_calls"][0]["function"]["name"] == "get_weather" + assert trace_messages[1]["tool_calls"][0]["function"]["name"] == "get_weather" assert ( city - in trace_messages[2]["tool_calls"][0]["function"]["arguments"][ + in trace_messages[1]["tool_calls"][0]["function"]["arguments"][ "location" ].lower() ) - assert trace_messages[3]["role"] == "tool" - assert trace_messages[4]["role"] == "assistant" - assert city in trace_messages[4]["content"].lower() @pytest.mark.skipif( @@ -249,13 +240,9 @@ async def test_streaming_response_with_tool_call( if len(response) == 2: 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 city in response[0][1].input["location"].lower() - - assert response[1][0].type == "text" - assert city in response[1][0].text.lower() + assert response[0][0].type == "tool_use" + assert response[0][0].name == "get_weather" + assert city in response[0][0].input["location"].lower() elif len(response) == 1: # expected output in this case is something like this: # [[TextBlock(text="I'll help you check the weather in New York using the get_weather function.", type='text', citations=None), ToolUseBlock(id='toolu_019VZsmxuUhShou2EpPBxvpe', input={'location': 'New York, NY', 'unit': 'celsius'}, name='get_weather', type='tool_use')]] @@ -290,18 +277,13 @@ async def test_streaming_response_with_tool_call( assert trace_messages[0]["role"] == "user" assert trace_messages[0]["content"] == query assert trace_messages[1]["role"] == "assistant" - assert city in trace_messages[1]["content"].lower() - assert trace_messages[2]["role"] == "assistant" - assert trace_messages[2]["tool_calls"][0]["function"]["name"] == "get_weather" + assert trace_messages[1]["tool_calls"][0]["function"]["name"] == "get_weather" assert ( city - in trace_messages[2]["tool_calls"][0]["function"]["arguments"][ + in trace_messages[1]["tool_calls"][0]["function"]["arguments"][ "location" ].lower() ) - assert trace_messages[3]["role"] == "tool" - assert trace_messages[4]["role"] == "assistant" - assert city in trace_messages[4]["content"].lower() @pytest.mark.skipif( diff --git a/tests/integration/anthropic/test_anthropic_without_tool_call.py b/tests/integration/anthropic/test_anthropic_without_tool_call.py index 727e76a..82cd9c3 100644 --- a/tests/integration/anthropic/test_anthropic_without_tool_call.py +++ b/tests/integration/anthropic/test_anthropic_without_tool_call.py @@ -42,7 +42,7 @@ async def test_response_without_tool_call( responses = [] for query in queries: response = client.messages.create( - model="claude-3-5-sonnet-20241022", + model="claude-sonnet-4-5-20250929", max_tokens=1024, messages=[{"role": "user", "content": query}], ) @@ -100,7 +100,7 @@ async def test_streaming_response_without_tool_call( response_text = "" with client.messages.stream( - model="claude-3-5-sonnet-20241022", + model="claude-sonnet-4-5-20250929", max_tokens=200, messages=messages, ) as response: diff --git a/tests/integration/guardrails/test_guardrails_anthropic.py b/tests/integration/guardrails/test_guardrails_anthropic.py index 77bcf1f..e5b0ea1 100644 --- a/tests/integration/guardrails/test_guardrails_anthropic.py +++ b/tests/integration/guardrails/test_guardrails_anthropic.py @@ -39,7 +39,7 @@ async def test_message_content_guardrail_from_file( ) request = { - "model": "claude-3-5-sonnet-20241022", + "model": "claude-sonnet-4-5-20250929", "max_tokens": 100, "messages": [{"role": "user", "content": "What is the capital of Spain?"}], } @@ -150,7 +150,7 @@ async def test_tool_call_guardrail_from_file( ], "tools": tools, "system": system_message, - "model": "claude-3-5-sonnet-20241022", + "model": "claude-sonnet-4-5-20250929", "max_tokens": 150, } @@ -250,7 +250,7 @@ async def test_input_from_guardrail_from_file( ) request = { - "model": "claude-3-5-sonnet-20241022", + "model": "claude-sonnet-4-5-20250929", "max_tokens": 100, "messages": [{"role": "user", "content": "Tell me more about Fight Club."}], } @@ -352,7 +352,7 @@ async def test_with_guardrails_from_explorer(explorer_api_url, gateway_url, do_s # This should not be blocked by the guardrails from the explorer when we push to explorer # because the file based guardrails are overridden by the explorer guardrails spain_request = { - "model": "claude-3-5-sonnet-20241022", + "model": "claude-sonnet-4-5-20250929", "messages": [{"role": "user", "content": "What is the capital of Spain?"}], "max_tokens": 100, } @@ -379,7 +379,7 @@ async def test_with_guardrails_from_explorer(explorer_api_url, gateway_url, do_s # This should be blocked by the guardrails from the explorer user_prompt = "What kind of a creature is Shrek? What is his Shrek's wife's name? Only answer these questions with single sentences, don't add any extra details." shrek_request = { - "model": "claude-3-5-sonnet-20241022", + "model": "claude-sonnet-4-5-20250929", "messages": [ { "role": "user", @@ -497,7 +497,7 @@ async def test_preguardrailing_with_guardrails_from_explorer( user_prompt = "Tell me a one sentence pun." request = { - "model": "claude-3-5-sonnet-20241022", + "model": "claude-sonnet-4-5-20250929", "messages": [ { "role": "user", diff --git a/tests/integration/litellm/test_chat_without_tool_call.py b/tests/integration/litellm/test_chat_without_tool_call.py index ee831a9..caa8d05 100644 --- a/tests/integration/litellm/test_chat_without_tool_call.py +++ b/tests/integration/litellm/test_chat_without_tool_call.py @@ -10,7 +10,6 @@ from litellm import completion MODEL_API_KEYS = { "openai/gpt-4o": "OPENAI_API_KEY", - "gemini/gemini-2.0-flash": "GEMINI_API_KEY", "anthropic/claude-3-5-haiku-20241022": "ANTHROPIC_API_KEY", } @@ -46,8 +45,6 @@ async def test_chat_completion( ) base_url += "/" + litellm_model.split("/")[0] # add provider name - if litellm_model.split("/")[0] == "gemini": - base_url += f"/v1beta/models/{litellm_model.split('/')[1]}" # gemini expects the model name in the url chat_response = completion( model=litellm_model,