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