diff --git a/proxy/requirements.txt b/proxy/requirements.txt index 193dbfb..a564b13 100644 --- a/proxy/requirements.txt +++ b/proxy/requirements.txt @@ -1,7 +1,5 @@ fastapi==0.115.7 httpx==0.28.1 -uvicorn==0.34.0 -invariant-sdk +invariant-sdk>=0.0.10 starlette-compress==1.4.0 -tavily-python -anthropic \ No newline at end of file +uvicorn==0.34.0 \ No newline at end of file diff --git a/proxy/routes/anthropic.py b/proxy/routes/anthropic.py index 1070af1..67d484e 100644 --- a/proxy/routes/anthropic.py +++ b/proxy/routes/anthropic.py @@ -1,37 +1,25 @@ """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 CLIENT_TIMEOUT, IGNORED_HEADERS from utils.explorer import push_trace from starlette.responses import StreamingResponse + 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", - "content-length" -] MISSING_INVARIANT_AUTH_HEADER = "Missing invariant-authorization header" MISSING_ANTHROPIC_AUTH_HEADER = "Missing athropic authorization header" NOT_SUPPORTED_ENDPOINT = "Not supported Anthropic 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"] + MESSAGE_START = "message_start" MESSGAE_DELTA = "message_delta" @@ -41,14 +29,15 @@ CONTENT_BLOCK_DELTA = "content_block_delta" CONTENT_BLOCK_STOP = "content_block_stop" 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)], @@ -71,13 +60,11 @@ async def anthropic_proxy( request_body_json = json.loads(request_body) anthropic_url = f"https://api.anthropic.com/{endpoint}" - client = httpx.AsyncClient(timeout=httpx.Timeout(60)) - + + client = httpx.AsyncClient(timeout=httpx.Timeout(CLIENT_TIMEOUT)) + 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") @@ -99,6 +86,7 @@ async def anthropic_proxy( ) return response.json() + async def push_to_explorer( dataset_name: str, merged_response: dict[str, Any], @@ -119,6 +107,7 @@ async def push_to_explorer( invariant_authorization=invariant_authorization, ) + async def handle_non_streaming_response( response: httpx.Response, dataset_name: str, @@ -239,6 +228,7 @@ def anthropic_to_invariant_messages( return output + def handle_user_message(message, keep_empty_tool_response): output = [] content = message["content"] @@ -257,7 +247,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"], } ) @@ -267,6 +259,7 @@ def handle_user_message(message, keep_empty_tool_response): output.append({"role": "user", "content": content}) return output + def handle_assistant_message(message): output = [] if isinstance(message["content"], list): diff --git a/proxy/routes/open_ai.py b/proxy/routes/open_ai.py index 37807f4..ba33d94 100644 --- a/proxy/routes/open_ai.py +++ b/proxy/routes/open_ai.py @@ -6,35 +6,21 @@ from typing import Any import httpx from fastapi import APIRouter, Depends, Header, HTTPException, Request, Response from starlette.responses import StreamingResponse +from utils.constants import CLIENT_TIMEOUT, 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() -MISSING_INVARIANT_AUTH_HEADER = "Missing invariant-authorization header" +MISSING_INVARIANT_AUTH_API_KEY = "Missing invariant api key" MISSING_AUTH_HEADER = "Missing authorization header" NOT_SUPPORTED_ENDPOINT = "Not supported OpenAI endpoint" FINISH_REASON_TO_PUSH_TRACE = ["stop", "length", "content_filter"] -def validate_headers( - invariant_authorization: str = Header(None), authorization: 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) +def validate_headers(authorization: str = Header(None)): + """Require the authorization header to be present""" if authorization is None: raise HTTPException(status_code=400, detail=MISSING_AUTH_HEADER) @@ -62,9 +48,32 @@ async def openai_proxy( # Check if the request is for streaming is_streaming = request_body_json.get("stream", False) - invariant_authorization = request.headers.get("invariant-authorization") - client = httpx.AsyncClient() + # The invariant-authorization header contains the Invariant API Key + # "invariant-authorization": "Bearer " + # The authorization header contains the OpenAI API Key + # "authorization": "Bearer " + # + # For some clients, it is not possible to pass a custom header + # In such cases, the Invariant API Key is passed as part of the + # authorization header with the OpenAI API key. + # The header in that case becomes: + # "authorization": "Bearer |invariant-auth: " + if request.headers.get( + "invariant-authorization" + ) is None and "|invariant-auth:" not in request.headers.get("authorization"): + raise HTTPException(status_code=400, detail=MISSING_INVARIANT_AUTH_API_KEY) + + if request.headers.get("invariant-authorization"): + invariant_authorization = request.headers.get("invariant-authorization") + else: + authorization = request.headers.get("authorization") + api_keys = authorization.split("|invariant-auth: ") + invariant_authorization = f"Bearer {api_keys[1].strip()}" + # Update the authorization header to pass the OpenAI API Key to the OpenAI API + headers["authorization"] = f"{api_keys[0].strip()}" + + client = httpx.AsyncClient(timeout=httpx.Timeout(CLIENT_TIMEOUT)) open_ai_request = client.build_request( "POST", f"https://api.openai.com/v1/{endpoint}", @@ -327,13 +336,9 @@ async def handle_non_streaming_response( dataset_name, json_response, request_body_json, invariant_authorization ) - response_headers = dict(response.headers) - response_headers.pop("Content-Encoding", None) - response_headers.pop("Content-Length", None) - return Response( content=json.dumps(json_response), status_code=response.status_code, media_type="application/json", - headers=response_headers, + headers=dict(response.headers), ) diff --git a/proxy/utils/constants.py b/proxy/utils/constants.py new file mode 100644 index 0000000..a8b61d3 --- /dev/null +++ b/proxy/utils/constants.py @@ -0,0 +1,16 @@ +"""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", + "content-length" +] + +CLIENT_TIMEOUT = 60.0 diff --git a/proxy/utils/explorer.py b/proxy/utils/explorer.py index d54b4e6..71108d7 100644 --- a/proxy/utils/explorer.py +++ b/proxy/utils/explorer.py @@ -3,53 +3,42 @@ import os from typing import Any, Dict, List -import httpx -from invariant_sdk.types.push_traces import PushTracesRequest +from invariant_sdk.async_client import AsyncClient +from invariant_sdk.types.push_traces import PushTracesRequest, PushTracesResponse DEFAULT_API_URL = "https://explorer.invariantlabs.ai" -PUSH_ENDPOINT = "/api/v1/push/trace" async def push_trace( messages: List[List[Dict[str, Any]]], dataset_name: str, invariant_authorization: str, -) -> Dict[str, str]: +) -> PushTracesResponse: """Pushes traces to the dataset on the Invariant Explorer. + If a dataset with the given name does not exist, it will be created. + Args: messages (List[List[Dict[str, Any]]]): List of messages to push. dataset_name (str): Name of the dataset. - invariant_authorization (str): Authorization token from the + invariant_authorization (str): Value of the invariant-authorization header. Returns: - Dict[str, str]: Response containing the trace ID. + PushTracesResponse: Response containing the trace ID details. """ - api_url = os.getenv("INVARIANT_API_URL", DEFAULT_API_URL).rstrip("/") # Remove any None values from the messages update_messages = [ [{k: v for k, v in msg.items() if v is not None} for msg in msg_list] for msg_list in messages ] request = PushTracesRequest(messages=update_messages, dataset=dataset_name) - async with httpx.AsyncClient() as client: - explorer_push_request = client.build_request( - "POST", - f"{api_url}{PUSH_ENDPOINT}", - json=request.to_json(), - headers={ - "Authorization": f"{invariant_authorization}", - "Accept": "application/json", - }, - ) - try: - response = await client.send(explorer_push_request) - response.raise_for_status() - return response.json() - except httpx.HTTPStatusError as e: - print(f"Failed to push trace: {e.response.text}") - return {"error": str(e)} - except Exception as e: - print(f"Unexpected error pushing trace: {str(e)}") - return {"error": str(e)} + client = AsyncClient( + api_url=os.getenv("INVARIANT_API_URL", DEFAULT_API_URL).rstrip("/"), + api_key=invariant_authorization.split("Bearer ")[1], + ) + try: + return await client.push_trace(request) + except Exception as e: + print(f"Failed to push trace: {e}") + return {"error": str(e)} diff --git a/run.sh b/run.sh index 3f20699..6561b48 100755 --- a/run.sh +++ b/run.sh @@ -62,7 +62,7 @@ tests() { # Make call to signup endpoint curl -k -X POST http://127.0.0.1/api/v1/user/signup - docker build -t 'explorer-proxy-test' -f ./tests/Dockerfile.test ./tests + docker build -t 'explorer-proxy-tests' -f ./tests/Dockerfile.test ./tests docker run \ --mount type=bind,source=./tests,target=/tests \ @@ -70,7 +70,7 @@ tests() { -e OPENAI_API_KEY="$OPENAI_API_KEY" \ -e ANTHROPIC_API_KEY="$ANTHROPIC_API_KEY"\ --env-file ./tests/.env.test \ - explorer-proxy-test $@ + explorer-proxy-tests $@ } # ----------------------------- diff --git a/tests/docker-compose.test.yml b/tests/docker-compose.test.yml index 43d6e4d..f39ef8b 100644 --- a/tests/docker-compose.test.yml +++ b/tests/docker-compose.test.yml @@ -1,4 +1,4 @@ -name: explorer-proxy-test +name: explorer-proxy-test-stack services: traefik: image: traefik:v2.0