Merge pull request #5192 from BerriAI/litellm_fix_pass_through_endpoints

feat(user_api_key_auth.py): support calling langfuse with litellm user_api_key_auth

litellm/proxy/_new_secret_config.yaml

@@ -3,4 +3,5 @@ model_list:
     litellm_params:
       model: "gpt-4"
     model_info:
-      my_custom_key: "my_custom_value"
+      my_custom_key: "my_custom_value"
+

litellm/proxy/auth/user_api_key_auth.py

@@ -85,6 +85,8 @@ def _get_bearer_token(
 ):
     if api_key.startswith("Bearer "):  # ensure Bearer token passed in
         api_key = api_key.replace("Bearer ", "")  # extract the token
+    elif api_key.startswith("Basic "):
+        api_key = api_key.replace("Basic ", "")  # handle langfuse input
     else:
         api_key = ""
     return api_key
@@ -138,7 +140,6 @@ async def user_api_key_auth(
         pass_through_endpoints: Optional[List[dict]] = general_settings.get(
             "pass_through_endpoints", None
         )
-
         if isinstance(api_key, str):
             passed_in_key = api_key
             api_key = _get_bearer_token(api_key=api_key)
@@ -367,6 +368,40 @@ async def user_api_key_auth(
                     parent_otel_span=parent_otel_span,
                 )
         #### ELSE ####
+
+        ## CHECK PASS-THROUGH ENDPOINTS ##
+        if pass_through_endpoints is not None:
+            for endpoint in pass_through_endpoints:
+                if endpoint.get("path", "") == route:
+                    ## IF AUTH DISABLED
+                    if endpoint.get("auth") is not True:
+                        return UserAPIKeyAuth()
+                    ## IF AUTH ENABLED
+                    ### IF CUSTOM PARSER REQUIRED
+                    if (
+                        endpoint.get("custom_auth_parser") is not None
+                        and endpoint.get("custom_auth_parser") == "langfuse"
+                    ):
+                        """
+                        - langfuse returns {'Authorization': 'Basic YW55dGhpbmc6YW55dGhpbmc'}
+                        - check the langfuse public key if it contains the litellm api key
+                        """
+                        import base64
+
+                        api_key = api_key.replace("Basic ", "").strip()
+                        decoded_bytes = base64.b64decode(api_key)
+                        decoded_str = decoded_bytes.decode("utf-8")
+                        api_key = decoded_str.split(":")[0]
+                    else:
+                        headers = endpoint.get("headers", None)
+                        if headers is not None:
+                            header_key = headers.get("litellm_user_api_key", "")
+                            if (
+                                isinstance(request.headers, dict)
+                                and request.headers.get(key=header_key) is not None
+                            ):
+                                api_key = request.headers.get(key=header_key)
+
         if master_key is None:
             if isinstance(api_key, str):
                 return UserAPIKeyAuth(
@@ -533,7 +568,11 @@ async def user_api_key_auth(
         if isinstance(
             api_key, str
         ):  # if generated token, make sure it starts with sk-.
-            assert api_key.startswith("sk-")  # prevent token hashes from being used
+            assert api_key.startswith(
+                "sk-"
+            ), "LiteLLM Virtual Key expected. Received={}, expected to start with 'sk-'.".format(
+                api_key
+            )  # prevent token hashes from being used
         else:
             verbose_logger.warning(
                 "litellm.proxy.proxy_server.user_api_key_auth(): Warning - Key={} is not a string.".format(

litellm/proxy/pass_through_endpoints/pass_through_endpoints.py

@@ -309,7 +309,7 @@ async def pass_through_request(
             json=_parsed_body,
         )
 
-        if response.status_code != 200:
+        if response.status_code >= 300:
             raise HTTPException(status_code=response.status_code, detail=response.text)
 
         content = await response.aread()

litellm/tests/test_pass_through_endpoints.py

@@ -1,5 +1,6 @@
 import os
 import sys
+from typing import Optional
 
 import pytest
 from fastapi import FastAPI
@@ -30,6 +31,7 @@ def client():
 async def test_pass_through_endpoint(client, monkeypatch):
     # Mock the httpx.AsyncClient.request method
     monkeypatch.setattr("httpx.AsyncClient.request", mock_request)
+    import litellm
 
     # Define a pass-through endpoint
     pass_through_endpoints = [
@@ -42,6 +44,11 @@ async def test_pass_through_endpoint(client, monkeypatch):
 
     # Initialize the pass-through endpoint
     await initialize_pass_through_endpoints(pass_through_endpoints)
+    general_settings: Optional[dict] = (
+        getattr(litellm.proxy.proxy_server, "general_settings", {}) or {}
+    )
+    general_settings.update({"pass_through_endpoints": pass_through_endpoints})
+    setattr(litellm.proxy.proxy_server, "general_settings", general_settings)
 
     # Make a request to the pass-through endpoint
     response = client.post("/test-endpoint", json={"prompt": "Hello, world!"})
@@ -54,6 +61,7 @@ async def test_pass_through_endpoint(client, monkeypatch):
 @pytest.mark.asyncio
 async def test_pass_through_endpoint_rerank(client):
     _cohere_api_key = os.environ.get("COHERE_API_KEY")
+    import litellm
 
     # Define a pass-through endpoint
     pass_through_endpoints = [
@@ -66,6 +74,11 @@ async def test_pass_through_endpoint_rerank(client):
 
     # Initialize the pass-through endpoint
     await initialize_pass_through_endpoints(pass_through_endpoints)
+    general_settings: Optional[dict] = (
+        getattr(litellm.proxy.proxy_server, "general_settings", {}) or {}
+    )
+    general_settings.update({"pass_through_endpoints": pass_through_endpoints})
+    setattr(litellm.proxy.proxy_server, "general_settings", general_settings)
 
     _json_data = {
         "model": "rerank-english-v3.0",
@@ -87,7 +100,7 @@ async def test_pass_through_endpoint_rerank(client):
 
 @pytest.mark.parametrize(
     "auth, rpm_limit, expected_error_code",
-    [(True, 0, 429), (True, 1, 200), (False, 0, 401)],
+    [(True, 0, 429), (True, 1, 200), (False, 0, 200)],
 )
 @pytest.mark.asyncio
 async def test_pass_through_endpoint_rpm_limit(auth, expected_error_code, rpm_limit):
@@ -123,6 +136,11 @@ async def test_pass_through_endpoint_rpm_limit(auth, expected_error_code, rpm_li
 
     # Initialize the pass-through endpoint
     await initialize_pass_through_endpoints(pass_through_endpoints)
+    general_settings: Optional[dict] = (
+        getattr(litellm.proxy.proxy_server, "general_settings", {}) or {}
+    )
+    general_settings.update({"pass_through_endpoints": pass_through_endpoints})
+    setattr(litellm.proxy.proxy_server, "general_settings", general_settings)
 
     _json_data = {
         "model": "rerank-english-v3.0",
@@ -146,6 +164,123 @@ async def test_pass_through_endpoint_rpm_limit(auth, expected_error_code, rpm_li
     assert response.status_code == expected_error_code
 
 
+@pytest.mark.parametrize(
+    "auth, rpm_limit, expected_error_code",
+    [(True, 0, 429), (True, 1, 207), (False, 0, 207)],
+)
+@pytest.mark.asyncio
+async def test_aaapass_through_endpoint_pass_through_keys_langfuse(
+    auth, expected_error_code, rpm_limit
+):
+
+    client = TestClient(app)
+    import litellm
+    from litellm.proxy._types import UserAPIKeyAuth
+    from litellm.proxy.proxy_server import ProxyLogging, hash_token, user_api_key_cache
+
+    # Store original values
+    original_user_api_key_cache = getattr(
+        litellm.proxy.proxy_server, "user_api_key_cache", None
+    )
+    original_master_key = getattr(litellm.proxy.proxy_server, "master_key", None)
+    original_prisma_client = getattr(litellm.proxy.proxy_server, "prisma_client", None)
+    original_proxy_logging_obj = getattr(
+        litellm.proxy.proxy_server, "proxy_logging_obj", None
+    )
+
+    try:
+
+        mock_api_key = "sk-my-test-key"
+        cache_value = UserAPIKeyAuth(
+            token=hash_token(mock_api_key), rpm_limit=rpm_limit
+        )
+
+        _cohere_api_key = os.environ.get("COHERE_API_KEY")
+
+        user_api_key_cache.set_cache(key=hash_token(mock_api_key), value=cache_value)
+
+        proxy_logging_obj = ProxyLogging(user_api_key_cache=user_api_key_cache)
+        proxy_logging_obj._init_litellm_callbacks()
+
+        setattr(litellm.proxy.proxy_server, "user_api_key_cache", user_api_key_cache)
+        setattr(litellm.proxy.proxy_server, "master_key", "sk-1234")
+        setattr(litellm.proxy.proxy_server, "prisma_client", "FAKE-VAR")
+        setattr(litellm.proxy.proxy_server, "proxy_logging_obj", proxy_logging_obj)
+
+        # Define a pass-through endpoint
+        pass_through_endpoints = [
+            {
+                "path": "/api/public/ingestion",
+                "target": "https://cloud.langfuse.com/api/public/ingestion",
+                "auth": auth,
+                "custom_auth_parser": "langfuse",
+                "headers": {
+                    "LANGFUSE_PUBLIC_KEY": "os.environ/LANGFUSE_PUBLIC_KEY",
+                    "LANGFUSE_SECRET_KEY": "os.environ/LANGFUSE_SECRET_KEY",
+                },
+            }
+        ]
+
+        # Initialize the pass-through endpoint
+        await initialize_pass_through_endpoints(pass_through_endpoints)
+        general_settings: Optional[dict] = (
+            getattr(litellm.proxy.proxy_server, "general_settings", {}) or {}
+        )
+        old_general_settings = general_settings
+        general_settings.update({"pass_through_endpoints": pass_through_endpoints})
+        setattr(litellm.proxy.proxy_server, "general_settings", general_settings)
+
+        _json_data = {
+            "batch": [
+                {
+                    "id": "80e2141f-0ca6-47b7-9c06-dde5e97de690",
+                    "type": "trace-create",
+                    "body": {
+                        "id": "0687af7b-4a75-4de8-a4f6-cba1cdc00865",
+                        "timestamp": "2024-08-14T02:38:56.092950Z",
+                        "name": "test-trace-litellm-proxy-passthrough",
+                    },
+                    "timestamp": "2024-08-14T02:38:56.093352Z",
+                }
+            ],
+            "metadata": {
+                "batch_size": 1,
+                "sdk_integration": "default",
+                "sdk_name": "python",
+                "sdk_version": "2.27.0",
+                "public_key": "anything",
+            },
+        }
+
+        # Make a request to the pass-through endpoint
+        response = client.post(
+            "/api/public/ingestion",
+            json=_json_data,
+            headers={"Authorization": "Basic c2stbXktdGVzdC1rZXk6YW55dGhpbmc="},
+        )
+
+        print("JSON response: ", _json_data)
+
+        print("RESPONSE RECEIVED - {}".format(response.text))
+
+        # Assert the response
+        assert response.status_code == expected_error_code
+
+        setattr(litellm.proxy.proxy_server, "general_settings", old_general_settings)
+    finally:
+        # Reset to original values
+        setattr(
+            litellm.proxy.proxy_server,
+            "user_api_key_cache",
+            original_user_api_key_cache,
+        )
+        setattr(litellm.proxy.proxy_server, "master_key", original_master_key)
+        setattr(litellm.proxy.proxy_server, "prisma_client", original_prisma_client)
+        setattr(
+            litellm.proxy.proxy_server, "proxy_logging_obj", original_proxy_logging_obj
+        )
+
+
 @pytest.mark.asyncio
 async def test_pass_through_endpoint_anthropic(client):
     import litellm
@@ -178,6 +313,11 @@ async def test_pass_through_endpoint_anthropic(client):
 
     # Initialize the pass-through endpoint
     await initialize_pass_through_endpoints(pass_through_endpoints)
+    general_settings: Optional[dict] = (
+        getattr(litellm.proxy.proxy_server, "general_settings", {}) or {}
+    )
+    general_settings.update({"pass_through_endpoints": pass_through_endpoints})
+    setattr(litellm.proxy.proxy_server, "general_settings", general_settings)
 
     _json_data = {
         "model": "gpt-3.5-turbo",