From ff80bf462cb010e7b227d3e07dba72a994cb4370 Mon Sep 17 00:00:00 2001 From: Zishan Date: Wed, 5 Feb 2025 15:36:52 +0100 Subject: [PATCH] add anthropic policy --- proxy/routes/anthropic.py | 145 ++++++++++++++++++++++++++++++++++++-- proxy/routes/open_ai.py | 1 - 2 files changed, 141 insertions(+), 5 deletions(-) diff --git a/proxy/routes/anthropic.py b/proxy/routes/anthropic.py index 544a6cb..cbb5e2b 100644 --- a/proxy/routes/anthropic.py +++ b/proxy/routes/anthropic.py @@ -1,15 +1,36 @@ """Proxy service to forward requests to the Anthropic APIs""" from fastapi import APIRouter, Header, HTTPException, Depends, Request +import json +import httpx +from typing import Any +from utils.explorer import push_trace +# from .open_ai import push_to_explorer proxy = APIRouter() ALLOWED_ANTHROPIC_ENDPOINTS = {"v1/messages"} +IGNORED_HEADERS = [ + "accept-encoding", + "host", + "invariant-authorization", + "x-forwarded-for", + "x-forwarded-host", + "x-forwarded-port", + "x-forwarded-proto", + "x-forwarded-server", + "x-real-ip", +] + MISSING_INVARIANT_AUTH_HEADER = "Missing invariant-authorization header" MISSING_AUTH_HEADER = "Missing authorization header" NOT_SUPPORTED_ENDPOINT = "Not supported OpenAI endpoint" FAILED_TO_PUSH_TRACE = "Failed to push trace to the dataset: " - +END_REASONS = [ + "end_turn", + "max_tokens", + "stop_sequence" +] def validate_headers( invariant_authorization: str = Header(None), authorization: str = Header(None) @@ -34,10 +55,126 @@ async def anthropic_proxy( raise HTTPException(status_code=404, detail=NOT_SUPPORTED_ENDPOINT) headers = { - k: v for k, v in request.headers.items() + k: v for k, v in request.headers.items() if k.lower() not in IGNORED_HEADERS } - headers["accept-encoding"] = "identity" + # headers["accept-encoding"] = "identity" request_body = await request.body() - print("request_body", request_body) + request_body_json = json.loads(request_body) + + anthropic_url = f"https://api.anthropic.com/{endpoint}" + client = httpx.AsyncClient() + + 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) + await handle_non_streaming_response( + response, dataset_name, request_body_json, invariant_authorization + ) + return response.json() + +async def push_to_explorer( + dataset_name: str, + merged_response: dict[str, Any], + request_body: dict[str, Any], + 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 + messages = request_body.get("messages", []) + if merged_response is not list: + merged_response = [merged_response] + messages += merged_response + if messages[-1].get("stop_reason") in END_REASONS: + messages = anthropic_to_invariant_messages(messages) + response = await push_trace( + dataset_name=dataset_name, + messages=[messages], + invariant_authorization=invariant_authorization, + ) + +async def handle_non_streaming_response( + response: httpx.Response, + dataset_name: str, + request_body_json: dict[str, Any], + invariant_authorization: str, +): + """Handles non-streaming Anthropic responses""" + json_response = response.json() + await push_to_explorer( + dataset_name, + json_response, + request_body_json, + invariant_authorization, + ) + +def anthropic_to_invariant_messages( + messages: list[dict], keep_empty_tool_response: bool = False +) -> list[dict]: + """Converts a list of messages from the Anthropic API to the Invariant API format.""" + output = [] + 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": + 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 \ No newline at end of file diff --git a/proxy/routes/open_ai.py b/proxy/routes/open_ai.py index 34f709d..74fbb44 100644 --- a/proxy/routes/open_ai.py +++ b/proxy/routes/open_ai.py @@ -300,7 +300,6 @@ async def push_to_explorer( # Combine the messages from the request body and the choices from the OpenAI response messages = request_body.get("messages", []) messages += [choice["message"] for choice in merged_response.get("choices", [])] - _ = await push_trace( dataset_name=dataset_name, messages=[messages],