From 7b6d77e0a5bd764cc312bbb27ec4bd9500076461 Mon Sep 17 00:00:00 2001 From: Marco Milanta Date: Wed, 2 Apr 2025 10:00:35 +0200 Subject: [PATCH] feat: add litellm support --- gateway/common/authorization.py | 21 +++- gateway/routes/gemini.py | 6 +- gateway/routes/open_ai.py | 2 +- .../litellm/test_chat_without_tool_call.py | 95 +++++++++++++++++++ tests/integration/requirements.txt | 1 + 5 files changed, 119 insertions(+), 6 deletions(-) create mode 100644 tests/integration/litellm/test_chat_without_tool_call.py diff --git a/gateway/common/authorization.py b/gateway/common/authorization.py index 6a4e844..aa7fffb 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,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: 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