CalvinBackup/invariant-gateway
1
0
Fork 0
mirror of https://github.com/invariantlabs-ai/invariant-gateway.git synced 2026-07-06 10:57:50 +02:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #71 from invariantlabs-ai/feat/use-cli-client-name

Browse Source 
feat: use cli client and srver name
This commit is contained in:
Hemang Sarkar
2025-11-06 17:09:53 +05:30
committed by GitHub
parent f6cb73cabe 1b08b8fb66
commit 89fb100be6
9 changed files with 33 additions and 53 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+1 -5
View File
@@ -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?"}],
gateway/mcp/mcp_sessions_manager.py
+10 -5
View File
@@ -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
gateway/routes/anthropic.py
+1 -1
View File
@@ -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
resources/streaming_chunk_text/anthropic.txt
+1 -1
View File
@@ -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'
tests/integration/anthropic/test_anthropic_header_with_invariant_key.py
+1 -1
View File
@@ -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=[
{
tests/integration/anthropic/test_anthropic_with_tool_call.py
+11 -29
View File
@@ -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(
tests/integration/anthropic/test_anthropic_without_tool_call.py
+2 -2
View File
@@ -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:
tests/integration/guardrails/test_guardrails_anthropic.py
+6 -6
View File
@@ -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",
tests/integration/litellm/test_chat_without_tool_call.py
-3
View File
@@ -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,