diff --git a/proxy/routes/anthropic.py b/proxy/routes/anthropic.py index 86706f9..f2f0a2b 100644 --- a/proxy/routes/anthropic.py +++ b/proxy/routes/anthropic.py @@ -1,50 +1,38 @@ """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 + +import httpx +from fastapi import APIRouter, Depends, Header, HTTPException, Request +from utils.constants import IGNORED_HEADERS 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_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 = [ - "end_turn", - "max_tokens", - "stop_sequence" -] +END_REASONS = ["end_turn", "max_tokens", "stop_sequence"] + def validate_headers( - invariant_authorization: str = Header(None), x_api_key: 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 x_api_key is None: raise HTTPException(status_code=400, detail=MISSING_ANTHROPIC_AUTH_HEADER) - + + @proxy.post( "/{dataset_name}/anthropic/{endpoint:path}", dependencies=[Depends(validate_headers)], -) +) async def anthropic_proxy( dataset_name: str, endpoint: str, @@ -65,10 +53,7 @@ async def anthropic_proxy( client = httpx.AsyncClient() anthropic_request = client.build_request( - "POST", - anthropic_url, - headers=headers, - data=request_body + "POST", anthropic_url, headers=headers, data=request_body ) invariant_authorization = request.headers.get("invariant-authorization") @@ -80,6 +65,7 @@ async def anthropic_proxy( ) return response.json() + async def push_to_explorer( dataset_name: str, merged_response: dict[str, Any], @@ -98,6 +84,7 @@ async def push_to_explorer( invariant_authorization=invariant_authorization, ) + async def handle_non_streaming_response( response: httpx.Response, dataset_name: str, @@ -115,6 +102,7 @@ async def handle_non_streaming_response( invariant_authorization, ) + def anthropic_to_invariant_messages( messages: list[dict], keep_empty_tool_response: bool = False ) -> list[dict]: @@ -133,6 +121,7 @@ def anthropic_to_invariant_messages( return output + def handle_user_message(message, keep_empty_tool_response): output = [] content = message["content"] @@ -151,7 +140,9 @@ def handle_user_message(message, keep_empty_tool_response): output.append( { "role": "tool", - "content": {"is_error": True} if sub_message["is_error"] else {}, + "content": {"is_error": True} + if sub_message["is_error"] + else {}, "tool_id": sub_message["tool_use_id"], } ) @@ -161,6 +152,7 @@ def handle_user_message(message, keep_empty_tool_response): output.append({"role": "user", "content": content}) return output + def handle_assistant_message(message): output = [] for sub_message in message["content"]: diff --git a/proxy/routes/open_ai.py b/proxy/routes/open_ai.py index 907638d..12c6b56 100644 --- a/proxy/routes/open_ai.py +++ b/proxy/routes/open_ai.py @@ -6,20 +6,10 @@ from typing import Any import httpx from fastapi import APIRouter, Depends, Header, HTTPException, Request, Response from starlette.responses import StreamingResponse +from utils.constants import IGNORED_HEADERS from utils.explorer import push_trace ALLOWED_OPEN_AI_ENDPOINTS = {"chat/completions"} -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", -] proxy = APIRouter() diff --git a/proxy/utils/constants.py b/proxy/utils/constants.py new file mode 100644 index 0000000..f690498 --- /dev/null +++ b/proxy/utils/constants.py @@ -0,0 +1,13 @@ +"""Common constants used in the proxy.""" + +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", +]