tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-18 00:23:48 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #26963 from stuxf/codex/file-endpoint-model-auth

Browse Source 
chore(proxy): align resource model auth checks
This commit is contained in:
yuneng-jiang
2026-05-01 14:22:27 -07:00
committed by GitHub
parent 42122b83f5 9e501edece
commit 2fca7ad3ef
9 changed files with 800 additions and 84 deletions
Download Patch File Download Diff File Expand all files Collapse all files
litellm/proxy/_lazy_openapi_snapshot.json
+31 -31
View File
@@ -14008,7 +14008,7 @@
"/mcp-rest/test/connection": {
"post": {
"description": "Test if we can connect to the provided MCP server before adding it",
"operationId": "test_connection_mcp_rest_test_connection_post",
"operationId": "test_connection_mcp_rest_test_connection_post_2",
"requestBody": {
"content": {
"application/json": {
@@ -14053,7 +14053,7 @@
"/mcp-rest/test/tools/list": {
"post": {
"description": "Preview tools available from MCP server before adding it",
"operationId": "test_tools_list_mcp_rest_test_tools_list_post",
"operationId": "test_tools_list_mcp_rest_test_tools_list_post_2",
"requestBody": {
"content": {
"application/json": {
@@ -14098,7 +14098,7 @@
"/mcp-rest/tools/call": {
"post": {
"description": "REST API to call a specific MCP tool with the provided arguments",
"operationId": "call_tool_rest_api_mcp_rest_tools_call_post",
"operationId": "call_tool_rest_api_mcp_rest_tools_call_post_2",
"responses": {
"200": {
"content": {
@@ -14123,7 +14123,7 @@
"/mcp-rest/tools/list": {
"get": {
"description": "List all available tools with information about the server they belong to.\n\nExample response:\n{\n \"tools\": [\n {\n \"name\": \"create_zap\",\n \"description\": \"Create a new zap\",\n \"inputSchema\": \"tool_input_schema\",\n \"mcp_info\": {\n \"server_name\": \"zapier\",\n \"logo_url\": \"https://www.zapier.com/logo.png\",\n }\n }\n ],\n \"error\": null,\n \"message\": \"Successfully retrieved tools\"\n}",
"operationId": "list_tool_rest_api_mcp_rest_tools_list_get",
"operationId": "list_tool_rest_api_mcp_rest_tools_list_get_2",
"parameters": [
{
"description": "The server id to list tools for",
@@ -21896,7 +21896,7 @@
"/policies/usage/overview": {
"get": {
"description": "Return policy performance overview for the dashboard.",
"operationId": "policies_usage_overview_policies_usage_overview_get",
"operationId": "policies_usage_overview_policies_usage_overview_get_2",
"parameters": [
{
"description": "YYYY-MM-DD",
@@ -22521,7 +22521,7 @@
"/policies/attachments/estimate-impact": {
"post": {
"description": "Estimate how many keys and teams would be affected by a policy attachment.\n\nUse this before creating an attachment to preview the blast radius.\n\nExample Request:\n```bash\ncurl -X POST \"http://localhost:4000/policies/attachments/estimate-impact\" \\\n -H \"Authorization: Bearer <your_api_key>\" \\\n -H \"Content-Type: application/json\" \\\n -d '{\n \"policy_name\": \"hipaa-compliance\",\n \"tags\": [\"healthcare\", \"health-*\"]\n }'\n```",
"operationId": "estimate_attachment_impact_policies_attachments_estimate_impact_post",
"operationId": "estimate_attachment_impact_policies_attachments_estimate_impact_post_2",
"requestBody": {
"content": {
"application/json": {
@@ -22568,7 +22568,7 @@
"/policies/resolve": {
"post": {
"description": "Resolve which policies and guardrails apply for a given context.\n\nUse this endpoint to debug \"what guardrails would apply to a request\nwith this team/key/model/tags combination?\"\n\nExample Request:\n```bash\ncurl -X POST \"http://localhost:4000/policies/resolve\" \\\n -H \"Authorization: Bearer <your_api_key>\" \\\n -H \"Content-Type: application/json\" \\\n -d '{\n \"tags\": [\"healthcare\"],\n \"model\": \"gpt-4\"\n }'\n```",
"operationId": "resolve_policies_for_context_policies_resolve_post",
"operationId": "resolve_policies_for_context_policies_resolve_post_2",
"parameters": [
{
"description": "Force a DB sync before resolving. Default uses in-memory cache.",
@@ -28329,7 +28329,7 @@
"/v1/vector_stores": {
"get": {
"description": "List vector stores.\n\nAPI Reference:\nhttps://platform.openai.com/docs/api-reference/vector-stores/list",
"operationId": "vector_store_list_v1_vector_stores_get",
"operationId": "vector_store_list_v1_vector_stores_get_2",
"parameters": [
{
"in": "query",
@@ -28430,7 +28430,7 @@
},
"post": {
"description": "Create a vector store.\n\nAPI Reference:\nhttps://platform.openai.com/docs/api-reference/vector-stores/create\n\nSupports target_model_names parameter for creating vector stores across multiple models:\n```json\n{\n \"name\": \"my-vector-store\",\n \"target_model_names\": \"gpt-4,gemini-2.0\"\n}\n```",
"operationId": "vector_store_create_v1_vector_stores_post",
"operationId": "vector_store_create_v1_vector_stores_post_2",
"responses": {
"200": {
"content": {
@@ -28455,7 +28455,7 @@
"/v1/vector_stores/{vector_store_id}": {
"delete": {
"description": "Delete a vector store.\n\nAPI Reference:\nhttps://platform.openai.com/docs/api-reference/vector-stores/delete",
"operationId": "vector_store_delete_v1_vector_stores__vector_store_id__delete",
"operationId": "vector_store_delete_v1_vector_stores__vector_store_id__delete_2",
"parameters": [
{
"in": "path",
@@ -28499,7 +28499,7 @@
},
"get": {
"description": "Retrieve a vector store.\n\nAPI Reference:\nhttps://platform.openai.com/docs/api-reference/vector-stores/retrieve",
"operationId": "vector_store_retrieve_v1_vector_stores__vector_store_id__get",
"operationId": "vector_store_retrieve_v1_vector_stores__vector_store_id__get_2",
"parameters": [
{
"in": "path",
@@ -28543,7 +28543,7 @@
},
"post": {
"description": "Update a vector store.\n\nAPI Reference:\nhttps://platform.openai.com/docs/api-reference/vector-stores/modify",
"operationId": "vector_store_update_v1_vector_stores__vector_store_id__post",
"operationId": "vector_store_update_v1_vector_stores__vector_store_id__post_2",
"parameters": [
{
"in": "path",
@@ -28588,7 +28588,7 @@
},
"/v1/vector_stores/{vector_store_id}/files": {
"get": {
"operationId": "vector_store_file_list_v1_vector_stores__vector_store_id__files_get",
"operationId": "vector_store_file_list_v1_vector_stores__vector_store_id__files_get_2",
"parameters": [
{
"in": "path",
@@ -28631,7 +28631,7 @@
]
},
"post": {
"operationId": "vector_store_file_create_v1_vector_stores__vector_store_id__files_post",
"operationId": "vector_store_file_create_v1_vector_stores__vector_store_id__files_post_2",
"parameters": [
{
"in": "path",
@@ -28676,7 +28676,7 @@
},
"/v1/vector_stores/{vector_store_id}/files/{file_id}": {
"delete": {
"operationId": "vector_store_file_delete_v1_vector_stores__vector_store_id__files__file_id__delete",
"operationId": "vector_store_file_delete_v1_vector_stores__vector_store_id__files__file_id__delete_2",
"parameters": [
{
"in": "path",
@@ -28728,7 +28728,7 @@
]
},
"get": {
"operationId": "vector_store_file_retrieve_v1_vector_stores__vector_store_id__files__file_id__get",
"operationId": "vector_store_file_retrieve_v1_vector_stores__vector_store_id__files__file_id__get_2",
"parameters": [
{
"in": "path",
@@ -28780,7 +28780,7 @@
]
},
"post": {
"operationId": "vector_store_file_update_v1_vector_stores__vector_store_id__files__file_id__post",
"operationId": "vector_store_file_update_v1_vector_stores__vector_store_id__files__file_id__post_2",
"parameters": [
{
"in": "path",
@@ -28834,7 +28834,7 @@
},
"/v1/vector_stores/{vector_store_id}/files/{file_id}/content": {
"get": {
"operationId": "vector_store_file_content_v1_vector_stores__vector_store_id__files__file_id__content_get",
"operationId": "vector_store_file_content_v1_vector_stores__vector_store_id__files__file_id__content_get_2",
"parameters": [
{
"in": "path",
@@ -28889,7 +28889,7 @@
"/v1/vector_stores/{vector_store_id}/search": {
"post": {
"description": "Search a vector store.\n\nAPI Reference:\nhttps://platform.openai.com/docs/api-reference/vector-stores/search",
"operationId": "vector_store_search_v1_vector_stores__vector_store_id__search_post",
"operationId": "vector_store_search_v1_vector_stores__vector_store_id__search_post_2",
"parameters": [
{
"in": "path",
@@ -28935,7 +28935,7 @@
"/vector_stores": {
"get": {
"description": "List vector stores.\n\nAPI Reference:\nhttps://platform.openai.com/docs/api-reference/vector-stores/list",
"operationId": "vector_store_list_vector_stores_get",
"operationId": "vector_store_list_vector_stores_get_2",
"parameters": [
{
"in": "query",
@@ -29036,7 +29036,7 @@
},
"post": {
"description": "Create a vector store.\n\nAPI Reference:\nhttps://platform.openai.com/docs/api-reference/vector-stores/create\n\nSupports target_model_names parameter for creating vector stores across multiple models:\n```json\n{\n \"name\": \"my-vector-store\",\n \"target_model_names\": \"gpt-4,gemini-2.0\"\n}\n```",
"operationId": "vector_store_create_vector_stores_post",
"operationId": "vector_store_create_vector_stores_post_2",
"responses": {
"200": {
"content": {
@@ -29061,7 +29061,7 @@
"/vector_stores/{vector_store_id}": {
"delete": {
"description": "Delete a vector store.\n\nAPI Reference:\nhttps://platform.openai.com/docs/api-reference/vector-stores/delete",
"operationId": "vector_store_delete_vector_stores__vector_store_id__delete",
"operationId": "vector_store_delete_vector_stores__vector_store_id__delete_2",
"parameters": [
{
"in": "path",
@@ -29105,7 +29105,7 @@
},
"get": {
"description": "Retrieve a vector store.\n\nAPI Reference:\nhttps://platform.openai.com/docs/api-reference/vector-stores/retrieve",
"operationId": "vector_store_retrieve_vector_stores__vector_store_id__get",
"operationId": "vector_store_retrieve_vector_stores__vector_store_id__get_2",
"parameters": [
{
"in": "path",
@@ -29149,7 +29149,7 @@
},
"post": {
"description": "Update a vector store.\n\nAPI Reference:\nhttps://platform.openai.com/docs/api-reference/vector-stores/modify",
"operationId": "vector_store_update_vector_stores__vector_store_id__post",
"operationId": "vector_store_update_vector_stores__vector_store_id__post_2",
"parameters": [
{
"in": "path",
@@ -29194,7 +29194,7 @@
},
"/vector_stores/{vector_store_id}/files": {
"get": {
"operationId": "vector_store_file_list_vector_stores__vector_store_id__files_get",
"operationId": "vector_store_file_list_vector_stores__vector_store_id__files_get_2",
"parameters": [
{
"in": "path",
@@ -29237,7 +29237,7 @@
]
},
"post": {
"operationId": "vector_store_file_create_vector_stores__vector_store_id__files_post",
"operationId": "vector_store_file_create_vector_stores__vector_store_id__files_post_2",
"parameters": [
{
"in": "path",
@@ -29282,7 +29282,7 @@
},
"/vector_stores/{vector_store_id}/files/{file_id}": {
"delete": {
"operationId": "vector_store_file_delete_vector_stores__vector_store_id__files__file_id__delete",
"operationId": "vector_store_file_delete_vector_stores__vector_store_id__files__file_id__delete_2",
"parameters": [
{
"in": "path",
@@ -29334,7 +29334,7 @@
]
},
"get": {
"operationId": "vector_store_file_retrieve_vector_stores__vector_store_id__files__file_id__get",
"operationId": "vector_store_file_retrieve_vector_stores__vector_store_id__files__file_id__get_2",
"parameters": [
{
"in": "path",
@@ -29386,7 +29386,7 @@
]
},
"post": {
"operationId": "vector_store_file_update_vector_stores__vector_store_id__files__file_id__post",
"operationId": "vector_store_file_update_vector_stores__vector_store_id__files__file_id__post_2",
"parameters": [
{
"in": "path",
@@ -29440,7 +29440,7 @@
},
"/vector_stores/{vector_store_id}/files/{file_id}/content": {
"get": {
"operationId": "vector_store_file_content_vector_stores__vector_store_id__files__file_id__content_get",
"operationId": "vector_store_file_content_vector_stores__vector_store_id__files__file_id__content_get_2",
"parameters": [
{
"in": "path",
@@ -29495,7 +29495,7 @@
"/vector_stores/{vector_store_id}/search": {
"post": {
"description": "Search a vector store.\n\nAPI Reference:\nhttps://platform.openai.com/docs/api-reference/vector-stores/search",
"operationId": "vector_store_search_vector_stores__vector_store_id__search_post",
"operationId": "vector_store_search_vector_stores__vector_store_id__search_post_2",
"parameters": [
{
"in": "path",
litellm/proxy/_lazy_openapi_snapshot.py
+4 -2
View File
@@ -10,7 +10,7 @@ any drift as a neutral check.
import json
import sys
from pathlib import Path
from typing import Dict, Optional
from typing import Dict, Optional, Set
SNAPSHOT_FILE = Path(__file__).parent / "_lazy_openapi_snapshot.json"
HTTP_METHODS = {"delete", "get", "head", "options", "patch", "post", "put"}
@@ -65,7 +65,7 @@ def generate_snapshot() -> Dict[str, Dict]:
from fastapi.openapi.utils import get_openapi
from litellm.proxy._lazy_features import LAZY_FEATURES
from litellm.proxy.proxy_server import app
from litellm.proxy.proxy_server import app, ensure_unique_openapi_operation_ids
for feat in LAZY_FEATURES:
if feat.module_path in sys.modules:
@@ -77,6 +77,7 @@ def generate_snapshot() -> Dict[str, Dict]:
sys.stderr.write(f"warning: skip {feat.name}: {exc}\n")
fragments: Dict[str, Dict] = {}
used_operation_ids: Set[str] = set()
for feat in LAZY_FEATURES:
feat_routes = [
r
@@ -93,6 +94,7 @@ def generate_snapshot() -> Dict[str, Dict]:
for op in path_ops.values():
if isinstance(op, dict):
op["tags"] = [feat.name]
full = ensure_unique_openapi_operation_ids(full, used_operation_ids)
fragments[feat.name] = {
"paths": paths,
"components": {"schemas": full.get("components", {}).get("schemas", {})},
litellm/proxy/auth/auth_checks.py
+8 -1
View File
@@ -60,6 +60,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,
@@ -486,7 +490,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
litellm/proxy/auth/auth_utils.py
+249 -12
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
@@ -976,20 +976,257 @@ def get_end_user_id_from_request_body(
return None
def get_model_from_request(
request_data: dict, route: 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")
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",
)
if model is not None:
model_names = model.split(",")
if len(model_names) == 1:
model = model_names[0].strip()
def _append_model_candidates(candidates: List[str], value: Any) -> None:
if value is None:
return
values = value if isinstance(value, (list, tuple, set)) else [value]
for item in values:
if item is None:
continue
if isinstance(item, str):
model_names = [model.strip() for model in item.split(",")]
else:
model = [m.strip() for m in model_names]
model_names = [str(item).strip()]
candidates.extend(model for model in model_names if model)
# If model not in request_data, try to extract from route
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, resource_id_field: Optional[str] = None
) -> 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)
)
if resource_id_field in ("video_id", "character_id"):
try:
from litellm.types.videos.utils import (
decode_character_id_with_provider,
decode_video_id_with_provider,
)
if resource_id_field == "video_id":
_append_model_candidates(
candidates=candidates,
value=decode_video_id_with_provider(resource_id).get("model_id"),
)
else:
_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), resource_id_field=field
),
)
return _dedupe_model_candidates(candidates)
def _format_model_candidates(
candidates: List[str],
) -> Optional[Union[str, List[str]]]:
if not candidates:
return None
if len(candidates) == 1:
return candidates[0]
return candidates
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)
litellm/proxy/auth/user_api_key_auth.py
+88 -29
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.user_api_key_cache import UserApiKeyCache
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
@@ -1254,6 +1282,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,
)
@@ -1279,7 +1308,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
@@ -1403,21 +1436,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
@@ -1425,14 +1466,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:
@@ -1863,7 +1905,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)
@@ -2134,6 +2180,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:
@@ -2152,7 +2199,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),
@@ -2239,11 +2290,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
@@ -2251,13 +2308,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
@@ -2265,14 +2323,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
litellm/proxy/proxy_server.py
+83
View File
@@ -6,6 +6,7 @@ import inspect
import io
import os
import random
import re
import secrets
import shutil
import subprocess
@@ -955,6 +956,85 @@ async def proxy_startup_event(app: FastAPI): # noqa: PLR0915
await proxy_shutdown_event() # type: ignore[reportGeneralTypeIssues]
def _generate_stable_operation_id(route: Any) -> str:
operation_id = re.sub(r"\W", "_", f"{route.name}{route.path_format}")
route_methods = sorted(route.methods or [])
if len(route_methods) == 1:
operation_id = f"{operation_id}_{route_methods[0].lower()}"
return operation_id
_OPENAPI_HTTP_METHODS = {
"delete",
"get",
"head",
"options",
"patch",
"post",
"put",
"trace",
}
def _strip_operation_id_method_suffix(operation_id: str) -> str:
base, separator, suffix = operation_id.rpartition("_")
if separator and suffix in _OPENAPI_HTTP_METHODS:
return base
return operation_id
def ensure_unique_openapi_operation_ids(
openapi_schema: Dict[str, Any],
reserved_operation_ids: Optional[Set[str]] = None,
) -> Dict[str, Any]:
operation_entries = []
operation_id_counts: Dict[str, int] = {}
for path_item in openapi_schema.get("paths", {}).values():
if not isinstance(path_item, dict):
continue
for method, operation in path_item.items():
if method not in _OPENAPI_HTTP_METHODS or not isinstance(operation, dict):
continue
operation_id = operation.get("operationId")
if not isinstance(operation_id, str):
continue
operation_entries.append((method, operation, operation_id))
operation_id_counts[operation_id] = (
operation_id_counts.get(operation_id, 0) + 1
)
used_operation_ids = set(reserved_operation_ids or set())
seen_operation_ids: Set[str] = set()
for method, operation, operation_id in operation_entries:
should_rewrite = (
operation_id_counts[operation_id] > 1
or operation_id in used_operation_ids
or operation_id in seen_operation_ids
)
if not should_rewrite:
seen_operation_ids.add(operation_id)
used_operation_ids.add(operation_id)
continue
base_operation_id = _strip_operation_id_method_suffix(operation_id)
new_operation_id = f"{base_operation_id}_{method}"
suffix = 2
while (
new_operation_id in used_operation_ids
or new_operation_id in seen_operation_ids
):
new_operation_id = f"{base_operation_id}_{method}_{suffix}"
suffix += 1
operation["operationId"] = new_operation_id
seen_operation_ids.add(new_operation_id)
used_operation_ids.add(new_operation_id)
if reserved_operation_ids is not None:
reserved_operation_ids.update(used_operation_ids)
return openapi_schema
app = FastAPI(
docs_url=_get_docs_url(),
redoc_url=_get_redoc_url(),
@@ -964,6 +1044,7 @@ app = FastAPI(
version=version,
root_path=server_root_path,
lifespan=proxy_startup_event, # type: ignore[reportGeneralTypeIssues]
generate_unique_id_function=_generate_stable_operation_id,
)
vertex_live_passthrough_vertex_base = VertexBase()
@@ -1043,6 +1124,7 @@ def get_openapi_schema():
from litellm.proxy._lazy_features import inject_lazy_stubs
openapi_schema = inject_lazy_stubs(openapi_schema)
openapi_schema = ensure_unique_openapi_operation_ids(openapi_schema)
# Fix Swagger UI execute path error when server_root_path is set
if server_root_path:
@@ -1074,6 +1156,7 @@ def custom_openapi():
from litellm.proxy._lazy_features import inject_lazy_stubs
openapi_schema = inject_lazy_stubs(openapi_schema)
openapi_schema = ensure_unique_openapi_operation_ids(openapi_schema)
# Fix Swagger UI execute path error when server_root_path is set
if server_root_path:
tests/test_litellm/proxy/auth/test_auth_utils.py
+203 -1
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
@@ -10,11 +11,12 @@ import pytest
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.auth.auth_utils import (
_get_customer_id_from_standard_headers,
abbreviate_api_key,
check_complete_credentials,
get_end_user_id_from_request_body,
get_model_from_request,
get_key_model_rpm_limit,
get_key_model_tpm_limit,
get_model_from_request,
get_project_model_rpm_limit,
get_project_model_tpm_limit,
is_request_body_safe,
@@ -258,6 +260,206 @@ 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_openai_deployment_route_still_works():
assert (
get_model_from_request(
request_data={},
route="/openai/deployments/my-azure-deployment/chat/completions",
)
== "my-azure-deployment"
)
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_model_from_request_only_runs_media_decoders_for_matching_fields():
with (
patch(
"litellm.types.videos.utils.decode_video_id_with_provider",
return_value={"model_id": "video-model"},
) as video_decoder,
patch(
"litellm.types.videos.utils.decode_character_id_with_provider",
return_value={"model_id": "character-model"},
) as character_decoder,
):
assert (
get_model_from_request(
request_data={"file_id": "file-provider-id"},
route="/v1/files/{file_id}",
)
is None
)
video_decoder.assert_not_called()
character_decoder.assert_not_called()
assert (
get_model_from_request(
request_data={"video_id": "video-provider-id"},
route="/v1/videos/{video_id}",
)
== "video-model"
)
video_decoder.assert_called_once_with("video-provider-id")
character_decoder.assert_not_called()
video_decoder.reset_mock()
character_decoder.reset_mock()
assert (
get_model_from_request(
request_data={"character_id": "character-provider-id"},
route="/v1/videos/{character_id}",
)
== "character-model"
)
video_decoder.assert_not_called()
character_decoder.assert_called_once_with("character-provider-id")
def test_get_model_from_request_handles_managed_id_decoder_failures():
with (
patch(
"litellm.proxy.openai_files_endpoints.common_utils.decode_model_from_file_id",
side_effect=Exception("decode failed"),
),
patch(
"litellm.llms.base_llm.managed_resources.utils.parse_unified_id",
side_effect=Exception("parse failed"),
),
patch(
"litellm.types.videos.utils.decode_video_id_with_provider",
side_effect=Exception("video decode failed"),
),
):
assert (
get_model_from_request(
request_data={"file_id": "not-a-managed-resource-id"},
route="/v1/files/{file_id}",
)
is None
)
assert (
get_model_from_request(
request_data={"video_id": "not-a-managed-resource-id"},
route="/v1/videos/{video_id}",
)
is None
)
def test_abbreviate_api_key():
assert abbreviate_api_key("sk-test-1234") == "sk-...1234"
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
tests/test_litellm/proxy/auth/test_user_api_key_auth.py
+52 -8
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
# ---------------------------------------------------------------------------
tests/test_litellm/proxy/test_lazy_openapi_snapshot.py
+82
View File
@@ -1,3 +1,85 @@
import sys
from types import ModuleType, SimpleNamespace
def test_generate_snapshot_uses_shared_operation_id_reservations(monkeypatch):
from litellm.proxy import _lazy_openapi_snapshot
route_a = SimpleNamespace(path="/feature-a/items")
route_b = SimpleNamespace(path="/feature-b/items")
fake_app = SimpleNamespace(
title="LiteLLM test",
version="0.0.0",
routes=[route_a, route_b],
)
fake_feature_a_module = ModuleType("fake_feature_a")
fake_feature_b_module = ModuleType("fake_feature_b")
monkeypatch.setitem(sys.modules, "fake_feature_a", fake_feature_a_module)
monkeypatch.setitem(sys.modules, "fake_feature_b", fake_feature_b_module)
fake_lazy_features_module = ModuleType("litellm.proxy._lazy_features")
fake_lazy_features_module.LAZY_FEATURES = [
SimpleNamespace(
name="feature-a",
module_path="fake_feature_a",
path_prefixes=("/feature-a",),
register_fn=lambda app, module: None,
),
SimpleNamespace(
name="feature-b",
module_path="fake_feature_b",
path_prefixes=("/feature-b",),
register_fn=lambda app, module: None,
),
]
monkeypatch.setitem(
sys.modules, "litellm.proxy._lazy_features", fake_lazy_features_module
)
def fake_get_openapi(title, version, routes):
path = routes[0].path
return {
"paths": {path: {"get": {"operationId": "shared_operation_id_get"}}},
"components": {"schemas": {"Example": {"type": "object"}}},
}
def fake_ensure_unique_openapi_operation_ids(schema, reserved_operation_ids):
for path_item in schema["paths"].values():
operation = path_item["get"]
operation_id = operation["operationId"]
if operation_id in reserved_operation_ids:
operation_id = f"{operation_id}_2"
operation["operationId"] = operation_id
reserved_operation_ids.add(operation_id)
return schema
fake_proxy_server_module = ModuleType("litellm.proxy.proxy_server")
fake_proxy_server_module.app = fake_app
fake_proxy_server_module.ensure_unique_openapi_operation_ids = (
fake_ensure_unique_openapi_operation_ids
)
monkeypatch.setitem(
sys.modules, "litellm.proxy.proxy_server", fake_proxy_server_module
)
monkeypatch.setattr("fastapi.openapi.utils.get_openapi", fake_get_openapi)
fragments = _lazy_openapi_snapshot.generate_snapshot()
assert (
fragments["feature-a"]["paths"]["/feature-a/items"]["get"]["operationId"]
== "shared_operation_id_get"
)
assert (
fragments["feature-b"]["paths"]["/feature-b/items"]["get"]["operationId"]
== "shared_operation_id_get_2"
)
assert fragments["feature-a"]["paths"]["/feature-a/items"]["get"]["tags"] == [
"feature-a"
]
assert fragments["feature-b"]["paths"]["/feature-b/items"]["get"]["tags"] == [
"feature-b"
]
from litellm.proxy._lazy_openapi_snapshot import _normalize_operation_ids