Merge branch 'main' into anthropic-implement

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
+uvicorn==0.34.0

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):

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 <Invariant API Key>"
+    # The authorization header contains the OpenAI API Key
+    # "authorization": "Bearer <OpenAI API Key>"
+    #
+    # 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 <OpenAI API Key>|invariant-auth: <Invariant API Key>"
+    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),
     )

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

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)}

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 $@
 }
 
 # -----------------------------

tests/docker-compose.test.yml

@@ -1,4 +1,4 @@
-name: explorer-proxy-test
+name: explorer-proxy-test-stack
 services:
   traefik:
     image: traefik:v2.0