diff --git a/proxy/routes/anthropic.py b/proxy/routes/anthropic.py index f62c044..801f7cd 100644 --- a/proxy/routes/anthropic.py +++ b/proxy/routes/anthropic.py @@ -7,7 +7,11 @@ import httpx from common.config_manager import ProxyConfig, ProxyConfigManager from fastapi import APIRouter, Depends, Header, HTTPException, Request, Response from starlette.responses import StreamingResponse -from utils.constants import CLIENT_TIMEOUT, IGNORED_HEADERS +from utils.constants import ( + CLIENT_TIMEOUT, + IGNORED_HEADERS, + INVARIANT_AUTHORIZATION_HEADER, +) from utils.explorer import push_trace proxy = APIRouter() @@ -24,7 +28,7 @@ CONTENT_BLOCK_START = "content_block_start" CONTENT_BLOCK_DELTA = "content_block_delta" CONTENT_BLOCK_STOP = "content_block_stop" -HEADER_AUTHORIZATION = "x-api-key" +ANTHROPIC_AUTHORIZATION_HEADER = "x-api-key" def validate_headers(x_api_key: str = Header(None)): @@ -43,7 +47,7 @@ def validate_headers(x_api_key: str = Header(None)): ) async def anthropic_v1_messages_proxy( request: Request, - dataset_name: str = None, + dataset_name: str = None, # This is None if the client doesn't want to push to Explorer config: ProxyConfig = Depends(ProxyConfigManager.get_config), # pylint: disable=unused-argument ): """Proxy calls to the Anthropic APIs""" @@ -51,19 +55,26 @@ async def anthropic_v1_messages_proxy( k: v for k, v in request.headers.items() if k.lower() not in IGNORED_HEADERS } headers["accept-encoding"] = "identity" - if request.headers.get( - "invariant-authorization" - ) is None and "|invariant-auth:" not in request.headers.get(HEADER_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(HEADER_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[HEADER_AUTHORIZATION] = f"{api_keys[0].strip()}" + # In case the user wants to push to Explorer, the request must contain the Invariant API Key + invariant_authorization = None + if dataset_name: + if request.headers.get( + INVARIANT_AUTHORIZATION_HEADER + ) is None and "|invariant-auth:" not in request.headers.get( + ANTHROPIC_AUTHORIZATION_HEADER + ): + raise HTTPException(status_code=400, detail=MISSING_INVARIANT_AUTH_API_KEY) + if request.headers.get(INVARIANT_AUTHORIZATION_HEADER): + invariant_authorization = request.headers.get( + INVARIANT_AUTHORIZATION_HEADER + ) + else: + header_value = request.headers.get(ANTHROPIC_AUTHORIZATION_HEADER) + api_keys = header_value.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[ANTHROPIC_AUTHORIZATION_HEADER] = f"{api_keys[0].strip()}" request_body = await request.body() @@ -109,9 +120,9 @@ async def push_to_explorer( async def handle_non_streaming_response( response: httpx.Response, - dataset_name: str, + dataset_name: Optional[str], request_body_json: dict[str, Any], - invariant_authorization: str, + invariant_authorization: Optional[str], ) -> Response: """Handles non-streaming Anthropic responses""" try: @@ -146,7 +157,7 @@ async def handle_streaming_response( client: httpx.AsyncClient, anthropic_request: httpx.Request, dataset_name: Optional[str], - invariant_authorization: str, + invariant_authorization: Optional[str], ) -> StreamingResponse: """Handles streaming Anthropic responses""" merged_response = [] diff --git a/proxy/routes/gemini.py b/proxy/routes/gemini.py index da6d060..7a5d42d 100644 --- a/proxy/routes/gemini.py +++ b/proxy/routes/gemini.py @@ -19,14 +19,14 @@ async def gemini_generate_content_proxy( api_version: str, model: str, endpoint: str, - dataset_name: str = None, + dataset_name: str = None, # This is None if the client doesn't want to push to Explorer alt: str = Query( None, title="Response Format", description="Set to 'sse' for streaming" ), config: ProxyConfig = Depends(ProxyConfigManager.get_config), # pylint: disable=unused-argument ) -> Response: """Proxy calls to the Gemini GenerateContent API""" - if "generateContent" != endpoint and "streamGenerateContent" != endpoint: + if endpoint not in ["generateContent", "streamGenerateContent"]: return Response( content="Invalid endpoint - the only endpoints supported are: \ /api/v1/proxy/gemini//models/:generateContent or \ diff --git a/proxy/routes/open_ai.py b/proxy/routes/open_ai.py index 8053b71..631f982 100644 --- a/proxy/routes/open_ai.py +++ b/proxy/routes/open_ai.py @@ -7,7 +7,11 @@ import httpx from common.config_manager import ProxyConfig, ProxyConfigManager from fastapi import APIRouter, Depends, Header, HTTPException, Request, Response from fastapi.responses import StreamingResponse -from utils.constants import CLIENT_TIMEOUT, IGNORED_HEADERS +from utils.constants import ( + CLIENT_TIMEOUT, + IGNORED_HEADERS, + INVARIANT_AUTHORIZATION_HEADER, +) from utils.explorer import push_trace proxy = APIRouter() @@ -15,6 +19,7 @@ proxy = APIRouter() MISSING_INVARIANT_AUTH_API_KEY = "Missing invariant api key" MISSING_AUTH_HEADER = "Missing authorization header" FINISH_REASON_TO_PUSH_TRACE = ["stop", "length", "content_filter"] +OPENAI_AUTHORIZATION_HEADER = "authorization" def validate_headers(authorization: str = Header(None)): @@ -33,7 +38,7 @@ def validate_headers(authorization: str = Header(None)): ) async def openai_chat_completions_proxy( request: Request, - dataset_name: str = None, + dataset_name: str = None, # This is None if the client doesn't want to push to Explorer config: ProxyConfig = Depends(ProxyConfigManager.get_config), # pylint: disable=unused-argument ) -> Response: """Proxy calls to the OpenAI APIs""" @@ -48,6 +53,7 @@ async def openai_chat_completions_proxy( # Check if the request is for streaming is_streaming = request_body_json.get("stream", False) + # In case the user wants to push to Explorer, the request must contain the Invariant API Key # The invariant-authorization header contains the Invariant API Key # "invariant-authorization": "Bearer " # The authorization header contains the OpenAI API Key @@ -58,19 +64,25 @@ async def openai_chat_completions_proxy( # authorization header with the OpenAI API key. # The header in that case becomes: # "authorization": "|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) + invariant_authorization = None + if dataset_name: + if request.headers.get( + INVARIANT_AUTHORIZATION_HEADER + ) is None and "|invariant-auth:" not in request.headers.get( + OPENAI_AUTHORIZATION_HEADER + ): + 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()}" + if request.headers.get(INVARIANT_AUTHORIZATION_HEADER): + invariant_authorization = request.headers.get( + INVARIANT_AUTHORIZATION_HEADER + ) + else: + header_value = request.headers.get(OPENAI_AUTHORIZATION_HEADER) + api_keys = header_value.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[OPENAI_AUTHORIZATION_HEADER] = f"{api_keys[0].strip()}" client = httpx.AsyncClient(timeout=httpx.Timeout(CLIENT_TIMEOUT)) open_ai_request = client.build_request( @@ -98,7 +110,7 @@ async def stream_response( open_ai_request: httpx.Request, dataset_name: Optional[str], request_body_json: dict[str, Any], - invariant_authorization: str, + invariant_authorization: Optional[str], ) -> Response: """ Handles streaming the OpenAI response to the client while building a merged_response @@ -325,7 +337,7 @@ async def handle_non_streaming_response( response: httpx.Response, dataset_name: Optional[str], request_body_json: dict[str, Any], - invariant_authorization: str, + invariant_authorization: Optional[str], ) -> Response: """Handles non-streaming OpenAI responses""" try: diff --git a/proxy/utils/constants.py b/proxy/utils/constants.py index 858a971..12801c2 100644 --- a/proxy/utils/constants.py +++ b/proxy/utils/constants.py @@ -13,3 +13,4 @@ IGNORED_HEADERS = [ ] CLIENT_TIMEOUT = 60.0 +INVARIANT_AUTHORIZATION_HEADER = "invariant-authorization"