feat: add litellm support

gateway/common/authorization.py

@@ -8,7 +8,10 @@ API_KEYS_SEPARATOR = ";invariant-auth="
 
 
 def extract_authorization_from_headers(
-    request: Request, dataset_name: Optional[str], llm_provider_api_key_header: str
+    request: Request,
+    dataset_name: Optional[str],
+    llm_provider_api_key_header: str,
+    llm_provider_fallback_api_key_headers: list[str] | None = None
 ) -> Tuple[str, str]:
     """
     Extracts the Invariant authorization and LLM Provider API key from the request headers.
@@ -19,8 +22,11 @@ def extract_authorization_from_headers(
     "invariant-authorization": "Bearer <Invariant API Key>"
     {llm_provider_api_key_header} contains the LLM Provider API Key as
     {llm_provider_api_key_header}: "<API Key>"
+    
+    If {llm_provider_api_key_header} is not among headers, we look for
+    any header among {llm_provider_fallback_api_key_headers}.
 
-    For some clients, it is not possible to pass a custom header
+    For some clients, it is not possible to pass a custom header at all,
     In such cases, the Invariant API Key is passed as part of the
     {llm_provider_api_key_header} with the LLM Provider API Key
     The header in that case becomes:
@@ -28,6 +34,17 @@ def extract_authorization_from_headers(
     """
     invariant_authorization = request.headers.get(INVARIANT_AUTHORIZATION_HEADER)
     llm_provider_api_key = request.headers.get(llm_provider_api_key_header)
+
+
+    if llm_provider_api_key is None and llm_provider_fallback_api_key_headers:
+        for header in llm_provider_fallback_api_key_headers:
+            llm_provider_api_key = request.headers.get(header)
+            if llm_provider_api_key:
+                break
+
+    if "Bearer " in llm_provider_api_key:
+        llm_provider_api_key = llm_provider_api_key.split("Bearer ")[1].strip()
+
     if dataset_name:
         if invariant_authorization is None:
             if llm_provider_api_key is None:

gateway/routes/gemini.py

@@ -34,6 +34,7 @@ from integrations.guardrails import (
 gateway = APIRouter()
 
 GEMINI_AUTHORIZATION_HEADER = "x-goog-api-key"
+GEMINI_AUTHORIZATION_FALLBACK_HEADER = "authorization"
 
 
 @gateway.post("/gemini/{api_version}/models/{model}:{endpoint}")
@@ -58,12 +59,11 @@ async def gemini_generate_content_gateway(
             status_code=400,
         )
     headers = {
-        k: v for k, v in request.headers.items() if k.lower() not in IGNORED_HEADERS
+        k: v for k, v in request.headers.items() if k.lower() not in IGNORED_HEADERS + [GEMINI_AUTHORIZATION_FALLBACK_HEADER]
     }
     headers["accept-encoding"] = "identity"
-
     invariant_authorization, gemini_api_key = extract_authorization_from_headers(
-        request, dataset_name, GEMINI_AUTHORIZATION_HEADER
+        request, dataset_name, GEMINI_AUTHORIZATION_HEADER, [GEMINI_AUTHORIZATION_FALLBACK_HEADER]
     )
     headers[GEMINI_AUTHORIZATION_HEADER] = gemini_api_key
 

gateway/routes/open_ai.py

@@ -64,7 +64,7 @@ async def openai_chat_completions_gateway(
     invariant_authorization, openai_api_key = extract_authorization_from_headers(
         request, dataset_name, OPENAI_AUTHORIZATION_HEADER
     )
-    headers[OPENAI_AUTHORIZATION_HEADER] = openai_api_key
+    headers[OPENAI_AUTHORIZATION_HEADER] = "Bearer " + openai_api_key
 
     request_body_bytes = await request.body()
     request_json = json.loads(request_body_bytes)

tests/integration/litellm/test_chat_without_tool_call.py

@@ -0,0 +1,95 @@
+import pytest
+import uuid
+from litellm import completion
+import litellm
+import time
+import requests
+import os
+
+MODEL_API_KEYS = {
+    "openai/gpt-4o": "OPENAI_API_KEY",
+    "gemini/gemini-2.0-flash": "GEMINI_API_KEY",
+    "anthropic/claude-3-5-haiku-20241022": "ANTHROPIC_API_KEY",
+}
+
+@pytest.mark.parametrize("litellm_model", MODEL_API_KEYS.keys(),)
+@pytest.mark.parametrize(
+    "do_stream, push_to_explorer",
+    [(False, False)],
+)
+async def test_chat_completion(
+    explorer_api_url: str, litellm_model: str, gateway_url: str, do_stream: bool, push_to_explorer: bool
+):
+    """Test the chat completions gateway calls with tool calling through litellm."""
+    # Check if the API key is set in the environment variables
+    api_key_env_var = MODEL_API_KEYS[litellm_model]
+    api_key = os.getenv(api_key_env_var)
+
+    if not api_key:
+        pytest.skip(f"Skipping {litellm_model} because {api_key_env_var} is not set")
+
+
+    dataset_name = f"test-dataset-litellm-{litellm_model}-{uuid.uuid4()}"
+    base_url = (f"{gateway_url}/api/v1/gateway/{dataset_name}"
+    if push_to_explorer
+    else f"{gateway_url}/api/v1/gateway")
+
+    base_url += "/" + litellm_model.split("/")[0] #add provider name
+    if litellm_model.split("/")[0] == "gemini":
+        base_url += f"/v1beta/models/{litellm_model.split('/')[1]}" #gemini expects the model name in the url
+
+    print(f"base_url: {base_url}")
+    chat_response = completion(
+        model=litellm_model,
+        messages=[{"role": "user", "content": "What is the capital of France?"}],
+        extra_headers= {"Invariant-Authorization": "Bearer <some-key>"},
+        stream=do_stream,
+        base_url=base_url,
+    )
+
+    # Verify the chat response
+    if not do_stream:
+        assert "PARIS" in chat_response.choices[0].message.content.upper()
+        expected_assistant_message = chat_response.choices[0].message.content
+    else:
+        full_response = ""
+        for chunk in chat_response:
+            if chunk.choices and chunk.choices[0].delta.content:
+                full_response += chunk.choices[0].delta.content
+        assert "PARIS" in full_response.upper()
+        expected_assistant_message = full_response
+
+    if push_to_explorer:
+        # Wait for the trace to be saved
+        # This is needed because the trace is saved asynchronously
+        time.sleep(2)
+        # Fetch the trace ids for the dataset
+        traces_response = requests.get(
+            f"{explorer_api_url}/api/v1/dataset/byuser/developer/{dataset_name}/traces",
+            timeout=5,
+        )
+        traces = traces_response.json()
+        assert len(traces) == 1
+        trace_id = traces[0]["id"]
+
+        # Fetch the trace
+        trace_response = requests.get(
+            f"{explorer_api_url}/api/v1/trace/{trace_id}",
+            timeout=5,
+        )
+        trace = trace_response.json()
+
+        for message in trace["messages"]:
+            message.pop("annotations", None)
+
+        # Verify the trace messages
+        assert trace["messages"] == [
+            {
+                "role": "user",
+                "content": "What is the capital of France?",
+            },
+            {
+                "role": "assistant",
+                "content": expected_assistant_message,
+            },
+        ]

tests/integration/requirements.txt

@@ -5,3 +5,4 @@ pillow
 pytest
 pytest-asyncio
 tavily-python
+litellm