diff --git a/README.md b/README.md index 23e173f..d035c81 100644 --- a/README.md +++ b/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 "}, + 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: diff --git a/gateway/common/authorization.py b/gateway/common/authorization.py index 6a4e844..74ba70b 100644 --- a/gateway/common/authorization.py +++ b/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 " {llm_provider_api_key_header} contains the LLM Provider API Key as {llm_provider_api_key_header}: "" + + 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 diff --git a/gateway/routes/gemini.py b/gateway/routes/gemini.py index 2643125..07aa36a 100644 --- a/gateway/routes/gemini.py +++ b/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 diff --git a/gateway/routes/open_ai.py b/gateway/routes/open_ai.py index ff565e0..85633d9 100644 --- a/gateway/routes/open_ai.py +++ b/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) diff --git a/tests/integration/litellm/test_chat_without_tool_call.py b/tests/integration/litellm/test_chat_without_tool_call.py new file mode 100644 index 0000000..879cd73 --- /dev/null +++ b/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 "}, + 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, + }, + ] \ No newline at end of file diff --git a/tests/integration/requirements.txt b/tests/integration/requirements.txt index 456cfbb..5a04ac8 100644 --- a/tests/integration/requirements.txt +++ b/tests/integration/requirements.txt @@ -5,3 +5,4 @@ pillow pytest pytest-asyncio tavily-python +litellm diff --git a/tests/unit_tests/common/test_authorization.py b/tests/unit_tests/common/test_authorization.py new file mode 100644 index 0000000..3742f35 --- /dev/null +++ b/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 \ No newline at end of file