Merge pull request #35 from invariantlabs-ai/gemini/litellm-support

feat: add litellm explicit support

README.md

@@ -149,6 +149,37 @@ print(response.messages[-1]["content"])
 # Output: "It seems to be sunny."
 ```
 
+### **🔹 LiteLLM Integration**
+
+LiteLLM is a python library that acts as a unified interface for calling multiple LLM providers. If you are using it, it is very convinient to connect to Gateway proxy. You just need to pass the correct `base_url`.
+
+```python
+from litellm import completion
+import random
+import os
+
+base_url = "/api/v1/gateway/litellm/{add-your-dataset-name-here}"
+EXAMPLE_MODELS = ["openai/gpt-4o", "gemini/gemini-2.0-flash", "anthropic/claude-3-5-haiku-20241022"]
+model = random.choice(SAMPLE_MODELS)
+
+base_url += "/" + model.split("/")[0] # append /gemini /openai or /anthropic. 
+
+if model.split("/")[0] == "gemini":
+    base_url += f"/v1beta/models/{model.split('/')[1]}" # gemini expects the model name in the url.
+
+
+chat_response = completion(
+    model=model,
+    messages=[{"role": "user", "content": "What is the capital of France?"}],
+    extra_headers= {"Invariant-Authorization": "Bearer <some-key>"},
+    stream=True,
+    base_url=base_url,
+)
+
+print(chat_response.choices[0].message.content)
+# Output: "Paris."
+```
+
 ### **🔹 Microsoft Autogen Integration**
 
 You can also easily integrate the Gateway with [Microsoft Autogen](https://github.com/microsoft/autogen) as follows:

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,16 @@ 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:
+                llm_provider_api_key_header = header
+                break
+
+
     if dataset_name:
         if invariant_authorization is None:
             if llm_provider_api_key is None:
@@ -49,4 +65,8 @@ def extract_authorization_from_headers(
 
             invariant_authorization = f"Bearer {api_keys[1].strip()}"
             llm_provider_api_key = f"{api_keys[0].strip()}"
+
+    if llm_provider_api_key and "Bearer " in llm_provider_api_key:
+        llm_provider_api_key = llm_provider_api_key.split("Bearer ")[1].strip()
+
     return invariant_authorization, llm_provider_api_key

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

tests/unit_tests/common/test_authorization.py

@@ -0,0 +1,81 @@
+"""Tests for the authorization header extractor."""
+
+import os
+import sys
+from fastapi import HTTPException
+import random
+import string
+import pytest
+
+
+# Add root folder (parent) to sys.path
+sys.path.append(
+    os.path.dirname(
+        os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+    )
+)
+
+from gateway.common.authorization import extract_authorization_from_headers, INVARIANT_AUTHORIZATION_HEADER, API_KEYS_SEPARATOR
+
+@pytest.mark.parametrize("push_to_explorer", [True, False])
+@pytest.mark.parametrize("invariant_authorization", [True, False])
+@pytest.mark.parametrize("invariant_authorization_appended_to_llm_provider_api_key", [True, False])
+@pytest.mark.parametrize("use_fallback_header", [True, False])
+def test_extract_authorization_from_headers(
+    push_to_explorer: bool,
+    invariant_authorization: bool,
+    invariant_authorization_appended_to_llm_provider_api_key: bool,
+    use_fallback_header: bool,
+):
+    """Test the extract_authorization_from_headers function."""
+
+    llm_apikey = ''.join(random.choices(string.ascii_letters + string.digits, k=10))
+    inv_apikey = ''.join(random.choices(string.ascii_letters + string.digits, k=10))
+    dataset_name = "test-dataset" if push_to_explorer else None
+
+    headers: dict[str, str] = {} 
+
+    llm_provider_api_key = "fallback-header" if use_fallback_header else "llm-provider-api-key"
+    headers[llm_provider_api_key] = llm_apikey
+
+    if invariant_authorization:
+        print("invariant_authorization - TRUE")
+        if invariant_authorization_appended_to_llm_provider_api_key:
+            headers[llm_provider_api_key] = f"{headers.get(llm_provider_api_key, '')}{API_KEYS_SEPARATOR}{inv_apikey}"
+        else:
+            headers[INVARIANT_AUTHORIZATION_HEADER] = f"Bearer {inv_apikey}"
+
+    # Mock request headers
+    class MockRequest:
+        def __init__(self, headers):
+            self.headers = headers
+
+    request = MockRequest(headers) 
+
+    # Call the function
+    try:
+        invariant_auth, llm_provider_api_key = extract_authorization_from_headers(
+            request,
+            dataset_name=dataset_name,
+            llm_provider_api_key_header="llm-provider-api-key",
+            llm_provider_fallback_api_key_headers=["fallback-header"],
+        )
+    except HTTPException as e:
+        # If an exception is raised, check if it is the expected one
+        if not invariant_authorization:
+            assert e.status_code == 400
+            assert e.detail == "Missing invariant api key"
+            return
+        else:
+            raise e
+    # Verify the results
+    if invariant_authorization:
+        if not push_to_explorer and invariant_authorization_appended_to_llm_provider_api_key:
+            assert llm_provider_api_key.split(API_KEYS_SEPARATOR)[0] == llm_apikey
+            assert llm_provider_api_key.split(API_KEYS_SEPARATOR)[1] == inv_apikey
+        else:
+            assert invariant_auth == ("Bearer " + inv_apikey)
+    else:
+        assert invariant_auth is None
+    if not(not push_to_explorer and invariant_authorization_appended_to_llm_provider_api_key):
+        assert llm_provider_api_key == llm_apikey