mirror of
https://github.com/invariantlabs-ai/invariant-gateway.git
synced 2026-07-07 11:27:50 +02:00
add test for anthropic agent
This commit is contained in:
Zishan
+67
-59
@@ -23,7 +23,7 @@ IGNORED_HEADERS = [
|
||||
]
|
||||
|
||||
MISSING_INVARIANT_AUTH_HEADER = "Missing invariant-authorization header"
|
||||
MISSING_AUTH_HEADER = "Missing authorization header"
|
||||
MISSING_ANTHROPIC_AUTH_HEADER = "Missing athropic authorization header"
|
||||
NOT_SUPPORTED_ENDPOINT = "Not supported OpenAI endpoint"
|
||||
FAILED_TO_PUSH_TRACE = "Failed to push trace to the dataset: "
|
||||
END_REASONS = [
|
||||
@@ -33,13 +33,13 @@ END_REASONS = [
|
||||
]
|
||||
|
||||
def validate_headers(
|
||||
invariant_authorization: str = Header(None), authorization: str = Header(None)
|
||||
invariant_authorization: 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)
|
||||
# if authorization is None:
|
||||
# raise HTTPException(status_code=400, detail=MISSING_AUTH_HEADER)
|
||||
if x_api_key is None:
|
||||
raise HTTPException(status_code=400, detail=MISSING_ANTHROPIC_AUTH_HEADER)
|
||||
|
||||
@proxy.post(
|
||||
"/{dataset_name}/anthropic/{endpoint:path}",
|
||||
@@ -53,11 +53,9 @@ async def anthropic_proxy(
|
||||
"""Proxy calls to the Anthropic APIs"""
|
||||
if endpoint not in ALLOWED_ANTHROPIC_ENDPOINTS:
|
||||
raise HTTPException(status_code=404, detail=NOT_SUPPORTED_ENDPOINT)
|
||||
|
||||
headers = {
|
||||
k: v for k, v in request.headers.items() if k.lower() not in IGNORED_HEADERS
|
||||
}
|
||||
# headers["accept-encoding"] = "identity"
|
||||
|
||||
request_body = await request.body()
|
||||
|
||||
@@ -89,14 +87,16 @@ async def push_to_explorer(
|
||||
invariant_authorization: str,
|
||||
) -> None:
|
||||
"""Pushes the full trace to the Invariant Explorer"""
|
||||
# Combine the messages from the request body and the choices from the OpenAI response
|
||||
# Combine the messages from the request body and Anthropic response
|
||||
messages = request_body.get("messages", [])
|
||||
if merged_response is not list:
|
||||
merged_response = [merged_response]
|
||||
messages += merged_response
|
||||
|
||||
# Only push the trace to explorer if the last message is an end turn message
|
||||
if messages[-1].get("stop_reason") in END_REASONS:
|
||||
messages = anthropic_to_invariant_messages(messages)
|
||||
response = await push_trace(
|
||||
_ = await push_trace(
|
||||
dataset_name=dataset_name,
|
||||
messages=[messages],
|
||||
invariant_authorization=invariant_authorization,
|
||||
@@ -122,59 +122,67 @@ def anthropic_to_invariant_messages(
|
||||
) -> list[dict]:
|
||||
"""Converts a list of messages from the Anthropic API to the Invariant API format."""
|
||||
output = []
|
||||
role_mapping = {
|
||||
"system": lambda msg: {"role": "system", "content": msg["content"]},
|
||||
"user": lambda msg: handle_user_message(msg, keep_empty_tool_response),
|
||||
"assistant": lambda msg: handle_assistant_message(msg),
|
||||
}
|
||||
|
||||
for message in messages:
|
||||
if message["role"] == "system":
|
||||
output.append({"role": "system", "content": message["content"]})
|
||||
if message["role"] == "user":
|
||||
if isinstance(message["content"], list):
|
||||
for sub_message in message["content"]:
|
||||
if sub_message["type"] == "tool_result":
|
||||
if sub_message["content"]:
|
||||
output.append(
|
||||
{
|
||||
"role": "tool",
|
||||
"content": sub_message["content"],
|
||||
"tool_id": sub_message["tool_use_id"],
|
||||
}
|
||||
)
|
||||
else:
|
||||
if keep_empty_tool_response and any(
|
||||
[sub_message[k] for k in sub_message]
|
||||
):
|
||||
output.append(
|
||||
{
|
||||
"role": "tool",
|
||||
"content": {"is_error": True}
|
||||
if sub_message["is_error"]
|
||||
else {},
|
||||
"tool_id": sub_message["tool_use_id"],
|
||||
}
|
||||
)
|
||||
elif sub_message["type"] == "text":
|
||||
output.append({"role": "user", "content": sub_message["text"]})
|
||||
else:
|
||||
output.append({"role": "user", "content": message["content"]})
|
||||
if message["role"] == "assistant":
|
||||
for sub_message in message["content"]:
|
||||
if sub_message["type"] == "text":
|
||||
output.append(
|
||||
{"role": "assistant", "content": sub_message.get("text")}
|
||||
)
|
||||
if sub_message["type"] == "tool_use":
|
||||
handler = role_mapping.get(message["role"])
|
||||
if handler:
|
||||
output.extend(handler(message))
|
||||
|
||||
return output
|
||||
|
||||
def handle_user_message(message, keep_empty_tool_response):
|
||||
output = []
|
||||
content = message["content"]
|
||||
if isinstance(content, list):
|
||||
for sub_message in content:
|
||||
if sub_message["type"] == "tool_result":
|
||||
if sub_message["content"]:
|
||||
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"),
|
||||
},
|
||||
}
|
||||
],
|
||||
"role": "tool",
|
||||
"content": sub_message["content"],
|
||||
"tool_id": sub_message["tool_use_id"],
|
||||
}
|
||||
)
|
||||
return output
|
||||
elif keep_empty_tool_response and any(sub_message.values()):
|
||||
output.append(
|
||||
{
|
||||
"role": "tool",
|
||||
"content": {"is_error": True} if sub_message["is_error"] else {},
|
||||
"tool_id": sub_message["tool_use_id"],
|
||||
}
|
||||
)
|
||||
elif sub_message["type"] == "text":
|
||||
output.append({"role": "user", "content": sub_message["text"]})
|
||||
else:
|
||||
output.append({"role": "user", "content": content})
|
||||
return output
|
||||
|
||||
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"),
|
||||
},
|
||||
}
|
||||
],
|
||||
}
|
||||
)
|
||||
return output
|
||||
|
||||
Reference in New Issue
Block a user