From d7bcfde3934bf0a57d6c7d0514b023fa637e055d Mon Sep 17 00:00:00 2001 From: Hemang Date: Thu, 13 Feb 2025 14:07:17 +0100 Subject: [PATCH 1/6] Remove response_headers.pop(...) in handle_non_streaming_response. --- proxy/routes/open_ai.py | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/proxy/routes/open_ai.py b/proxy/routes/open_ai.py index 37807f4..907638d 100644 --- a/proxy/routes/open_ai.py +++ b/proxy/routes/open_ai.py @@ -327,13 +327,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), ) From 2789ec2b98496cfd8fba729875026b6c961a3f3a Mon Sep 17 00:00:00 2001 From: Hemang Date: Thu, 13 Feb 2025 14:11:47 +0100 Subject: [PATCH 2/6] Use the AsyncClient from the sdk to push traces to Explorer. --- proxy/requirements.txt | 6 ++---- proxy/utils/explorer.py | 43 +++++++++++++++-------------------------- 2 files changed, 18 insertions(+), 31 deletions(-) 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/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)} From 1a7cebbb52b09291d319bfe11dd304f8500cfaa0 Mon Sep 17 00:00:00 2001 From: Hemang Date: Thu, 13 Feb 2025 14:25:20 +0100 Subject: [PATCH 3/6] Moves ignored_headers to constants and reformatted anthropic.py --- proxy/routes/anthropic.py | 46 ++++++++++++++++----------------------- proxy/routes/open_ai.py | 12 +--------- proxy/utils/constants.py | 13 +++++++++++ 3 files changed, 33 insertions(+), 38 deletions(-) create mode 100644 proxy/utils/constants.py 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", +] From 891abf29524ea195cb175be4456c49e318f53765 Mon Sep 17 00:00:00 2001 From: Hemang Date: Thu, 13 Feb 2025 15:46:09 +0100 Subject: [PATCH 4/6] Change some names for the test stack. --- run.sh | 4 ++-- tests/docker-compose.test.yml | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/run.sh b/run.sh index 7b7a251..ed995ee 100755 --- a/run.sh +++ b/run.sh @@ -62,14 +62,14 @@ 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 \ --network invariant-proxy-web-test \ -e OPENAI_API_KEY="$OPENAI_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 From 4890aa56a64332ccf5550c551caaf0d445177bc2 Mon Sep 17 00:00:00 2001 From: Hemang Date: Fri, 14 Feb 2025 17:23:30 +0100 Subject: [PATCH 5/6] Add timeout to httpx.AsyncClient --- proxy/routes/anthropic.py | 4 ++-- proxy/routes/open_ai.py | 4 ++-- proxy/utils/constants.py | 2 ++ 3 files changed, 6 insertions(+), 4 deletions(-) diff --git a/proxy/routes/anthropic.py b/proxy/routes/anthropic.py index f2f0a2b..b407c6b 100644 --- a/proxy/routes/anthropic.py +++ b/proxy/routes/anthropic.py @@ -5,7 +5,7 @@ from typing import Any import httpx from fastapi import APIRouter, Depends, Header, HTTPException, Request -from utils.constants import IGNORED_HEADERS +from utils.constants import CLIENT_TIMEOUT, IGNORED_HEADERS from utils.explorer import push_trace proxy = APIRouter() @@ -50,7 +50,7 @@ async def anthropic_proxy( request_body_json = json.loads(request_body) anthropic_url = f"https://api.anthropic.com/{endpoint}" - client = httpx.AsyncClient() + client = httpx.AsyncClient(timeout=httpx.Timeout(CLIENT_TIMEOUT)) anthropic_request = client.build_request( "POST", anthropic_url, headers=headers, data=request_body diff --git a/proxy/routes/open_ai.py b/proxy/routes/open_ai.py index 12c6b56..8bce55b 100644 --- a/proxy/routes/open_ai.py +++ b/proxy/routes/open_ai.py @@ -6,7 +6,7 @@ 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.constants import CLIENT_TIMEOUT, IGNORED_HEADERS from utils.explorer import push_trace ALLOWED_OPEN_AI_ENDPOINTS = {"chat/completions"} @@ -54,7 +54,7 @@ async def openai_proxy( is_streaming = request_body_json.get("stream", False) invariant_authorization = request.headers.get("invariant-authorization") - client = httpx.AsyncClient() + client = httpx.AsyncClient(timeout=httpx.Timeout(CLIENT_TIMEOUT)) open_ai_request = client.build_request( "POST", f"https://api.openai.com/v1/{endpoint}", diff --git a/proxy/utils/constants.py b/proxy/utils/constants.py index f690498..74dd72e 100644 --- a/proxy/utils/constants.py +++ b/proxy/utils/constants.py @@ -11,3 +11,5 @@ IGNORED_HEADERS = [ "x-forwarded-server", "x-real-ip", ] + +CLIENT_TIMEOUT = 60.0 \ No newline at end of file From df2c2fe8213f47d729d3595c47bced44a9f2d2a3 Mon Sep 17 00:00:00 2001 From: Hemang Date: Mon, 17 Feb 2025 11:07:28 +0100 Subject: [PATCH 6/6] For some users like OpenHands, it is not possible to pass in a custom header containing the Invariant API Key. In that case, modify the authorization header to include both the OpenAI API Key and the Invariant API Key. --- proxy/routes/open_ai.py | 35 +++++++++++++++++++++++++++-------- proxy/utils/constants.py | 2 +- 2 files changed, 28 insertions(+), 9 deletions(-) diff --git a/proxy/routes/open_ai.py b/proxy/routes/open_ai.py index 8bce55b..ba33d94 100644 --- a/proxy/routes/open_ai.py +++ b/proxy/routes/open_ai.py @@ -13,18 +13,14 @@ ALLOWED_OPEN_AI_ENDPOINTS = {"chat/completions"} 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) @@ -52,7 +48,30 @@ 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") + + # 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( diff --git a/proxy/utils/constants.py b/proxy/utils/constants.py index 74dd72e..858a971 100644 --- a/proxy/utils/constants.py +++ b/proxy/utils/constants.py @@ -12,4 +12,4 @@ IGNORED_HEADERS = [ "x-real-ip", ] -CLIENT_TIMEOUT = 60.0 \ No newline at end of file +CLIENT_TIMEOUT = 60.0