scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-07 08:26:10 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 36

chore(proxy): align resource model auth checks

Browse source
This commit is contained in:
user 2026-04-30 21:59:56 -07:00
parent efa33bfe50
commit 336fe8276f
5 changed files with 501 additions and 50 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

9
litellm/proxy/auth/auth_checks.py
View file

@ -61,6 +61,10 @@ from litellm.proxy._types import (
UserAPIKeyAuth,
)
from litellm.proxy.auth.route_checks import RouteChecks
from litellm.proxy.common_utils.http_parsing_utils import (
_safe_get_request_headers,
_safe_get_request_query_params,
)
from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler
from litellm.proxy.guardrails.tool_name_extraction import (
TOOL_CAPABLE_CALL_TYPES,
@ -485,7 +489,10 @@ async def common_checks( # noqa: PLR0915
from litellm.proxy.proxy_server import prisma_client, user_api_key_cache
_model: Optional[Union[str, List[str]]] = get_model_from_request(
request_body, route
request_data=request_body,
route=route,
request_headers=_safe_get_request_headers(request=request),
request_query_params=_safe_get_request_query_params(request=request),
)
# 1. If team is blocked

253
litellm/proxy/auth/auth_utils.py
View file

@ -2,7 +2,7 @@ import os
import re
import sys
from functools import lru_cache
from typing import Any, List, Optional, Tuple
from typing import Any, Dict, List, Mapping, Optional, Tuple, Union
from fastapi import HTTPException, Request, status
@ -942,20 +942,249 @@ def get_end_user_id_from_request_body(
return None
def get_model_from_request(
request_data: dict, route: str
MODEL_ROUTING_HEADER_NAME = "x-litellm-model"
_MODEL_ROUTING_ROUTE_MARKERS = (
"/files",
"/batches",
"/vector_stores",
"/skills",
"/evals",
"/fine_tuning",
"/videos",
)
_MODEL_ROUTING_HEADER_OR_QUERY_ROUTE_MARKERS = (
"/files",
"/batches",
"/skills",
"/evals",
)
_MODEL_ROUTING_QUERY_TARGET_MODEL_ROUTE_MARKERS = (
"/files",
"/batches",
"/fine_tuning",
)
_MODEL_ROUTING_BODY_TARGET_MODEL_ROUTE_MARKERS = (
"/files",
"/batches",
"/vector_stores",
)
_MODEL_ROUTING_COMPLETION_MODEL_ROUTE_MARKERS = ("/evals",)
_MODEL_ROUTING_ID_FIELDS = (
"file_id",
"input_file_id",
"output_file_id",
"error_file_id",
"batch_id",
"fine_tuning_job_id",
"training_file",
"validation_file",
"vector_store_id",
"video_id",
"character_id",
)
def _append_model_candidates(candidates: List[str], value: Any) -> None:
if value is None:
return
if isinstance(value, str):
model_names = [model.strip() for model in value.split(",")]
elif isinstance(value, (list, tuple, set)):
for item in value:
_append_model_candidates(candidates=candidates, value=item)
return
else:
model_names = [str(value).strip()]
candidates.extend(model for model in model_names if model)
def _dedupe_model_candidates(candidates: List[str]) -> List[str]:
deduped: List[str] = []
for model in candidates:
if model not in deduped:
deduped.append(model)
return deduped
def _get_case_insensitive_mapping_value(
mapping: Optional[Mapping[str, Any]], key: str
) -> Any:
if not mapping:
return None
if key in mapping:
return mapping[key]
key_lower = key.lower()
for mapping_key, value in mapping.items():
if str(mapping_key).lower() == key_lower:
return value
return None
def _route_matches_any_marker(route: str, markers: Tuple[str, ...]) -> bool:
normalized_route = route.lower()
return any(marker in normalized_route for marker in markers)
def _route_uses_model_routing_sources(route: str) -> bool:
return _route_matches_any_marker(route=route, markers=_MODEL_ROUTING_ROUTE_MARKERS)
def _extract_models_from_managed_resource_id(resource_id: Any) -> List[str]:
if not isinstance(resource_id, str) or not resource_id:
return []
candidates: List[str] = []
try:
from litellm.proxy.openai_files_endpoints.common_utils import (
_is_base64_encoded_unified_file_id,
decode_model_from_file_id,
get_model_id_from_unified_batch_id,
get_models_from_unified_file_id,
)
_append_model_candidates(
candidates=candidates, value=decode_model_from_file_id(resource_id)
)
unified_file_id = _is_base64_encoded_unified_file_id(resource_id)
if unified_file_id:
_append_model_candidates(
candidates=candidates,
value=get_models_from_unified_file_id(unified_file_id),
)
_append_model_candidates(
candidates=candidates,
value=get_model_id_from_unified_batch_id(unified_file_id),
)
except Exception as e:
verbose_proxy_logger.debug(
"Unable to extract model from managed file/batch ID: %s", str(e)
)
try:
from litellm.llms.base_llm.managed_resources.utils import parse_unified_id
parsed_id = parse_unified_id(resource_id)
if parsed_id:
_append_model_candidates(
candidates=candidates, value=parsed_id.get("model_id")
)
_append_model_candidates(
candidates=candidates, value=parsed_id.get("target_model_names")
)
except Exception as e:
verbose_proxy_logger.debug(
"Unable to extract model from unified managed resource ID: %s", str(e)
)
try:
from litellm.types.videos.utils import (
decode_character_id_with_provider,
decode_video_id_with_provider,
)
_append_model_candidates(
candidates=candidates,
value=decode_video_id_with_provider(resource_id).get("model_id"),
)
_append_model_candidates(
candidates=candidates,
value=decode_character_id_with_provider(resource_id).get("model_id"),
)
except Exception as e:
verbose_proxy_logger.debug(
"Unable to extract model from managed video/character ID: %s", str(e)
)
return _dedupe_model_candidates(candidates)
def _extract_model_candidates_from_request(
request_data: dict,
route: str,
request_headers: Optional[Mapping[str, Any]] = None,
request_query_params: Optional[Mapping[str, Any]] = None,
) -> List[str]:
candidates: List[str] = []
uses_model_routing_sources = _route_uses_model_routing_sources(route=route)
uses_header_or_query_model_sources = _route_matches_any_marker(
route=route, markers=_MODEL_ROUTING_HEADER_OR_QUERY_ROUTE_MARKERS
)
uses_query_target_model_sources = _route_matches_any_marker(
route=route, markers=_MODEL_ROUTING_QUERY_TARGET_MODEL_ROUTE_MARKERS
)
uses_body_target_model_sources = _route_matches_any_marker(
route=route, markers=_MODEL_ROUTING_BODY_TARGET_MODEL_ROUTE_MARKERS
)
uses_completion_model_sources = _route_matches_any_marker(
route=route, markers=_MODEL_ROUTING_COMPLETION_MODEL_ROUTE_MARKERS
)
body_model = request_data.get("model")
_append_model_candidates(candidates, body_model)
if uses_body_target_model_sources or not body_model:
_append_model_candidates(candidates, request_data.get("target_model_names"))
if uses_completion_model_sources and isinstance(
request_data.get("completion"), dict
):
_append_model_candidates(candidates, request_data["completion"].get("model"))
if uses_model_routing_sources:
if uses_header_or_query_model_sources:
_append_model_candidates(
candidates,
_get_case_insensitive_mapping_value(request_query_params, "model"),
)
_append_model_candidates(
candidates,
_get_case_insensitive_mapping_value(
request_headers, MODEL_ROUTING_HEADER_NAME
),
)
if uses_query_target_model_sources:
_append_model_candidates(
candidates,
_get_case_insensitive_mapping_value(
request_query_params, "target_model_names"
),
)
for field in _MODEL_ROUTING_ID_FIELDS:
_append_model_candidates(
candidates,
_extract_models_from_managed_resource_id(request_data.get(field)),
)
return _dedupe_model_candidates(candidates)
def _format_model_candidates(
candidates: List[str],
) -> Optional[Union[str, List[str]]]:
# First try to get model from request_data
model = request_data.get("model") or request_data.get("target_model_names")
if not candidates:
return None
if len(candidates) == 1:
return candidates[0]
return candidates
if model is not None:
model_names = model.split(",")
if len(model_names) == 1:
model = model_names[0].strip()
else:
model = [m.strip() for m in model_names]
# If model not in request_data, try to extract from route
def get_model_from_request(
request_data: dict,
route: str,
request_headers: Optional[Mapping[str, Any]] = None,
request_query_params: Optional[Mapping[str, Any]] = None,
) -> Optional[Union[str, List[str]]]:
candidates = _extract_model_candidates_from_request(
request_data=request_data,
route=route,
request_headers=request_headers,
request_query_params=request_query_params,
)
model = _format_model_candidates(candidates)
# If no explicit model was found, try to extract from route
if model is None:
# Parse model from route that follows the pattern /openai/deployments/{model}/*
match = re.match(r"/openai/deployments/([^/]+)", route)

117
litellm/proxy/auth/user_api_key_auth.py
View file

@ -11,7 +11,7 @@ import asyncio
import re
import secrets
from datetime import datetime, timezone
from typing import Any, List, Optional, Tuple, cast
from typing import Any, List, Optional, Tuple, Union, cast
import fastapi
from fastapi import HTTPException, Request, WebSocket, status
@ -63,6 +63,7 @@ from litellm.proxy.common_utils.cache_coordinator import EventDrivenCacheCoordin
from litellm.proxy.common_utils.http_parsing_utils import (
_read_request_body,
_safe_get_request_headers,
_safe_get_request_query_params,
populate_request_with_path_params,
)
from litellm.proxy.common_utils.realtime_utils import _realtime_request_body
@ -118,6 +119,29 @@ azure_apim_header = APIKeyHeader(
)
def _get_model_from_request_context(
request_data: dict,
route: str,
request: Optional[Request],
) -> Optional[Union[str, List[str]]]:
return get_model_from_request(
request_data=request_data,
route=route,
request_headers=_safe_get_request_headers(request=request),
request_query_params=_safe_get_request_query_params(request=request),
)
def _get_model_names_for_budget_checks(
model: Optional[Union[str, List[str]]],
) -> List[str]:
if model is None:
return []
if isinstance(model, str):
return [model]
return model
def _get_bearer_token_or_received_api_key(api_key: str) -> str:
if api_key.startswith("Bearer "): # ensure Bearer token passed in
api_key = api_key.replace("Bearer ", "") # extract the token
@ -884,7 +908,11 @@ async def _user_api_key_auth_builder( # noqa: PLR0915
)
# Check if model has zero cost - if so, skip all budget checks
model = get_model_from_request(request_data, route)
model = _get_model_from_request_context(
request_data=request_data,
route=route,
request=request,
)
skip_budget_checks = False
if model is not None and llm_router is not None:
from litellm.proxy.auth.auth_checks import _is_model_cost_zero
@ -1252,6 +1280,7 @@ async def _user_api_key_auth_builder( # noqa: PLR0915
valid_token=valid_token,
request_data=request_data,
route=route,
request=request,
llm_model_list=llm_model_list,
llm_router=llm_router,
)
@ -1277,7 +1306,11 @@ async def _user_api_key_auth_builder( # noqa: PLR0915
user_obj = None
# Check 2a. Check if model has zero cost - if so, skip all budget checks
model = get_model_from_request(request_data, route)
model = _get_model_from_request_context(
request_data=request_data,
route=route,
request=request,
)
skip_budget_checks = False
if model is not None and llm_router is not None:
from litellm.proxy.auth.auth_checks import _is_model_cost_zero
@ -1395,21 +1428,29 @@ async def _user_api_key_auth_builder( # noqa: PLR0915
# Check 5. Token Model Spend is under Model budget
max_budget_per_model = valid_token.model_max_budget
current_model = request_data.get("model", None)
current_model = _get_model_from_request_context(
request_data=request_data,
route=route,
request=request,
)
current_models = _get_model_names_for_budget_checks(
model=current_model
)
if (
max_budget_per_model is not None
and isinstance(max_budget_per_model, dict)
and len(max_budget_per_model) > 0
and prisma_client is not None
and current_model is not None
and current_models
and valid_token.token is not None
):
## GET THE SPEND FOR THIS MODEL
await model_max_budget_limiter.is_key_within_model_budget(
user_api_key_dict=valid_token,
model=current_model,
)
for model_name in current_models:
await model_max_budget_limiter.is_key_within_model_budget(
user_api_key_dict=valid_token,
model=model_name,
)
# Check 5b. End-user model max budget
end_user_mmb = valid_token.end_user_model_max_budget
@ -1417,14 +1458,15 @@ async def _user_api_key_auth_builder( # noqa: PLR0915
end_user_mmb is not None
and isinstance(end_user_mmb, dict)
and len(end_user_mmb) > 0
and current_model is not None
and current_models
and valid_token.end_user_id is not None
):
await model_max_budget_limiter.is_end_user_within_model_budget(
end_user_id=valid_token.end_user_id,
end_user_model_max_budget=end_user_mmb,
model=current_model,
)
for model_name in current_models:
await model_max_budget_limiter.is_end_user_within_model_budget(
end_user_id=valid_token.end_user_id,
end_user_model_max_budget=end_user_mmb,
model=model_name,
)
# Check 6: Additional Common Checks across jwt + key auth
if valid_token.team_id is not None:
@ -1851,7 +1893,11 @@ async def _run_centralized_common_checks(
user_api_key_auth_obj.project_alias = project_object.project_alias
skip_budget_checks = False
model = get_model_from_request(request_data, route)
model = _get_model_from_request_context(
request_data=request_data,
route=route,
request=request,
)
if model is not None and llm_router is not None:
skip_budget_checks = _is_model_cost_zero(model=model, llm_router=llm_router)
@ -2122,6 +2168,7 @@ async def _enforce_key_and_fallback_model_access(
valid_token: UserAPIKeyAuth,
request_data: dict,
route: str,
request: Optional[Request],
llm_model_list: Optional[list],
llm_router: Optional[Any],
) -> None:
@ -2140,7 +2187,11 @@ async def _enforce_key_and_fallback_model_access(
):
pass
else:
model = get_model_from_request(request_data, route)
model = _get_model_from_request_context(
request_data=request_data,
route=route,
request=request,
)
fallback_models = cast(
Optional[List[ALL_FALLBACK_MODEL_VALUES]],
request_data.get("fallbacks", None),
@ -2227,11 +2278,17 @@ async def _run_post_custom_auth_checks(
valid_token=valid_token,
request_data=request_data,
route=route,
request=request,
llm_model_list=llm_model_list,
llm_router=llm_router,
)
current_model = request_data.get("model", None)
current_model = _get_model_from_request_context(
request_data=request_data,
route=route,
request=request,
)
current_models = _get_model_names_for_budget_checks(model=current_model)
# 3. Check key-level model_max_budget
max_budget_per_model = valid_token.model_max_budget
@ -2239,13 +2296,14 @@ async def _run_post_custom_auth_checks(
max_budget_per_model is not None
and isinstance(max_budget_per_model, dict)
and len(max_budget_per_model) > 0
and current_model is not None
and current_models
and valid_token.token is not None
):
await model_max_budget_limiter.is_key_within_model_budget(
user_api_key_dict=valid_token,
model=current_model,
)
for model_name in current_models:
await model_max_budget_limiter.is_key_within_model_budget(
user_api_key_dict=valid_token,
model=model_name,
)
# 4. Check end-user model_max_budget
end_user_mmb = valid_token.end_user_model_max_budget
@ -2253,14 +2311,15 @@ async def _run_post_custom_auth_checks(
end_user_mmb is not None
and isinstance(end_user_mmb, dict)
and len(end_user_mmb) > 0
and current_model is not None
and current_models
and valid_token.end_user_id is not None
):
await model_max_budget_limiter.is_end_user_within_model_budget(
end_user_id=valid_token.end_user_id,
end_user_model_max_budget=end_user_mmb,
model=current_model,
)
for model_name in current_models:
await model_max_budget_limiter.is_end_user_within_model_budget(
end_user_id=valid_token.end_user_id,
end_user_model_max_budget=end_user_mmb,
model=model_name,
)
# team / user / end_user / project context objects are fetched by
# the centralized common_checks gate in user_api_key_auth after

112
tests/test_litellm/proxy/auth/test_auth_utils.py
View file

@ -2,6 +2,7 @@
Unit tests for auth_utils functions related to rate limiting and customer ID extraction.
"""
import base64
from typing import Optional
from unittest.mock import MagicMock, patch
@ -258,6 +259,117 @@ def test_get_model_from_request_vertex_passthrough_still_works():
assert get_model_from_request(request_data={}, route=route) == "gemini-1.5-pro"
def test_get_model_from_request_includes_file_endpoint_header_model():
assert (
get_model_from_request(
request_data={},
route="/v1/files",
request_headers={"X-LiteLLM-Model": "restricted-model"},
)
== "restricted-model"
)
def test_get_model_from_request_ignores_routing_header_on_standard_llm_routes():
assert (
get_model_from_request(
request_data={"model": "allowed-model"},
route="/v1/chat/completions",
request_headers={"x-litellm-model": "restricted-model"},
)
== "allowed-model"
)
def test_get_model_from_request_authorizes_all_file_routing_model_sources():
models = get_model_from_request(
request_data={"model": "body-model"},
route="/v1/files",
request_headers={"x-litellm-model": "header-model"},
request_query_params={"target_model_names": "query-model-a,query-model-b"},
)
assert isinstance(models, list)
assert set(models) == {
"body-model",
"query-model-a",
"query-model-b",
"header-model",
}
def test_get_model_from_request_extracts_simple_encoded_file_id_model():
from litellm.proxy.openai_files_endpoints.common_utils import (
encode_file_id_with_model,
)
file_id = encode_file_id_with_model(
file_id="file-provider-id",
model="restricted-model",
)
assert (
get_model_from_request(
request_data={"file_id": file_id},
route="/v1/files/{file_id}",
)
== "restricted-model"
)
def test_get_model_from_request_extracts_unified_file_id_models():
raw_unified_file_id = (
"litellm_proxy:application/octet-stream;unified_id,test-id;"
"target_model_names,model-a,model-b;llm_output_file_id,file-provider-id"
)
encoded_unified_file_id = (
base64.urlsafe_b64encode(raw_unified_file_id.encode()).decode().rstrip("=")
)
assert get_model_from_request(
request_data={"file_id": encoded_unified_file_id},
route="/v1/files/{file_id}",
) == ["model-a", "model-b"]
def test_get_model_from_request_extracts_eval_completion_model():
assert (
get_model_from_request(
request_data={"completion": {"model": "judge-model"}},
route="/v1/evals/{eval_id}/runs",
)
== "judge-model"
)
def test_get_model_from_request_includes_fine_tuning_target_model_query():
assert (
get_model_from_request(
request_data={},
route="/v1/fine_tuning/jobs",
request_query_params={"target_model_names": "fine-tune-model"},
)
== "fine-tune-model"
)
def test_get_model_from_request_extracts_video_id_model():
from litellm.types.videos.utils import encode_video_id_with_provider
video_id = encode_video_id_with_provider(
video_id="video-provider-id",
provider="openai",
model_id="video-model",
)
assert (
get_model_from_request(
request_data={"video_id": video_id},
route="/v1/videos/{video_id}",
)
== "video-model"
)
def test_get_customer_user_header_returns_none_when_no_customer_role():
from litellm.proxy.auth.auth_utils import get_customer_user_header_from_mapping

60
tests/test_litellm/proxy/auth/test_user_api_key_auth.py
View file

@ -1,8 +1,7 @@
import asyncio
import json
import os
import sys
from typing import Tuple
from types import SimpleNamespace
from unittest.mock import ANY, AsyncMock, MagicMock, patch
sys.path.insert(
@ -32,6 +31,13 @@ from litellm.proxy.auth.user_api_key_auth import (
)
class _RoutingRequest:
def __init__(self, headers=None, query_params=None):
self.headers = headers or {}
self.query_params = query_params or {}
self.state = SimpleNamespace()
def test_get_api_key():
bearer_token = "Bearer sk-12345678"
api_key = "sk-12345678"
@ -107,6 +113,39 @@ async def test_custom_auth_honors_key_level_model_access_restriction_allowed_wit
)
@pytest.mark.asyncio
async def test_custom_auth_enforces_key_model_access_from_file_route_header_with_opt_in():
valid_token = UserAPIKeyAuth(token="test_token", models=["allowed-model"])
request = _RoutingRequest(headers={"x-litellm-model": "restricted-model"})
with (
patch(
"litellm.proxy.auth.user_api_key_auth.can_key_call_model",
new_callable=AsyncMock,
) as mock_can_key,
patch(
"litellm.proxy.auth.user_api_key_auth.common_checks", new_callable=AsyncMock
),
patch(
"litellm.proxy.proxy_server.general_settings",
{"custom_auth_run_common_checks": True},
),
):
await _run_post_custom_auth_checks(
valid_token=valid_token,
request=request,
request_data={},
route="/v1/files",
parent_otel_span=None,
)
mock_can_key.assert_awaited_once_with(
model="restricted-model",
llm_model_list=ANY,
valid_token=valid_token,
llm_router=ANY,
)
@pytest.mark.asyncio
async def test_custom_auth_honors_key_level_model_access_restriction_denied_with_opt_in():
valid_token = UserAPIKeyAuth(token="test_token", models=["gpt-4o-mini"])
@ -1752,7 +1791,11 @@ async def test_team_metadata_refreshed_from_team_object_during_auth():
from starlette.datastructures import URL
from starlette.requests import Request
from litellm.proxy._types import LiteLLM_TeamTableCachedObj, LitellmUserRoles, UserAPIKeyAuth
from litellm.proxy._types import (
LiteLLM_TeamTableCachedObj,
LitellmUserRoles,
UserAPIKeyAuth,
)
from litellm.proxy.auth.user_api_key_auth import _user_api_key_auth_builder
api_key = "sk-test-team-metadata-refresh"
@ -1833,16 +1876,17 @@ async def test_team_metadata_refreshed_from_team_object_during_auth():
request_data={},
)
assert result.team_metadata == {"guardrails": ["test-guardrail-333"]}, (
f"team_metadata was not updated from fresh team object. Got: {result.team_metadata}"
)
assert result.team_metadata == {
"guardrails": ["test-guardrail-333"]
}, f"team_metadata was not updated from fresh team object. Got: {result.team_metadata}"
finally:
for k, v in _originals.items():
setattr(_proxy_server_mod, k, v)
# ---------------------------------------------------------------------------
# _run_centralized_common_checks — centralized authz gate
# ---------------------------------------------------------------------------