mirror of
https://github.com/BerriAI/litellm.git
synced 2025-04-26 11:14:04 +00:00
feat(user_api_key_auth.py): support calling langfuse with litellm user_api_key_auth
This commit is contained in:
Krrish Dholakia
parent
66d77f177f
commit
742e3cbccf
4 changed files with 160 additions and 7 deletions
|
|
@ -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,93 @@ 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_pass_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
|
||||
|
||||
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://us.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 {}
|
||||
)
|
||||
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
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pass_through_endpoint_anthropic(client):
|
||||
import litellm
|
||||
|
|
@ -178,6 +283,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",
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue