ai-workspace-services/litellm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Move UI to encrypted token usage (#10302)

Browse Source 
* test(test_auth_checks.py): add unit tests for ExperimentalUIJWTToken

* test: add appropriate flag

* fix: fix ruff check

* test: add autouse fixture to test salt key

* fix(user_api_key_auth.py): fix auth flow logic

* test: skip flaky test - anthropic does not reliably return 'redacted_thinking'
This commit is contained in:
Krish Dholakia 2025-04-25 18:20:41 -07:00 committed by GitHub
parent fee3dc970c
commit 70f7c73def
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
10 changed files with 230 additions and 43 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
litellm/proxy/_new_secret_config.yaml
View File

@ -4,6 +4,11 @@ model_list:
model: azure/gpt-4o
api_key: os.environ/AZURE_API_KEY
api_base: os.environ/AZURE_API_BASE
- model_name: fake-openai-endpoint
litellm_params:
model: openai/fake
api_key: fake-key
api_base: https://exampleopenaiendpoint-production.up.railway.app/
- model_name: "gpt-4o-mini-openai"
litellm_params:
model: gpt-4o-mini
@ -33,7 +38,6 @@ model_list:
litellm_settings:
num_retries: 0
callbacks: ["datadog_llm_observability"]
check_provider_endpoint: true
files_settings:

56
litellm/proxy/auth/auth_checks.py
View File

@ -953,6 +953,62 @@ async def get_team_object(
)
class ExperimentalUIJWTToken:
@staticmethod
def get_experimental_ui_login_jwt_auth_token(user_info: LiteLLM_UserTable) -> str:
from datetime import UTC, datetime, timedelta
from litellm.proxy.common_utils.encrypt_decrypt_utils import (
encrypt_value_helper,
)
if user_info.user_role is None:
raise Exception("User role is required for experimental UI login")
# Calculate expiration time (10 minutes from now)
expiration_time = datetime.now(UTC) + timedelta(minutes=10)
# Format the expiration time as ISO 8601 string
expires = expiration_time.strftime("%Y-%m-%dT%H:%M:%S.%f")[:-3] + "+00:00"
valid_token = UserAPIKeyAuth(
token="ui-token",
key_name="ui-token",
key_alias="ui-token",
max_budget=litellm.max_ui_session_budget,
rpm_limit=100, # allow user to have a conversation on test key pane of UI
expires=expires,
user_id=user_info.user_id,
team_id="litellm-dashboard",
models=user_info.models,
max_parallel_requests=None,
user_role=LitellmUserRoles(user_info.user_role),
)
return encrypt_value_helper(valid_token.model_dump_json(exclude_none=True))
@staticmethod
def get_key_object_from_ui_hash_key(
hashed_token: str,
) -> Optional[UserAPIKeyAuth]:
import json
from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth
from litellm.proxy.common_utils.encrypt_decrypt_utils import (
decrypt_value_helper,
)
decrypted_token = decrypt_value_helper(hashed_token, exception_type="debug")
if decrypted_token is None:
return None
try:
return UserAPIKeyAuth(**json.loads(decrypted_token))
except Exception as e:
raise Exception(
f"Invalid hash key. Hash key={hashed_token}. Decrypted token={decrypted_token}. Error: {e}"
)
@log_db_metrics
async def get_key_object(
hashed_token: str,

8
litellm/proxy/auth/user_api_key_auth.py
View File

@ -24,6 +24,7 @@ from litellm.caching import DualCache
from litellm.litellm_core_utils.dd_tracing import tracer
from litellm.proxy._types import *
from litellm.proxy.auth.auth_checks import (
ExperimentalUIJWTToken,
_cache_key_object,
_get_user_role,
_is_user_proxy_admin,
@ -51,6 +52,7 @@ from litellm.proxy.auth.oauth2_check import check_oauth2_token
from litellm.proxy.auth.oauth2_proxy_hook import handle_oauth2_proxy_request
from litellm.proxy.common_utils.http_parsing_utils import _read_request_body
from litellm.proxy.utils import PrismaClient, ProxyLogging
from litellm.secret_managers.main import get_secret_bool
from litellm.types.services import ServiceTypes
user_api_key_service_logger_obj = ServiceLogging() # used for tracking latency on OTEL
@ -553,6 +555,12 @@ async def _user_api_key_auth_builder( # noqa: PLR0915
verbose_logger.debug("api key not found in cache.")
valid_token = None
## Check UI Hash Key
if valid_token is None and get_secret_bool("EXPERIMENTAL_UI_LOGIN"):
valid_token = ExperimentalUIJWTToken.get_key_object_from_ui_hash_key(
api_key
)
if (
valid_token is not None
and isinstance(valid_token, UserAPIKeyAuth)

16
litellm/proxy/common_utils/encrypt_decrypt_utils.py
View File

@ -1,6 +1,6 @@
import base64
import os
from typing import Optional
from typing import Literal, Optional
from litellm._logging import verbose_proxy_logger
@ -39,7 +39,9 @@ def encrypt_value_helper(value: str, new_encryption_key: Optional[str] = None):
raise e
def decrypt_value_helper(value: str):
def decrypt_value_helper(
value: str, exception_type: Literal["debug", "error"] = "error"
):
signing_key = _get_salt_key()
try:
@ -51,11 +53,13 @@ def decrypt_value_helper(value: str):
# if it's not str - do not decrypt it, return the value
return value
except Exception as e:
verbose_proxy_logger.error(
f"Error decrypting value, Did your master_key/salt key change recently? \nError: {str(e)}\nSet permanent salt key - https://docs.litellm.ai/docs/proxy/prod#5-set-litellm-salt-key"
)
error_message = f"Error decrypting value, Did your master_key/salt key change recently? \nError: {str(e)}\nSet permanent salt key - https://docs.litellm.ai/docs/proxy/prod#5-set-litellm-salt-key"
if exception_type == "debug":
verbose_proxy_logger.debug(error_message)
return None
verbose_proxy_logger.error(error_message)
# [Non-Blocking Exception. - this should not block decrypting other values]
pass
return None
def encrypt_value(value: str, signing_key: str):

1
litellm/proxy/management_endpoints/team_endpoints.py
View File

@ -538,6 +538,7 @@ async def update_team(
detail={"error": "Team doesn't exist. Got={}".format(team_row)},
)
verbose_proxy_logger.info("Successfully updated team - %s, info", team_row.team_id)
await _cache_team_object(
team_id=team_row.team_id,
team_table=LiteLLM_TeamTableCachedObj(**team_row.model_dump()),

4
litellm/proxy/management_endpoints/ui_sso.py
View File

@ -642,7 +642,9 @@ async def auth_callback(request: Request): # noqa: PLR0915
litellm_dashboard_ui += "?login=success"
verbose_proxy_logger.info(f"Redirecting to {litellm_dashboard_ui}")
redirect_response = RedirectResponse(url=litellm_dashboard_ui, status_code=303)
redirect_response.set_cookie(key="token", value=jwt_token, secure=True)
redirect_response.set_cookie(
key="token", value=jwt_token, secure=True, httponly=True
)
return redirect_response

40
litellm/proxy/proxy_server.py
View File

@ -150,7 +150,11 @@ from litellm.proxy.analytics_endpoints.analytics_endpoints import (
router as analytics_router,
)
from litellm.proxy.anthropic_endpoints.endpoints import router as anthropic_router
from litellm.proxy.auth.auth_checks import get_team_object, log_db_metrics
from litellm.proxy.auth.auth_checks import (
ExperimentalUIJWTToken,
get_team_object,
log_db_metrics,
)
from litellm.proxy.auth.auth_utils import check_response_size_is_safe
from litellm.proxy.auth.handle_jwt import JWTHandler
from litellm.proxy.auth.litellm_license import LicenseCheck
@ -6724,7 +6728,7 @@ async def login(request: Request): # noqa: PLR0915
)
# check if we can find the `username` in the db. on the ui, users can enter username=their email
_user_row = None
_user_row: Optional[LiteLLM_UserTable] = None
user_role: Optional[
Literal[
LitellmUserRoles.PROXY_ADMIN,
@ -6734,8 +6738,11 @@ async def login(request: Request): # noqa: PLR0915
]
] = None
if prisma_client is not None:
_user_row = await prisma_client.db.litellm_usertable.find_first(
where={"user_email": {"equals": username}}
_user_row = cast(
Optional[LiteLLM_UserTable],
await prisma_client.db.litellm_usertable.find_first(
where={"user_email": {"equals": username}}
),
)
disabled_non_admin_personal_key_creation = (
get_disabled_non_admin_personal_key_creation()
@ -6800,6 +6807,31 @@ async def login(request: Request): # noqa: PLR0915
litellm_dashboard_ui += "/ui/"
import jwt
if get_secret_bool("EXPERIMENTAL_UI_LOGIN"):
user_info: Optional[LiteLLM_UserTable] = None
if _user_row is not None:
user_info = _user_row
elif (
user_id is not None
): # if user_id is not None, we are using the UI_USERNAME and UI_PASSWORD
user_info = LiteLLM_UserTable(
user_id=user_id,
user_role=user_role,
models=[],
max_budget=litellm.max_ui_session_budget,
)
if user_info is None:
raise HTTPException(
status_code=401,
detail={
"error": "User Information is required for experimental UI login"
},
)
key = ExperimentalUIJWTToken.get_experimental_ui_login_jwt_auth_token(
user_info
)
jwt_token = jwt.encode( # type: ignore
{
"user_id": user_id,

4
test-results/.last-run.json
View File

@ -1,4 +0,0 @@
{
"status": "failed",
"failedTests": []
}

111
tests/litellm/proxy/auth/test_auth_checks.py Normal file
View File

@ -0,0 +1,111 @@
import asyncio
import json
import os
import sys
from unittest.mock import AsyncMock, MagicMock, patch
sys.path.insert(
0, os.path.abspath("../../..")
) # Adds the parent directory to the system path
from datetime import UTC, datetime, timedelta
import pytest
import litellm
from litellm.proxy._types import (
LiteLLM_UserTable,
LitellmUserRoles,
SSOUserDefinedValues,
)
from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken
from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper
@pytest.fixture(autouse=True)
def set_salt_key(monkeypatch):
"""Automatically set LITELLM_SALT_KEY for all tests"""
monkeypatch.setenv("LITELLM_SALT_KEY", "sk-1234")
@pytest.fixture
def valid_sso_user_defined_values():
return LiteLLM_UserTable(
user_id="test_user",
user_email="test@example.com",
user_role=LitellmUserRoles.PROXY_ADMIN.value,
models=["gpt-3.5-turbo"],
max_budget=100.0,
)
@pytest.fixture
def invalid_sso_user_defined_values():
return LiteLLM_UserTable(
user_id="test_user",
user_email="test@example.com",
user_role=None, # Missing user role
models=["gpt-3.5-turbo"],
max_budget=100.0,
)
def test_get_experimental_ui_login_jwt_auth_token_valid(valid_sso_user_defined_values):
"""Test generating JWT token with valid user role"""
token = ExperimentalUIJWTToken.get_experimental_ui_login_jwt_auth_token(
valid_sso_user_defined_values
)
# Decrypt and verify token contents
decrypted_token = decrypt_value_helper(token, exception_type="debug")
token_data = json.loads(decrypted_token)
assert token_data["user_id"] == "test_user"
assert token_data["user_role"] == LitellmUserRoles.PROXY_ADMIN.value
assert token_data["models"] == ["gpt-3.5-turbo"]
assert token_data["max_budget"] == litellm.max_ui_session_budget
# Verify expiration time is set and valid
assert "expires" in token_data
expires = datetime.fromisoformat(token_data["expires"].replace("Z", "+00:00"))
assert expires > datetime.now(UTC)
assert expires <= datetime.now(UTC) + timedelta(minutes=10)
def test_get_experimental_ui_login_jwt_auth_token_invalid(
invalid_sso_user_defined_values,
):
"""Test generating JWT token with missing user role"""
with pytest.raises(Exception) as exc_info:
ExperimentalUIJWTToken.get_experimental_ui_login_jwt_auth_token(
invalid_sso_user_defined_values
)
assert str(exc_info.value) == "User role is required for experimental UI login"
def test_get_key_object_from_ui_hash_key_valid(
valid_sso_user_defined_values, monkeypatch
):
"""Test getting key object from valid UI hash key"""
monkeypatch.setenv("EXPERIMENTAL_UI_LOGIN", "True")
# Generate a valid token
token = ExperimentalUIJWTToken.get_experimental_ui_login_jwt_auth_token(
valid_sso_user_defined_values
)
# Get key object
key_object = ExperimentalUIJWTToken.get_key_object_from_ui_hash_key(token)
assert key_object is not None
assert key_object.user_id == "test_user"
assert key_object.user_role == LitellmUserRoles.PROXY_ADMIN
assert key_object.models == ["gpt-3.5-turbo"]
assert key_object.max_budget == litellm.max_ui_session_budget
def test_get_key_object_from_ui_hash_key_invalid():
"""Test getting key object from invalid UI hash key"""
# Test with invalid token
key_object = ExperimentalUIJWTToken.get_key_object_from_ui_hash_key("invalid_token")
assert key_object is None

27
tests/llm_translation/test_anthropic_completion.py
View File

@ -997,33 +997,6 @@ def test_anthropic_thinking_output(model):
assert resp.choices[0].message.thinking_blocks[0]["signature"] is not None
@pytest.mark.parametrize(
"model",
[
"anthropic/claude-3-7-sonnet-20250219",
# "bedrock/us.anthropic.claude-3-7-sonnet-20250219-v1:0",
],
)
def test_anthropic_redacted_thinking_output(model):
from litellm import completion
litellm._turn_on_debug()
resp = completion(
model=model,
messages=[{"role": "user", "content": "ANTHROPIC_MAGIC_STRING_TRIGGER_REDACTED_THINKING_46C9A13E193C177646C7398A98432ECCCE4C1253D5E2D82641AC0E52CC2876CB"}],
thinking={"type": "enabled", "budget_tokens": 1024},
)
print(resp)
assert resp.choices[0].message.thinking_blocks is not None
assert isinstance(resp.choices[0].message.thinking_blocks, list)
assert len(resp.choices[0].message.thinking_blocks) > 0
assert resp.choices[0].message.thinking_blocks[0]["type"] == "redacted_thinking"
assert resp.choices[0].message.thinking_blocks[0]["data"] is not None
@pytest.mark.parametrize(
"model",