diff --git a/enterprise/litellm_enterprise/integrations/custom_guardrail.py b/enterprise/litellm_enterprise/integrations/custom_guardrail.py index cbac4cb5b15..58c24219b47 100644 --- a/enterprise/litellm_enterprise/integrations/custom_guardrail.py +++ b/enterprise/litellm_enterprise/integrations/custom_guardrail.py @@ -20,9 +20,15 @@ class EnterpriseCustomGuardrailHelper: the tag's value (the mode for that tag). - If no request tag matches, fall back to default mode(s). """ + from litellm.litellm_core_utils.core_helpers import ( + get_metadata_variable_name_from_kwargs, + ) from litellm.litellm_core_utils.litellm_logging import ( StandardLoggingPayloadSetup, ) + from litellm.router_strategy.tag_based_routing import ( + without_baggage_mapped_tags, + ) from litellm.proxy._types import CommonProxyErrors from litellm.proxy.proxy_server import premium_user @@ -36,9 +42,13 @@ class EnterpriseCustomGuardrailHelper: proxy_server_request = data.get("proxy_server_request", {}) - request_tags = StandardLoggingPayloadSetup.get_request_tags( - litellm_params=data, - proxy_server_request=proxy_server_request, + request_tags = without_baggage_mapped_tags( + StandardLoggingPayloadSetup.get_request_tags( + litellm_params=data, + proxy_server_request=proxy_server_request, + ), + data.get(get_metadata_variable_name_from_kwargs(data)), + keep_serving_deployment_tags=True, ) # Check if any request tag matches a Mode tag key diff --git a/litellm/constants.py b/litellm/constants.py index c7673db6976..bc94760b187 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1645,6 +1645,7 @@ OUTPUT_TOKEN_CEILING_PARAMS: Final = frozenset({"max_tokens", "max_completion_to CLIENT_OUTPUT_CEILING_METADATA_KEY: Final = "_client_output_ceiling" CONSUMED_REQUEST_TAGS_METADATA_KEY: Final = "_consumed_request_tags" ROUTING_REQUEST_TAGS_METADATA_KEY: Final = "_routing_request_tags" +BAGGAGE_MAPPED_TAGS_METADATA_KEY: Final = "_baggage_mapped_tags" VECTOR_STORE_CALLER_METADATA_KWARG: Final = "litellm_vector_store_caller_metadata" ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY: Final = "_litellm_router_usage_counted_tokens" INTERNAL_CALL_ORIGIN_METADATA_KEY: Final = "internal_call_origin" diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index e4f77ac0005..523f478d771 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -86,6 +86,7 @@ from litellm.types.proxy.carried_budget_state import ( UserBudgetSnapshot as UserBudgetSnapshot, # noqa: PLC0414 # public re-export ) from litellm.types.proxy.control_plane_endpoints import WorkerRegistryEntry +from litellm.types.proxy.request_tag_mappings import RequestTagMapping from litellm.types.proxy.spend_capture_rate import SpendCaptureRateCheckSettings from litellm.types.router import AllowedModelRegion, RouterErrors, UpdateRouterConfig from litellm.types.router_weights import validate_router_settings_dict @@ -2457,6 +2458,20 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase): Documents all the fields supported by `general_settings` in config.yaml """ + request_tag_mappings: Sequence[RequestTagMapping] | None = Field( + None, + description="Map W3C baggage values to request tags. Exact value/tag rules take precedence over tag_prefix fallbacks.", + ) + + @field_validator("request_tag_mappings") + @classmethod + def validate_request_tag_mapping_duplicates( + cls, rules: Sequence[RequestTagMapping] | None + ) -> Sequence[RequestTagMapping] | None: + if rules is not None and len({(rule.key, rule.value) for rule in rules}) != len(rules): + raise ValueError("Only one mapping is allowed for each baggage key/value, including its any-value fallback") + return rules + completion_model: str | None = Field(None, description="proxy level default model for all chat completion calls") max_in_flight_requests_per_worker: int | None = Field( None, gt=0, description="maximum concurrent requests handled by each worker" diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 0f94d2952ec..5a09a0efbf6 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -877,18 +877,27 @@ def _enforce_user_param_check(general_settings: dict, request: Request, request_ raise Exception(f"'user' param not passed in. 'enforce_user_param'={general_settings['enforce_user_param']}") -def _reject_clientside_metadata_tags_check(general_settings: dict, request_body: dict, route: str) -> None: +def _reject_clientside_metadata_tags_check( + general_settings: dict, request_body: dict, route: str, request: Request | None = None +) -> None: if not general_settings.get("reject_clientside_metadata_tags", False): return - if ( - RouteChecks.is_llm_api_route(route=route) - and "metadata" in request_body - and isinstance(request_body["metadata"], dict) - and "tags" in request_body["metadata"] - ): + original_client_tags: Final = ( + getattr(request.state, "litellm_baggage_client_metadata_tags", None) if request is not None else None + ) + has_client_tags: Final = ( + original_client_tags + if isinstance(original_client_tags, bool) + else ( + "metadata" in request_body + and isinstance(request_body["metadata"], dict) + and "tags" in request_body["metadata"] + ) + ) + if RouteChecks.is_llm_api_route(route=route) and has_client_tags: raise ProxyException( - message=f"Client-side 'metadata.tags' not allowed in request. 'reject_clientside_metadata_tags'={general_settings['reject_clientside_metadata_tags']}. Tags can only be set via API key metadata.", + message=f"Client-side 'metadata.tags' not allowed in request. 'reject_clientside_metadata_tags'={general_settings['reject_clientside_metadata_tags']}. Tags can only be set via key or team metadata, or by configured request_tag_mappings.", type=ProxyErrorTypes.bad_request_error, param="metadata.tags", code=status.HTTP_400_BAD_REQUEST, @@ -1337,24 +1346,19 @@ async def common_checks( ) # Run before apply_key_tags_pre_auth injects key metadata.tags into request_body. - _reject_clientside_metadata_tags_check(general_settings, request_body, route) + _reject_clientside_metadata_tags_check(general_settings, request_body, route, request) # If this is a free model, skip all budget checks if not skip_all_budget_checks: # Key metadata.tags are injected into request_body here so the tag budget # check can read them; this mutation must run before the gathered checks. + from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup + if valid_token is not None: - from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup - - LiteLLMProxyRequestSetup.pre_seed_litellm_metadata_for_route( - request_data=request_body, - route=route, - ) - - LiteLLMProxyRequestSetup.apply_key_tags_pre_auth( - request_data=request_body, - user_api_key_dict=valid_token, - ) + for tag_target in LiteLLMProxyRequestSetup.pre_auth_tag_merge_targets(request, request_body): + LiteLLMProxyRequestSetup.pre_seed_litellm_metadata_for_route(request_data=tag_target, route=route) + LiteLLMProxyRequestSetup.apply_key_tags_pre_auth(request_data=tag_target, user_api_key_dict=valid_token) + LiteLLMProxyRequestSetup.apply_authenticated_baggage_tags(request, request_body) async def _user_max_budget_check() -> None: # 4.1 personal budget @@ -2456,7 +2460,7 @@ async def _fetch_uncached_tags( uncached_tags: Sequence[str], prisma_client: PrismaClient, user_api_key_cache: UserApiKeyCache, -) -> tuple[tuple[str, LiteLLM_TagTable], ...]: +) -> tuple[tuple[str, LiteLLM_TagTable], ...] | Exception: """Rows for the tags a cache probe missed; names absent from the registry never reach the DB.""" if not uncached_tags: return () @@ -2484,13 +2488,46 @@ async def _fetch_uncached_tags( model_type=LiteLLM_TagTable, ttl=get_management_object_ttl(user_api_key_cache), ) - except Exception as e: # noqa: BLE001 # fail-safe: a tag fetch error must yield "no budget objects", never break auth - verbose_proxy_logger.debug("Error batch fetching tags from database: %s", e) - return () + except Exception as e: # noqa: BLE001 # returned as a value: budget checks fail open, the baggage mapper fails closed + return e else: return fetched +async def _probe_and_fetch_tags( + tag_names: Sequence[str], + prisma_client: PrismaClient, + user_api_key_cache: UserApiKeyCache, +) -> tuple[Mapping[str, LiteLLM_TagTable], Exception | None]: + probed: Final = [ + ( + tag_name, + await user_api_key_cache.async_get_cache(key=tag_cache_key(tag_name), model_type=LiteLLM_TagTable), + ) + for tag_name in tag_names + ] + fetched: Final = await _fetch_uncached_tags( + uncached_tags=tuple(tag_name for tag_name, tag_obj in probed if tag_obj is None), + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + ) + rows: Final = () if isinstance(fetched, Exception) else fetched + found: Final = {tag_name: tag_obj for tag_name, tag_obj in (*probed, *rows) if tag_obj is not None} + return found, fetched if isinstance(fetched, Exception) else None + + +@log_db_metrics +@with_service_target(AUTH_OBJECTS_TARGET) +async def read_tag_objects_batch( + tag_names: Sequence[str], + prisma_client: PrismaClient, + user_api_key_cache: UserApiKeyCache, +) -> Mapping[str, LiteLLM_TagTable] | Exception: + """Like ``get_tag_objects_batch``, but a failed DB read comes back as the error instead of missing tags.""" + found, error = await _probe_and_fetch_tags(tag_names, prisma_client, user_api_key_cache) + return found if error is None else error + + @log_db_metrics @with_service_target(AUTH_OBJECTS_TARGET) async def get_tag_objects_batch( @@ -2521,19 +2558,10 @@ async def get_tag_objects_batch( if prisma_client is None or not tag_names: return {} - probed: Final = [ - ( - tag_name, - await user_api_key_cache.async_get_cache(key=tag_cache_key(tag_name), model_type=LiteLLM_TagTable), - ) - for tag_name in tag_names - ] - fetched: Final = await _fetch_uncached_tags( - uncached_tags=tuple(tag_name for tag_name, tag_obj in probed if tag_obj is None), - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - ) - return {tag_name: tag_obj for tag_name, tag_obj in (*probed, *fetched) if tag_obj is not None} + found, error = await _probe_and_fetch_tags(tag_names, prisma_client, user_api_key_cache) + if error is not None: + verbose_proxy_logger.debug("Error batch fetching tags from database: %s", error) + return dict(found) @log_db_metrics diff --git a/litellm/proxy/auth/route_checks.py b/litellm/proxy/auth/route_checks.py index f5e8304c83c..720bfabb95a 100644 --- a/litellm/proxy/auth/route_checks.py +++ b/litellm/proxy/auth/route_checks.py @@ -91,6 +91,21 @@ _SELF_SCOPED_CREDENTIAL_ROUTES: Final = ( ) +_UNBILLED_ROUTES: Final[frozenset[str]] = frozenset( + { + "/models", + "/v1/models", + "/utils/token_counter", + "/utils/transform_request", + "/responses/input_tokens", + "/v1/responses/input_tokens", + "/openai/v1/responses/input_tokens", + } +) +_TOKEN_COUNTING_SEGMENTS: Final[frozenset[str]] = frozenset({"count_tokens", "count-tokens"}) +_TOKEN_COUNTING_ACTION: Final = "countTokens" + + class RouteChecks: @staticmethod def should_call_route( @@ -412,6 +427,11 @@ class RouteChecks: detail=f"user not allowed to access this route. Route={route} is an admin only route", ) + @staticmethod + def is_unbilled_route(route: str) -> bool: + resource, _, action = route.rsplit("/", 1)[-1].partition(":") + return route in _UNBILLED_ROUTES or resource in _TOKEN_COUNTING_SEGMENTS or action == _TOKEN_COUNTING_ACTION + @staticmethod def is_llm_api_route(route: str) -> bool: """ diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index a1abed202a2..27fa4731baf 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -18,6 +18,7 @@ from typing import Final, NamedTuple, Protocol, Union, cast import fastapi from fastapi import HTTPException, Request, WebSocket, status from fastapi.security.api_key import APIKeyHeader +from pydantic import TypeAdapter from starlette.exceptions import WebSocketException import litellm @@ -268,6 +269,7 @@ _CLAUDE_MODEL_ROUTES: Final = frozenset( f"/{prefix}{endpoint}" for prefix in ("", "v1/") for endpoint in ("messages", "chat/completions", "responses") ) _CLAUDE_MODEL_NORMALIZED: Final = "litellm.claude_model_normalized" +_GENERAL_SETTINGS_ADAPTER: Final = TypeAdapter(dict[str, object]) async def _normalize_claude_model( @@ -1185,6 +1187,18 @@ async def user_api_key_auth_live_websocket(websocket: WebSocket) -> UserAPIKeyAu return await user_api_key_auth_websocket_for_model(websocket, model=None) +def _has_required_baggage_tag_mapping() -> bool: + from litellm.proxy import proxy_server + from litellm.proxy.common_utils.request_tag_mappings import REQUEST_TAG_MAPPINGS_ADAPTER + + raw_mappings: Final = _GENERAL_SETTINGS_ADAPTER.validate_python(proxy_server.general_settings).get( + "request_tag_mappings" + ) + return bool(raw_mappings) and any( + rule.required for rule in REQUEST_TAG_MAPPINGS_ADAPTER.validate_python(raw_mappings) + ) + + async def user_api_key_auth_websocket_for_model( websocket: WebSocket, model: str | None, *, close_on_auth_error: bool = True ) -> UserAPIKeyAuth: @@ -1240,7 +1254,11 @@ async def user_api_key_auth_websocket_for_model( # Call user_api_key_auth with the extracted API key # Note: You'll need to modify this to work with WebSocket context if needed try: - return await user_api_key_auth(request=request, api_key=f"Bearer {api_key}") + authorized: Final = await user_api_key_auth(request=request, api_key=f"Bearer {api_key}") + if _has_required_baggage_tag_mapping(): + raise HTTPException( + 403, "Required baggage tag mappings need HTTP inference, WebSocket inference is disabled" + ) except Exception as e: if not close_on_auth_error: raise @@ -1249,6 +1267,8 @@ async def user_api_key_auth_websocket_for_model( log_model_access_denial(e) await websocket.close(code=status.WS_1008_POLICY_VIOLATION) raise HTTPException(status_code=403, detail=str(e)) + else: + return authorized async def user_api_key_auth_accepted_websocket_for_model(websocket: WebSocket, model: str) -> UserAPIKeyAuth: @@ -3759,20 +3779,17 @@ async def run_centralized_common_checks( # later seed in common_checks pushes key tags and the # _tag_max_budget_check read into `litellm_metadata`, hiding header # tags from per-tag budget enforcement on LITELLM_METADATA_ROUTES. - LiteLLMProxyRequestSetup.pre_seed_litellm_metadata_for_route( - request_data=request_data, - route=route, - ) - # Merge x-litellm-tags into request_data BEFORE common_checks runs. # _tag_max_budget_check inside common_checks only inspects request_data; # without this pre-merge, header-supplied tags bypass tag-budget # enforcement. - LiteLLMProxyRequestSetup.apply_client_tag_policy_pre_auth( - request=request, - request_data=request_data, - user_api_key_dict=user_api_key_auth_obj, - ) + for tag_target in LiteLLMProxyRequestSetup.pre_auth_tag_merge_targets(request, request_data): + LiteLLMProxyRequestSetup.pre_seed_litellm_metadata_for_route(request_data=tag_target, route=route) + LiteLLMProxyRequestSetup.apply_client_tag_policy_pre_auth( + request=request, + request_data=tag_target, + user_api_key_dict=user_api_key_auth_obj, + ) bind_admission_counter_keys(user_api_key_auth_obj, end_user_id=end_user_id) try: @@ -4001,7 +4018,7 @@ def _resolve_request_principal(request: Request, valid_token: UserAPIKeyAuth) -> async def _authorize_authenticated_request( user_api_key_auth_obj: UserAPIKeyAuth, request: Request, - request_data: dict, + request_data: dict[str, object], route: str, api_key: str, *, @@ -4061,6 +4078,19 @@ async def _authorize_authenticated_request( billable=request_data.get("method") in (None, "message/send", "message/stream", "SendMessage", "SendStreamingMessage"), ) + if "request_tag_mappings" in general_settings and request.method == "POST": + from litellm.proxy.common_utils.request_tag_mappings import is_baggage_mapped_route + from litellm.proxy.proxy_server import proxy_logging_obj, user_api_key_cache + + if is_baggage_mapped_route(route): + LiteLLMProxyRequestSetup.pre_seed_litellm_metadata_for_route(request_data=authorized_data, route=route) + await LiteLLMProxyRequestSetup.apply_baggage_tag_mappings_pre_auth( + request=request, + request_data=authorized_data, + general_settings=_GENERAL_SETTINGS_ADAPTER.validate_python(general_settings), + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + ) await run_centralized_common_checks( user_api_key_auth_obj=user_api_key_auth_obj, request=request, diff --git a/litellm/proxy/common_utils/callback_utils.py b/litellm/proxy/common_utils/callback_utils.py index 4fa7709e5c7..46434a997b7 100644 --- a/litellm/proxy/common_utils/callback_utils.py +++ b/litellm/proxy/common_utils/callback_utils.py @@ -12,6 +12,7 @@ import litellm from litellm import get_secret from litellm._logging import verbose_proxy_logger from litellm.constants import ( + BAGGAGE_MAPPED_TAGS_METADATA_KEY, CLIENT_OUTPUT_CEILING_METADATA_KEY, CONSUMED_REQUEST_TAGS_METADATA_KEY, MAX_GUARDRAIL_SCAN_METADATA_HEADER_LENGTH, @@ -543,6 +544,7 @@ LITELLM_PROXY_INTERNAL_METADATA_KEYS: Final = frozenset( PRE_CALL_EXECUTED_GUARDRAILS_KEY, SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY, CONSUMED_REQUEST_TAGS_METADATA_KEY, + BAGGAGE_MAPPED_TAGS_METADATA_KEY, CLIENT_OUTPUT_CEILING_METADATA_KEY, ROUTING_REQUEST_TAGS_METADATA_KEY, "disable_global_guardrails", diff --git a/litellm/proxy/common_utils/http_parsing_utils.py b/litellm/proxy/common_utils/http_parsing_utils.py index ea33bc763d4..521a9a15bd0 100644 --- a/litellm/proxy/common_utils/http_parsing_utils.py +++ b/litellm/proxy/common_utils/http_parsing_utils.py @@ -372,6 +372,16 @@ def safe_set_request_parsed_body( _safe_set_request_parsed_body: Final = safe_set_request_parsed_body +def detach_request_parsed_body( + request: Request, request_data: Mapping[str, object] +) -> dict[str, object]: # mutable-ok: the client's body copy, which the base tag merges write into + client_body: Final = dict(request_data) + cached: Final[object] = request.scope.get("parsed_body") + if isinstance(cached, tuple) and cached[1] is request_data: + request.scope["parsed_body"] = (cached[0], client_body) # rebind-ok: cache holds the client copy + return client_body + + def rewrite_request_model( request_data: dict[str, object], # mutable-ok: the request body is rewritten in place for every downstream reader request: Request | None, diff --git a/litellm/proxy/common_utils/request_tag_mappings.py b/litellm/proxy/common_utils/request_tag_mappings.py new file mode 100644 index 00000000000..dca269dca16 --- /dev/null +++ b/litellm/proxy/common_utils/request_tag_mappings.py @@ -0,0 +1,67 @@ +from collections.abc import Mapping, Sequence +from typing import Final +from urllib.parse import unquote_plus + +from fastapi import status +from opentelemetry import baggage +from opentelemetry.baggage.propagation import W3CBaggagePropagator +from opentelemetry.context import Context +from pydantic import TypeAdapter + +from litellm.proxy._types import LiteLLMRoutes, ProxyErrorTypes, ProxyException +from litellm.proxy.auth.route_checks import RouteChecks +from litellm.types.proxy.request_tag_mappings import RequestTagMapping + +REQUEST_TAG_MAPPINGS_ADAPTER: Final = TypeAdapter(tuple[RequestTagMapping, ...]) + + +def is_baggage_mapped_route(route: str) -> bool: + from litellm.proxy.pass_through_endpoints.pass_through_endpoints import InitPassThroughEndpointHelpers + + return ( + (RouteChecks.is_llm_api_route(route) or InitPassThroughEndpointHelpers.is_registered_pass_through_route(route)) + and not RouteChecks.is_unbilled_route(route) + and not RouteChecks.check_route_access(route=route, allowed_routes=LiteLLMRoutes.mcp_inference_routes.value) + ) + + +def mapping_rejection(message: str) -> ProxyException: + return ProxyException( + message=message, type=ProxyErrorTypes.bad_request_error, param=None, code=status.HTTP_400_BAD_REQUEST + ) + + +def resolve_baggage_tag_mappings(header: str, rules: Sequence[RequestTagMapping]) -> tuple[tuple[str, bool], ...]: + """Exact mappings win over the any-value fallback for the same baggage key.""" + members: Final = tuple(_member_without_properties(member) for member in header.split(",")) + context: Final = W3CBaggagePropagator().extract({"baggage": ",".join(members)}, context=Context()) + values: Final = baggage.get_all(context) + member_keys: Final = tuple(unquote_plus(member.partition("=")[0]).strip() for member in members) + resolved: Final = ( + _resolve_key(key, values, member_keys, rules) for key in dict.fromkeys(rule.key for rule in rules) + ) + return tuple(mapping for mapping in resolved if mapping is not None) + + +def _member_without_properties(member: str) -> str: + key, separator, value = member.partition(";")[0].partition("=") + return f"{key.strip()}{separator}{value.strip()}" + + +def _resolve_key( + key: str, values: Mapping[str, object], member_keys: Sequence[str], rules: Sequence[RequestTagMapping] +) -> tuple[str, bool] | None: + key_rules: Final = tuple(rule for rule in rules if rule.key == key) + key_spellings: Final = frozenset((key, unquote_plus(key))) + sent_keys: Final = tuple(member_key for member_key in member_keys if member_key in key_spellings) + value: Final = values.get(sent_keys[0]) if len(sent_keys) == 1 else None + selected: Final = next( + (rule for rule in key_rules if rule.value == value and rule.value is not None), None + ) or next((rule for rule in key_rules if rule.value is None), None) + usable: Final = isinstance(value, str) and value.isprintable() and bool(value) + if not usable or selected is None: + if any(rule.required for rule in key_rules): + raise mapping_rejection(f"Missing, invalid or unmapped baggage value for {key}") + return None + tag: Final = selected.tag if selected.tag is not None else f"{selected.tag_prefix}{value}" + return tag, selected.require_budgeted_tag diff --git a/litellm/proxy/guardrails/auto_router_compression.py b/litellm/proxy/guardrails/auto_router_compression.py index ecb182ab33c..1267fb15db8 100644 --- a/litellm/proxy/guardrails/auto_router_compression.py +++ b/litellm/proxy/guardrails/auto_router_compression.py @@ -156,15 +156,13 @@ async def arm_pre_call( if not isinstance(model_alias, str) or not model_alias: return - from litellm.router_strategy.tag_based_routing import ( - get_tags_from_request_kwargs, - ) + from litellm.router_strategy.tag_based_routing import get_routing_tags_from_request_kwargs policy: Final = policy_for_model( llm_router=llm_router, model_alias=model_alias, request_kwargs=data, - request_tags=get_tags_from_request_kwargs(data), + request_tags=get_routing_tags_from_request_kwargs(data), ) if policy is None: return diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 67a1d640103..c0222dbcd0b 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -5,7 +5,8 @@ import json import re import time from collections import OrderedDict -from collections.abc import Mapping, MutableMapping, Sequence +from collections.abc import Iterable, Mapping, MutableMapping, Sequence +from dataclasses import dataclass from datetime import datetime from types import MappingProxyType from typing import ( @@ -17,7 +18,7 @@ from typing import ( cast, # noqa: TID251 # narrows untyped deployment and credential payloads ) -from fastapi import HTTPException, Request +from fastapi import HTTPException, Request, status from pydantic import TypeAdapter from pydantic import ValidationError as PydanticValidationError from starlette.datastructures import Headers @@ -27,6 +28,7 @@ from litellm._logging import verbose_logger, verbose_proxy_logger from litellm._service_logger import ServiceLogging from litellm._uuid import uuid from litellm.constants import ( + BAGGAGE_MAPPED_TAGS_METADATA_KEY, CALLER_SESSION_ID_METADATA_KEY, CLIENT_OUTPUT_CEILING_METADATA_KEY, CONSUMED_REQUEST_TAGS_METADATA_KEY, @@ -80,6 +82,7 @@ from litellm.proxy.common_utils.callback_utils import ( ) from litellm.proxy.common_utils.http_parsing_utils import ( # noqa: F401 # legacy module exports _safe_get_request_headers, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export + detach_request_parsed_body, safe_get_request_headers, ) from litellm.proxy.spend_tracking.carried_budget_state import carried_budget_metadata @@ -90,6 +93,7 @@ _SPECIAL_HEADERS_CACHE: Final = frozenset(str(v.value).lower() for v in SpecialH _REDACTED_HEADER_VALUE: Final = "***REDACTED***" _DECODED_JSON: Final = TypeAdapter(object) +_OBJECT_LIST_ADAPTER: Final = TypeAdapter(list[object]) _CREDENTIAL_HEADER_NAMES: Final = SpecialHeaders.litellm_credential_header_names() | frozenset( {"cookie", "proxy-authorization"} ) @@ -209,8 +213,19 @@ def _drop_client_caller_session(request_body: Mapping[str, object]) -> None: metadata.pop(CALLER_SESSION_ID_METADATA_KEY, None) +_PROXY_SET_METADATA_KEYS: Final = frozenset({CALLER_SESSION_ID_METADATA_KEY, BAGGAGE_MAPPED_TAGS_METADATA_KEY}) +_METADATA_TAGS_ADAPTER: Final = TypeAdapter(tuple[object, ...]) + + def _requester_metadata_snapshot(metadata: Mapping[str, object]) -> Mapping[str, object]: - return copy.deepcopy({key: value for key, value in metadata.items() if key != CALLER_SESSION_ID_METADATA_KEY}) + requester: Final = {key: value for key, value in metadata.items() if key not in _PROXY_SET_METADATA_KEYS} + stamp: Final = metadata.get(BAGGAGE_MAPPED_TAGS_METADATA_KEY) + tags: Final = requester.get("tags") + if not isinstance(stamp, BaggageMappedTagsStamp) or not isinstance(tags, list): + return copy.deepcopy(requester) + sent_tags: Final = [tag for tag in _METADATA_TAGS_ADAPTER.validate_python(tags) if tag not in stamp.tags] + without_tags: Final = {key: value for key, value in requester.items() if key != "tags"} + return copy.deepcopy({**without_tags, "tags": sent_tags} if sent_tags else without_tags) def _stampable_key_hash(user_api_key_dict: UserAPIKeyAuth) -> str | None: @@ -254,7 +269,7 @@ from litellm.router import Router from litellm.secret_managers.main import get_secret_bool from litellm.types.litellm_params import AGENTIC_LOOP_KWARG_NAMES from litellm.types.llms.anthropic import ANTHROPIC_API_HEADERS -from litellm.types.router import holds_api_key_secret_reference +from litellm.types.router import BaggageMappedTagsStamp, holds_api_key_secret_reference from litellm.types.services import ServiceTypes from litellm.types.utils import ( API_KEY_SECRET_REFERENCE_REJECTION, @@ -285,8 +300,11 @@ if TYPE_CHECKING: from opentelemetry.trace import Span as OtelSpan from litellm.integrations.otel.model.destination import OtelDestination + from litellm.models.tag import LiteLLM_TagTable + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.policy_engine.attachment_registry import AttachmentRegistry from litellm.proxy.proxy_server import ProxyConfig as _ProxyConfig + from litellm.proxy.utils import PrismaClient from litellm.types.proxy.policy_engine import Policy, PolicyMatchContext ProxyConfig = _ProxyConfig @@ -537,6 +555,39 @@ def _normalized_metadata_object(field: str, value: object) -> Mapping[str, objec raise _invalid_metadata_type_error(field=field, value=value) +def _serves_through_db_outage() -> bool: + from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler + + return PrismaDBExceptionHandler.should_allow_request_on_db_unavailable() + + +def _tag_budgets_unavailable(message: str) -> ProxyException: + return ProxyException( + message=message, type=ProxyErrorTypes.no_db_connection, param=None, code=status.HTTP_503_SERVICE_UNAVAILABLE + ) + + +async def _unbudgeted_mapped_tag( + budgeted_tags: tuple[str, ...], prisma_client: "PrismaClient | None", user_api_key_cache: "UserApiKeyCache" +) -> str | ProxyException | None: + from litellm.proxy.auth.auth_checks import read_tag_objects_batch + + if prisma_client is None: + return _tag_budgets_unavailable("Baggage tag mappings require a database to check budgeted tags") + tag_objects: Final = await read_tag_objects_batch( + tag_names=budgeted_tags, prisma_client=prisma_client, user_api_key_cache=user_api_key_cache + ) + if isinstance(tag_objects, Exception): + return _tag_budgets_unavailable("Baggage tag mappings could not read tag budgets from the database") + return next((tag for tag in budgeted_tags if not _has_tag_budget(tag_objects.get(tag))), None) + + +def _has_tag_budget(tag: "LiteLLM_TagTable | None") -> bool: + if tag is None or tag.litellm_budget_table is None: + return False + return tag.litellm_budget_table.max_budget is not None + + def _normalized_metadata_slot( request_data: MutableMapping[str, object], metadata_variable_name: str ) -> dict[str, object]: @@ -554,6 +605,43 @@ def _normalized_metadata_slot( return normalized +def _metadata_object_or_error(field: str, value: object) -> Mapping[str, object] | ProxyException: + if value is None: + return {} + try: + return dict(_normalized_metadata_object(field, value)) + except ProxyException as invalid_type: + return invalid_type + + +@dataclass(frozen=True, slots=True) +class _BaggageClientBody: + body: dict[str, object] # mutable-ok: the client's cached body, which the base tag merges write into + + +def _detach_auth_metadata_from_client_body( + request: Request, + request_data: MutableMapping[str, object], # mutable-ok: auth's own copy takes the normalized metadata objects +) -> None: + client_body: Final = detach_request_parsed_body(request, request_data) + request.state.litellm_baggage_client_body = _BaggageClientBody(client_body) # rebind-ok: mirror tag merges + slots: Final = tuple( + (field, _metadata_object_or_error(field, request_data[field])) + for field in ("metadata", "litellm_metadata") + if field in request_data + ) + for field, normalized in slots: + request_data[field] = {} if isinstance(normalized, ProxyException) else normalized # rebind-ok: auth's copy + + +def _metadata_string_tags(metadata_objects: Iterable[Mapping[str, object]]) -> tuple[str, ...]: + tag_lists: Final = (metadata.get("tags") for metadata in metadata_objects) + tags: Final = itertools.chain.from_iterable( + _OBJECT_LIST_ADAPTER.validate_python(tag_list) for tag_list in tag_lists if isinstance(tag_list, list) + ) + return tuple(tag for tag in tags if isinstance(tag, str)) + + def _strip_untrusted_request_header_controls( headers: Any, *, @@ -1132,6 +1220,8 @@ _get_dynamic_logging_metadata: Final = get_dynamic_logging_metadata _TENANT_OTEL_PARAMS: Final = TypeAdapter(StandardCallbackDynamicParams) +_STRING_TAGS_ADAPTER: Final = TypeAdapter(tuple[str, ...]) +_CLIENT_TAGS_ADAPTER: Final = TypeAdapter(tuple[object, ...]) def _tenant_otel_params(callback_vars: Mapping[str, str]) -> StandardCallbackDynamicParams: @@ -2045,7 +2135,7 @@ class LiteLLMProxyRequestSetup: @staticmethod def pre_seed_litellm_metadata_for_route( - request_data: dict, + request_data: MutableMapping[str, object], # mutable-ok: seed shared request metadata before admission route: str, ) -> None: """Pre-seed ``litellm_metadata`` for routes that track tags there. @@ -2064,6 +2154,87 @@ class LiteLLMProxyRequestSetup: if any(metadata_route in route for metadata_route in LITELLM_METADATA_ROUTES): request_data.setdefault("litellm_metadata", {}) + @staticmethod + def pre_auth_tag_merge_targets( + request: Request | None, + request_data: dict[str, object], # mutable-ok: auth's request body, which the tag merges write into + ) -> tuple[dict[str, object], ...]: # mutable-ok: every body the base tag merges write into + client_body: Final[object] = getattr(getattr(request, "state", None), "litellm_baggage_client_body", None) + if not isinstance(client_body, _BaggageClientBody) or client_body.body is request_data: + return (request_data,) + return (request_data, client_body.body) + + @staticmethod + async def apply_baggage_tag_mappings_pre_auth( + request: Request, + request_data: MutableMapping[str, object], # mutable-ok: enrich the shared request before budget admission + general_settings: Mapping[str, object], + prisma_client: "PrismaClient | None", + user_api_key_cache: "UserApiKeyCache", + ) -> None: + """Resolve configured baggage tags before the existing tag-budget gate.""" + from litellm.proxy.common_utils.request_tag_mappings import ( + REQUEST_TAG_MAPPINGS_ADAPTER, + mapping_rejection, + resolve_baggage_tag_mappings, + ) + + raw_rules: Final = general_settings.get("request_tag_mappings") + if not isinstance(raw_rules, list) or not raw_rules: + return + rules: Final = REQUEST_TAG_MAPPINGS_ADAPTER.validate_python(raw_rules) + header: Final = ",".join(request.headers.getlist("baggage")) + try: + mapped: Final = resolve_baggage_tag_mappings(header, rules) + except ProxyException: + _detach_auth_metadata_from_client_body(request, request_data) + raise + if not mapped: + return + client_metadata: Final[object] = request_data.get("metadata") + _detach_auth_metadata_from_client_body(request, request_data) + budgeted_tags: Final = tuple(tag for tag, require_budget in mapped if require_budget) + unbudgeted_tag: Final = ( + await _unbudgeted_mapped_tag(budgeted_tags, prisma_client, user_api_key_cache) if budgeted_tags else None + ) + if isinstance(unbudgeted_tag, ProxyException) and not _serves_through_db_outage(): + raise unbudgeted_tag + if isinstance(unbudgeted_tag, str): + raise mapping_rejection(f"Mapped tag {unbudgeted_tag} requires an existing tag budget") + client_has_tags: Final = (isinstance(client_metadata, dict) and "tags" in client_metadata) or ( + get_metadata_variable_name_from_kwargs(request_data) == "metadata" + and any(member.strip() for member in request.headers.get("x-litellm-tags", "").split(",")) + ) + mapped_tags: Final = tuple(tag for tag, _ in mapped) + request.state.litellm_baggage_client_metadata_tags = client_has_tags # rebind-ok: request input flag + request.state.litellm_baggage_mapped_tags = mapped_tags # rebind-ok: per-request immutable mapped tags + LiteLLMProxyRequestSetup.apply_authenticated_baggage_tags(request, request_data) + + @staticmethod + def apply_authenticated_baggage_tags( + request: Request, + request_data: MutableMapping[str, object], # mutable-ok: restore authenticated tags in downstream metadata + metadata_key: str | None = None, + ) -> None: + request_state: Final[object] = getattr(request, "state", None) + raw_tags: Final[object] = getattr(request_state, "litellm_baggage_mapped_tags", None) + if not isinstance(raw_tags, tuple) or not raw_tags: + return + tags: Final = _STRING_TAGS_ADAPTER.validate_python(raw_tags) + slot: Final = metadata_key if metadata_key is not None else get_metadata_variable_name_from_kwargs(request_data) + metadata: Final = _normalized_metadata_slot(request_data, slot) + existing: Final = metadata.get("tags") + if existing and not isinstance(existing, list): + return + existing_tags: Final = _CLIENT_TAGS_ADAPTER.validate_python(existing) if isinstance(existing, list) else () + prior_stamp: Final = metadata.get(BAGGAGE_MAPPED_TAGS_METADATA_KEY) + prior_mapped: Final = prior_stamp.tags if isinstance(prior_stamp, BaggageMappedTagsStamp) else () + sent_tags: Final = tuple(tag for tag in existing_tags if tag not in prior_mapped) + metadata["tags"] = [*existing_tags, *(tag for tag in tags if tag not in existing_tags)] + metadata[BAGGAGE_MAPPED_TAGS_METADATA_KEY] = BaggageMappedTagsStamp( + tags=tuple(tag for tag in tags if tag not in sent_tags) + ) + @staticmethod def apply_key_tags_pre_auth( request_data: dict, @@ -2221,12 +2392,16 @@ async def add_litellm_data_to_request( data.pop(_internal_key, None) _reject_url_valued_destinations(data) _raw_metadata_by_field: Final = { - _metadata_field: data.pop(_metadata_field) + _metadata_field: _DECODED_JSON.validate_python(data.pop(_metadata_field)) for _metadata_field in ("metadata", "litellm_metadata") if data.get(_metadata_field) is not None } - for _metadata_field, _raw_metadata in _raw_metadata_by_field.items(): - data[_metadata_field] = _normalized_metadata_object(_metadata_field, _raw_metadata) + _normalized_metadata_by_field: Final = { + _metadata_field: _normalized_metadata_object(_metadata_field, _raw_metadata) + for _metadata_field, _raw_metadata in _raw_metadata_by_field.items() + } + data.update(_normalized_metadata_by_field) + _client_sent_metadata_tags: Final = _metadata_string_tags(_normalized_metadata_by_field.values()) # Strip spoofable auth metadata from user-supplied metadata dict _user_metadata = data.get("metadata") if isinstance(_user_metadata, dict): @@ -2654,15 +2829,14 @@ async def add_litellm_data_to_request( tags_to_add=_caller_body_tags, ) - # caller_tags: exactly what this request itself supplied (x-litellm-tags header, - # body "tags", or body "metadata.tags" on litellm_metadata routes), never - # anything from key/team/project metadata. Read directly from the header and - # body values here, the same way inherited_tags above is read directly from - # key/team/project metadata -- neither is derived by inspecting the shared - # "tags" list, which a pre-auth pass (apply_client_tag_policy_pre_auth) may - # have already merged caller header tags into before this function runs. + # caller_tags: the tags the request's headers and body carry (x-litellm-tags, body "tags", and body + # "metadata.tags" / "litellm_metadata.tags" read before the mapped tags are added below), + # never derived from the shared "tags" list. data[_metadata_variable_name]["caller_tags"] = tuple( # rebind-ok: matches file idiom - dict.fromkeys((*(tags or ()), *(_caller_body_tags or ()))) + dict.fromkeys((*(tags or ()), *_client_sent_metadata_tags)) + ) + LiteLLMProxyRequestSetup.apply_authenticated_baggage_tags( + request=request, request_data=data, metadata_key=_metadata_variable_name ) # Team Callbacks controls @@ -3958,6 +4132,16 @@ def _apply_resolved_guardrails_to_metadata( verbose_proxy_logger.debug("Policy engine: added guardrails to request metadata: %s", combined) +def _policy_match_tags(request_data: Mapping[str, object], metadata_variable_name: str) -> list[str] | None: + from litellm.proxy.common_utils.http_parsing_utils import get_tags_from_request_body + from litellm.router_strategy.tag_based_routing import without_baggage_mapped_tags + + tags: Final = without_baggage_mapped_tags( + get_tags_from_request_body(request_data), request_data.get(metadata_variable_name) + ) + return list(tags) if tags else None + + async def add_guardrails_from_policy_engine( data: dict, metadata_variable_name: str, @@ -3982,7 +4166,6 @@ async def add_guardrails_from_policy_engine( user_api_key_dict: The user's API key authentication info """ from litellm._logging import verbose_proxy_logger - from litellm.proxy.common_utils.http_parsing_utils import get_tags_from_request_body from litellm.proxy.policy_engine.policy_registry import get_policy_registry from litellm.types.proxy.policy_engine import PolicyMatchContext @@ -4000,7 +4183,7 @@ async def add_guardrails_from_policy_engine( return # Extract tags and build context - all_tags: Final = get_tags_from_request_body(data) or None + all_tags: Final = _policy_match_tags(data, metadata_variable_name) _team_alias: Final = user_api_key_dict.team_alias _key_alias: Final = user_api_key_dict.key_alias context: Final = PolicyMatchContext( diff --git a/litellm/proxy/management_endpoints/tag_management_endpoints.py b/litellm/proxy/management_endpoints/tag_management_endpoints.py index 9a1cf701142..c12a90ed6a2 100644 --- a/litellm/proxy/management_endpoints/tag_management_endpoints.py +++ b/litellm/proxy/management_endpoints/tag_management_endpoints.py @@ -275,6 +275,7 @@ async def new_tag( """ from litellm.proxy._types import CommonProxyErrors from litellm.proxy.proxy_server import ( + invalidate_spend_counter, litellm_proxy_admin_name, llm_router, prisma_client, @@ -316,6 +317,7 @@ async def new_tag( ) await _evict_tag_cache_keys((tag_cache_key(tag.name), tag_registry_cache_key())) + await invalidate_spend_counter(counter_key=f"spend:tag:{tag.name}") # Update models with new tag if tag.models: diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index f2605a8ec6f..d9d3b6f8c92 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -222,7 +222,12 @@ async def _json_request_body(request: Request) -> Mapping[str, object]: return await request.json() -def get_passthrough_router_request_metadata(user_api_key_dict: UserAPIKeyAuth) -> Mapping[str, object]: +_METADATA_ADAPTER: Final = TypeAdapter(Mapping[str, object]) + + +def get_passthrough_router_request_metadata( + user_api_key_dict: UserAPIKeyAuth, request: Request | None = None +) -> Mapping[str, object]: """ Build the request metadata carrying key-level spend attribution and the pre-call budget reservation for a router-model passthrough request. @@ -243,13 +248,17 @@ def get_passthrough_router_request_metadata(user_api_key_dict: UserAPIKeyAuth) - """ from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup - request_data: Final = {"litellm_metadata": {}} + request_data: Final[dict[str, object]] = { # mutable-ok: legacy enrichment writes this wrapper + "litellm_metadata": {} + } LiteLLMProxyRequestSetup.add_user_api_key_auth_to_request_metadata( data=request_data, user_api_key_dict=user_api_key_dict, _metadata_variable_name="litellm_metadata", ) - return request_data["litellm_metadata"] + if request is not None: + LiteLLMProxyRequestSetup.apply_authenticated_baggage_tags(request, request_data) + return _METADATA_ADAPTER.validate_python(request_data["litellm_metadata"]) async def llm_passthrough_factory_proxy_route( @@ -549,7 +558,7 @@ async def vllm_proxy_route( params=None, headers=None, cookies=None, - litellm_metadata=get_passthrough_router_request_metadata(user_api_key_dict), + litellm_metadata=get_passthrough_router_request_metadata(user_api_key_dict, request), ), ) @@ -2229,7 +2238,7 @@ async def _relay_router_model( params=None, headers=None, cookies=None, - litellm_metadata=get_passthrough_router_request_metadata(user_api_key_dict), + litellm_metadata=get_passthrough_router_request_metadata(user_api_key_dict, request), ) except httpx.HTTPStatusError as upstream_error: return await _relay_upstream_response(upstream_error.response) diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 311dab058a4..fe7ac30e434 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -2021,6 +2021,9 @@ def _apply_key_team_project_controls( ) +_METADATA_ADAPTER: Final = TypeAdapter(dict[str, object]) + + def _update_metadata_with_tags_in_header(request: Request, metadata: dict) -> dict: """ If tags are in the request headers, add them to the metadata @@ -2046,7 +2049,12 @@ def _update_metadata_with_tags_in_header(request: Request, metadata: dict) -> di tags_to_add=tags_to_add, ) - return metadata + restored_metadata: Final = _METADATA_ADAPTER.validate_python(metadata) + request_data: Final[dict[str, object]] = { # mutable-ok: legacy enrichment writes this wrapper + "metadata": restored_metadata + } + LiteLLMProxyRequestSetup.apply_authenticated_baggage_tags(request, request_data) + return restored_metadata class _PassThroughRequestEnvelope(TypedDict, total=False): diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index ee8f20be243..c5128ddab87 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -7104,6 +7104,7 @@ class ProxyConfig: "vector_store_deny_by_default", "search_tool_deny_by_default", "vertex_ai_pass_through_drop_anthropic_betas", + "request_tag_mappings", ) if name in typed_general_settings } @@ -8348,6 +8349,7 @@ class ProxyConfig: previous_cleanup_schedule: tuple[object, ...], ) -> None: effects: Final = ( + self._apply_request_tag_mappings_setting, self._apply_alerting_settings, self._apply_pass_through_settings, self._apply_boolean_settings, @@ -8359,6 +8361,19 @@ class ProxyConfig: for effect in effects: await effect(db_values) + async def _apply_request_tag_mappings_setting(self, db_values: Mapping[str, SettingsJsonValue]) -> None: + if "request_tag_mappings" not in db_values: + return + mappings: Final = self.settings.get("request_tag_mappings") + try: + ConfigGeneralSettings.model_validate(MappingProxyType({"request_tag_mappings": mappings})) + except ValidationError as error: + self.settings.pop("request_tag_mappings", None) + verbose_proxy_logger.error( + "Ignoring invalid general_settings.request_tag_mappings from the DB, so no baggage maps to tags: %s", + error, + ) + async def _apply_alerting_settings(self, db_values: Mapping[str, SettingsJsonValue]) -> None: alerting: Final = self.settings.get("alerting") if "alerting" in db_values and isinstance(alerting, list): @@ -19640,6 +19655,7 @@ _PLUGIN_KEY_REDACTED: Final = "***" _GENERAL_SETTINGS_CONFIG_LIST_FIELD_TYPES: Final[Mapping[str, str]] = MappingProxyType( { + "request_tag_mappings": "RequestTagMappings", "max_parallel_requests": "Integer", "global_max_parallel_requests": "Integer", "max_in_flight_requests_per_worker": "Integer", diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index 4c9df95cd8d..803df9f83ca 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -228,29 +228,6 @@ def _raise_counter_budget_exceeded( ) -_UNBILLED_ROUTES: Final[frozenset[str]] = frozenset( - { - "/models", - "/v1/models", - "/utils/token_counter", - "/responses/input_tokens", - "/v1/responses/input_tokens", - "/openai/v1/responses/input_tokens", - } -) -_TOKEN_COUNTING_SEGMENTS: Final[frozenset[str]] = frozenset({"count_tokens", "count-tokens"}) -_TOKEN_COUNTING_ACTION: Final = "countTokens" - - -def _is_token_counting_route(route: str) -> bool: - resource, _, action = route.rsplit("/", 1)[-1].partition(":") - return resource in _TOKEN_COUNTING_SEGMENTS or action == _TOKEN_COUNTING_ACTION - - -def _is_unbilled_route(route: str) -> bool: - return route in _UNBILLED_ROUTES or _is_token_counting_route(route) - - async def reserve_budget_for_request( request_body: dict, route: str, @@ -269,7 +246,7 @@ async def reserve_budget_for_request( ) -> dict | None: if valid_token is None or not RouteChecks.is_llm_api_route(route=route): return None - if _is_unbilled_route(route): + if RouteChecks.is_unbilled_route(route): return None if get_model_from_request(request_body, route, llm_router=llm_router) is None: return None diff --git a/litellm/router.py b/litellm/router.py index ba3b4d53725..aadbdaa8520 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -62,6 +62,7 @@ from litellm.caching.caching import ( ) from litellm.caching.redis_cache import log_redis_failure from litellm.constants import ( + BAGGAGE_MAPPED_TAGS_METADATA_KEY, CLIENT_OUTPUT_CEILING_METADATA_KEY, CONSUMED_REQUEST_TAGS_METADATA_KEY, DEFAULT_AUTO_ROUTER_MAX_INPUT_CHARS, @@ -162,9 +163,12 @@ from litellm.router_strategy.lowest_tpm_rpm import LowestTPMLoggingHandler from litellm.router_strategy.lowest_tpm_rpm_v2 import LowestTPMLoggingHandler_v2, PrefetchedUsage from litellm.router_strategy.simple_shuffle import simple_shuffle from litellm.router_strategy.tag_based_routing import ( + baggage_stamp_for_serving_deployment, get_deployments_for_tag, + get_routing_tags_from_request_kwargs, get_tags_from_request_kwargs, is_valid_deployment_tag, + tags_with_serving_deployment_tags, ) from litellm.router_utils.access_windows import access_windows_config_error, filter_reserved_deployments from litellm.router_utils.add_retry_fallback_headers import ( @@ -4681,13 +4685,13 @@ class Router: ## DEPLOYMENT-LEVEL TAGS deployment_tags: Final = deployment.get("litellm_params", {}).get("tags") - if deployment_tags: - existing_tags = kwargs[metadata_variable_name].get("tags") or [] - merged_tags: Final = list(existing_tags) - for tag in deployment_tags: - if tag not in merged_tags: - merged_tags.append(tag) - kwargs[metadata_variable_name]["tags"] = merged_tags + merged_tags: Final = tags_with_serving_deployment_tags(kwargs[metadata_variable_name], deployment_tags) + if merged_tags is not None: + kwargs[metadata_variable_name]["tags"] = list(merged_tags) + + baggage_stamp: Final = baggage_stamp_for_serving_deployment(kwargs[metadata_variable_name], deployment_tags) + if baggage_stamp is not None: + kwargs[metadata_variable_name][BAGGAGE_MAPPED_TAGS_METADATA_KEY] = baggage_stamp ## CREDENTIAL NAME AS TAG credential_name: Final = deployment.get("litellm_params", {}).get("litellm_credential_name") @@ -15292,7 +15296,7 @@ class Router: if not candidates: return None - request_tags: Final = get_tags_from_request_kwargs(request_kwargs) + request_tags: Final = get_routing_tags_from_request_kwargs(request_kwargs) if request_tags: for tagged in candidates: if tagged.tags and is_valid_deployment_tag( @@ -15525,7 +15529,7 @@ class Router: llm_router=self, model_alias=registered_model_name, request_kwargs=request_kwargs, - request_tags=get_tags_from_request_kwargs(request_kwargs), + request_tags=get_routing_tags_from_request_kwargs(request_kwargs), ) # Shared compression already ran in the pre-call hook, so reuse it rather than # compressing twice. Conditional on arming having actually happened: only the @@ -15574,7 +15578,7 @@ class Router: value=self._consumed_request_tags_stamp( selected_strategy=selected_strategy, pre_routing_hook_response=pre_routing_hook_response, - request_tags=get_tags_from_request_kwargs(request_kwargs), + request_tags=get_routing_tags_from_request_kwargs(request_kwargs), ), ) diff --git a/litellm/router_strategy/tag_based_routing.py b/litellm/router_strategy/tag_based_routing.py index 73236294020..f13fd3fc630 100644 --- a/litellm/router_strategy/tag_based_routing.py +++ b/litellm/router_strategy/tag_based_routing.py @@ -10,13 +10,18 @@ Use this to route requests between Teams import re from collections.abc import Mapping, Sequence +from dataclasses import replace from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, overload from litellm._logging import verbose_logger -from litellm.constants import CONSUMED_REQUEST_TAGS_METADATA_KEY, ROUTING_REQUEST_TAGS_METADATA_KEY +from litellm.constants import ( + BAGGAGE_MAPPED_TAGS_METADATA_KEY, + CONSUMED_REQUEST_TAGS_METADATA_KEY, + ROUTING_REQUEST_TAGS_METADATA_KEY, +) from litellm.litellm_core_utils.core_helpers import get_metadata_variable_name_from_kwargs -from litellm.types.router import ConsumedRequestTagsStamp, DeploymentTypedDict, RouterErrors +from litellm.types.router import BaggageMappedTagsStamp, ConsumedRequestTagsStamp, DeploymentTypedDict, RouterErrors if TYPE_CHECKING: from litellm.router import Router as _Router @@ -484,6 +489,64 @@ def _request_tags_after_router_consumption(metadata: object, model: str) -> Sequ return tuple(dict.fromkeys((*leftover, *(tag for tag in typed_inherited_tags if isinstance(tag, str))))) +def without_baggage_mapped_tags( + request_tags: Sequence[str] | None, metadata: object, *, keep_serving_deployment_tags: bool = False +) -> Sequence[str] | None: + """Drop baggage-mapped tags that neither key, team or project policy nor the caller itself asked for.""" + if request_tags is None or not isinstance(metadata, Mapping): + return request_tags + typed_metadata: Final[Mapping[str, object]] = metadata + stamp: Final = typed_metadata.get(BAGGAGE_MAPPED_TAGS_METADATA_KEY) + if not isinstance(stamp, BaggageMappedTagsStamp): + return request_tags + asked_for_tags: Final = ( + *_metadata_tags(typed_metadata, "inherited_tags"), + *_metadata_tags(typed_metadata, "caller_tags"), + *(stamp.serving_deployment_tags if keep_serving_deployment_tags else ()), + ) + return tuple(tag for tag in request_tags if tag not in stamp.tags or tag in asked_for_tags) or None + + +def baggage_stamp_for_serving_deployment(metadata: object, deployment_tags: object) -> BaggageMappedTagsStamp | None: + if not isinstance(metadata, Mapping): + return None + typed_metadata: Final[Mapping[str, object]] = metadata + stamp: Final = typed_metadata.get(BAGGAGE_MAPPED_TAGS_METADATA_KEY) + if not isinstance(stamp, BaggageMappedTagsStamp): + return None + return replace( + stamp, serving_deployment_tags=tuple(tag for tag in _tag_sequence(deployment_tags) if isinstance(tag, str)) + ) + + +def tags_with_serving_deployment_tags(metadata: object, deployment_tags: object) -> tuple[object, ...] | None: + """Merge the serving deployment's tags ahead of baggage-mapped tags, so baggage never changes tag precedence.""" + serving_tags: Final = _tag_sequence(deployment_tags) + if not isinstance(metadata, Mapping) or not serving_tags: + return None + typed_metadata: Final[Mapping[str, object]] = metadata + existing_tags: Final = _metadata_tags(typed_metadata, "tags") + stamp: Final = typed_metadata.get(BAGGAGE_MAPPED_TAGS_METADATA_KEY) + mapped_tags: Final = stamp.tags if isinstance(stamp, BaggageMappedTagsStamp) else () + kept_tags: Final = tuple(tag for tag in existing_tags if tag not in mapped_tags) + added_serving_tags: Final = tuple( + tag for index, tag in enumerate(serving_tags) if tag not in kept_tags and tag not in serving_tags[:index] + ) + readded_mapped_tags: Final = tuple( + dict.fromkeys(tag for tag in existing_tags if tag in mapped_tags and tag not in added_serving_tags) + ) + return (*kept_tags, *added_serving_tags, *readded_mapped_tags) + + +def _metadata_tags(metadata: Mapping[str, object], key: str) -> Sequence[object]: + return _tag_sequence(metadata.get(key)) + + +def _tag_sequence(tags: object) -> Sequence[object]: + typed_tags: Final[Sequence[object]] = tags if isinstance(tags, (list, tuple)) else () + return typed_tags + + async def get_deployments_for_tag( llm_router_instance: LitellmRouter, model: str, # used to raise the correct error @@ -524,7 +587,9 @@ async def get_deployments_for_tag( verbose_logger.debug("request metadata: %s", request_kwargs.get(metadata_variable_name)) if metadata_variable_name in request_kwargs: metadata: Final = request_kwargs[metadata_variable_name] - request_tags: Final = _request_tags_after_router_consumption(metadata, model) + request_tags: Final = without_baggage_mapped_tags( + _request_tags_after_router_consumption(metadata, model), metadata + ) match_any: Final = llm_router_instance.tag_filtering_match_any routing_prefix: Final = llm_router_instance.tag_routing_prefix or "" @@ -707,3 +772,21 @@ def get_tags_from_request_kwargs( _get_tags_from_request_kwargs = get_tags_from_request_kwargs + + +def _request_metadata(request_kwargs: Mapping[str, object]) -> object: + metadata_variable_name: Final = get_metadata_variable_name_from_kwargs(request_kwargs) + if metadata_variable_name in request_kwargs: + return request_kwargs[metadata_variable_name] + litellm_params: Final = request_kwargs.get("litellm_params") + if not isinstance(litellm_params, Mapping): + return None + typed_litellm_params: Final[Mapping[str, object]] = litellm_params + return typed_litellm_params.get(metadata_variable_name) + + +def get_routing_tags_from_request_kwargs(request_kwargs: Mapping[str, object]) -> Sequence[str]: + return ( + without_baggage_mapped_tags(get_tags_from_request_kwargs(request_kwargs), _request_metadata(request_kwargs)) + or () + ) diff --git a/litellm/types/proxy/request_tag_mappings.py b/litellm/types/proxy/request_tag_mappings.py new file mode 100644 index 00000000000..17cb4f4048d --- /dev/null +++ b/litellm/types/proxy/request_tag_mappings.py @@ -0,0 +1,27 @@ +from typing import Literal + +from pydantic import BaseModel, ConfigDict, Field, model_validator +from typing_extensions import Self + + +class RequestTagMapping(BaseModel): + """An exact baggage value mapping, or a prefix for any value of a key.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + + header: Literal["baggage"] = "baggage" + key: str = Field(min_length=1, pattern=r"^[!#$%&'*+\-.^_`|~0-9A-Za-z]+$") + value: str | None = Field(None, min_length=1) + tag: str | None = Field(None, min_length=1) + tag_prefix: str | None = Field(None, min_length=1) + required: bool = False + require_budgeted_tag: bool = False + + @model_validator(mode="after") + def validate_mapping(self) -> Self: + if self.value is not None: + if self.tag is None or self.tag_prefix is not None: + raise ValueError("An exact value mapping requires tag and cannot set tag_prefix") + elif self.tag_prefix is None or self.tag is not None: + raise ValueError("An any-value mapping requires tag_prefix and cannot set tag") + return self diff --git a/litellm/types/router.py b/litellm/types/router.py index 47dafcf275d..6cd758d864e 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -1332,6 +1332,18 @@ class ConsumedRequestTagsStamp: tags: tuple[str, ...] +@dataclass(frozen=True, slots=True) +class BaggageMappedTagsStamp: + """Request tags added by baggage mappings, which classify spend and budgets but never select deployments. + + The serving deployment's own tags are kept beside them, so a mapped tag that the deployment also carries + still picks that deployment's guardrail mode. + """ + + tags: tuple[str, ...] + serving_deployment_tags: tuple[str, ...] = () + + @runtime_checkable class PreRoutingStrategy(Protocol): """Structural interface shared by the auto / complexity / adaptive / quality routers.""" diff --git a/tests/integration/spend/test_baggage_tag_mappings.py b/tests/integration/spend/test_baggage_tag_mappings.py new file mode 100644 index 00000000000..1bf90269484 --- /dev/null +++ b/tests/integration/spend/test_baggage_tag_mappings.py @@ -0,0 +1,847 @@ +import itertools +import json +import uuid +from collections.abc import Generator, Iterator, Mapping, Sequence +from contextlib import ExitStack, contextmanager +from dataclasses import dataclass +from hashlib import sha256 +from pathlib import Path +from types import MappingProxyType +from typing import Final, Literal + +import httpx +import pytest +import yaml +from integration._support.client import ( + JSON_OBJECT, + Gateway, + eventually, + gateway_from_environment, + list_value, + object_value, +) +from integration._support.database import scratch_database +from integration._support.database_relay import dropped_connection_relay, plaintext_database_url +from integration._support.process import graceful_stop_seconds, owned_proxy +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import JsonValue, TypeAdapter +from websockets.exceptions import InvalidStatus +from websockets.sync.client import connect + +pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) + +PROXY_CONFIG: Final = Path(__file__).resolve().parents[1] / "proxy_config.yaml" +PRICED: Final[dict[str, JsonValue]] = {"input_cost_per_token": 0.01, "output_cost_per_token": 0.01} +WORKER_PROBES: Final = 6 +JSON_LIST: Final = TypeAdapter(list[JsonValue]) +GUARDRAIL_VERDICT: Final = b'{"action":"GUARDRAIL_INTERVENED","texts":["masked baggage prompt"]}' +GUARDRAIL_PATH: Final = "/beta/litellm_basic_guardrail_api" +NO_FIELDS: Final[Mapping[str, JsonValue]] = MappingProxyType({}) + +Route = Literal["chat", "chat_litellm_metadata", "messages", "responses"] +CallerTag = Literal["none", "header", "metadata_json_with_header", "litellm_metadata"] +RejectedClientTag = Literal["none", "header", "body", "metadata_json", "litellm_metadata"] + + +@dataclass(frozen=True, slots=True) +class Rig: + proxy: Gateway + provider: Wire + callbacks: Wire + marker: str + + @property + def chat_model(self) -> str: + return f"chat-{self.marker}" + + @property + def exact_tag(self) -> str: + return f"usecase-a-{self.marker}" + + def usecase_tag(self, value: str) -> str: + return f"usecase-{self.marker}:{value}" + + def tenant_tag(self, value: str) -> str: + return f"tenant-{self.marker}:{value}" + + def budgeted_tag(self, value: str) -> str: + return f"budgeted-{self.marker}:{value}" + + def team_tag(self, value: str) -> str: + return f"team-{self.marker}:{value}" + + @property + def guarded_tag(self) -> str: + return f"guarded-{self.marker}:on" + + @property + def policy_tag(self) -> str: + return f"policy-{self.marker}:sensitive" + + +def _chat_reply() -> Reply: + marker: Final = uuid.uuid4().hex + return Reply( + body=json.dumps( + { + "id": f"chatcmpl-{marker}", + "object": "chat.completion", + "created": 1, + "model": "gpt-4o-mini", + "choices": [{"index": 0, "message": {"role": "assistant", "content": marker}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 5, "completion_tokens": 3, "total_tokens": 8}, + } + ).encode() + ) + + +def _responses_reply() -> Reply: + marker: Final = uuid.uuid4().hex + return Reply( + body=json.dumps( + { + "id": f"resp_{marker}", + "object": "response", + "created_at": 1, + "status": "completed", + "model": "gpt-4o-mini", + "output": [ + { + "id": f"msg_{marker}", + "type": "message", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": marker, "annotations": []}], + } + ], + "usage": {"input_tokens": 5, "output_tokens": 3, "total_tokens": 8}, + } + ).encode() + ) + + +def _provider(request: Request) -> Reply: + if request.target == GUARDRAIL_PATH: + return Reply(body=b'{"action":"NONE"}') + if request.target == "/v1/responses": + return _responses_reply() + return _chat_reply() + + +def _deployment(model_name: str, provider: Wire, **parameters: JsonValue) -> dict[str, JsonValue]: + return { + "model_name": model_name, + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_base": f"{provider.url}/v1", + "api_key": "synthetic-baggage-key", + **PRICED, + **parameters, + }, + } + + +def _write_config( + directory: Path, + model_list: list[JsonValue], + general_settings: Mapping[str, JsonValue], + litellm_settings: Mapping[str, JsonValue], + guardrails: Sequence[JsonValue] = (), +) -> Path: + base: Final = JSON_OBJECT.validate_python(yaml.safe_load(PROXY_CONFIG.read_text())) + config: Final = { + **base, + "model_list": model_list, + "general_settings": {**object_value(base["general_settings"]), **general_settings}, + "litellm_settings": {**object_value(base["litellm_settings"]), **litellm_settings}, + **({"guardrails": list(guardrails)} if guardrails else {}), + } + path: Final = directory / "baggage-tag-mappings.yaml" + path.write_text(yaml.safe_dump(config)) + return path + + +def _routed_deployment(marker: str, provider: Wire, name: str, tag: str) -> dict[str, JsonValue]: + return { + **_deployment(f"routed-{marker}", provider, tags=[tag]), + "model_info": {"id": f"{name}-{marker}", "enable_tag_filtering": True}, + } + + +def _mapped_models(marker: str, provider: Wire, upstream_url: str) -> list[JsonValue]: + routed: Final = ( + _routed_deployment(marker, provider, "tagged", f"usecase-a-{marker}"), + _routed_deployment(marker, provider, "default", "default"), + ) + realtime: Final[dict[str, JsonValue]] = { + "model_name": f"realtime-{marker}", + "litellm_params": { + "model": "openai/gpt-realtime-2.1", + "api_base": f"{upstream_url}/v1", + "api_key": f"realtime-{marker}", + }, + } + guarded: Final = _deployment(f"guarded-{marker}", provider, tags=[f"guarded-{marker}:on"]) + pre_guarded: Final = _deployment(f"guarded-pre-{marker}", provider, tags=[f"guarded-{marker}:pre"]) + return [_deployment(f"chat-{marker}", provider), guarded, pre_guarded, *routed, realtime] + + +def _mapped_rules(marker: str, *, tenant_required: bool) -> list[JsonValue]: + return [ + {"key": "entrypoint", "value": "a", "tag": f"usecase-a-{marker}"}, + {"key": "entrypoint", "tag_prefix": f"usecase-{marker}:"}, + {"key": "tenant", "tag_prefix": f"tenant-{marker}:", "required": tenant_required}, + {"key": "budgeted", "tag_prefix": f"budgeted-{marker}:", "require_budgeted_tag": True}, + {"key": "team+id", "tag_prefix": f"team-{marker}:"}, + {"key": "guard", "tag_prefix": f"guarded-{marker}:"}, + ] + + +def _deployment_tag_mode_guardrail(marker: str, provider: Wire) -> dict[str, JsonValue]: + return { + "guardrail_name": f"deployment-tag-mode-{marker}", + "litellm_params": { + "guardrail": "generic_guardrail_api", + "mode": { + "tags": { + f"guarded-{marker}:on": "post_call", + f"guarded-{marker}:pre": "pre_call", + f"policy-{marker}:sensitive": "post_call", + } + }, + "default_on": True, + "api_base": f"{provider.url}{GUARDRAIL_PATH}", + "api_key": "synthetic-guardrail-key", + }, + } + + +@contextmanager +def _rig( + tmp_path_factory: pytest.TempPathFactory, + general_settings: Mapping[str, JsonValue], + *, + tenant_required: bool, + workers: int, + tag_mode_guardrail: bool = False, +) -> Generator[Rig]: + marker: Final = uuid.uuid4().hex[:12] + directory: Final = tmp_path_factory.mktemp("baggage-tag-mappings") + with ( + gateway_from_environment() as base, + wire_server(_provider) as provider, + wire_server(lambda _request: Reply()) as callbacks, + ): + config: Final = _write_config( + directory, + _mapped_models(marker, provider, base.upstream_url), + {"request_tag_mappings": _mapped_rules(marker, tenant_required=tenant_required), **general_settings}, + {"callbacks": ["generic_api"], "DEFAULT_FLUSH_INTERVAL_SECONDS": 1}, + [_deployment_tag_mode_guardrail(marker, provider)] if tag_mode_guardrail else (), + ) + with owned_proxy( + base, directory, {"GENERIC_LOGGER_ENDPOINT": callbacks.url}, config=config, workers=workers + ) as proxy: + yield Rig(proxy, provider, callbacks, marker) + + +@pytest.fixture(scope="module") +def mapped(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]: + with _rig(tmp_path_factory, {}, tenant_required=True, workers=2) as rig: + yield rig + + +@pytest.fixture(scope="module") +def optional(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]: + with _rig(tmp_path_factory, {}, tenant_required=False, workers=1, tag_mode_guardrail=True) as rig: + yield rig + + +@pytest.fixture(scope="module") +def rejecting(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]: + with _rig(tmp_path_factory, {"reject_clientside_metadata_tags": True}, tenant_required=False, workers=1) as rig: + yield rig + + +def _digest(key: str) -> str: + return sha256(key.encode()).hexdigest() + + +def _caller_tags(raw: JsonValue) -> frozenset[str]: + tags: Final = JSON_LIST.validate_json(raw) if isinstance(raw, str) else list_value(raw) + return frozenset(tag for tag in tags if isinstance(tag, str) and not tag.startswith("User-Agent: ")) + + +def _spend_tags(proxy: Gateway, key: str) -> frozenset[str]: + rows: Final = eventually( + lambda: proxy.rows('SELECT request_tags FROM "LiteLLM_SpendLogs" WHERE api_key=%s', (_digest(key),)), + lambda values: len(values) == 1, + seconds=70, + ) + return _caller_tags(rows[0]["request_tags"]) + + +def _batch_events(batch: Request) -> Iterator[dict[str, JsonValue]]: + return (object_value(item) for item in JSON_LIST.validate_json(batch.body)) + + +def _logged_event(callbacks: Wire, call_id: str) -> dict[str, JsonValue]: + def delivered() -> tuple[dict[str, JsonValue], ...]: + events: Final = itertools.chain.from_iterable(_batch_events(batch) for batch in callbacks.drain()) + return tuple(event for event in events if event.get("litellm_call_id") == call_id) + + events: Final = eventually(delivered, lambda values: len(values) == 1, seconds=30) + return events[0] + + +def _inference_body(route: Route, model: str, text: str, caller_tag: str) -> dict[str, JsonValue]: + match route: + case "chat": + return {"model": model, "messages": [{"role": "user", "content": text}], "metadata": {"tags": [caller_tag]}} + case "chat_litellm_metadata": + return { + "model": model, + "messages": [{"role": "user", "content": text}], + "litellm_metadata": {"tags": [caller_tag]}, + } + case "messages": + return { + "model": model, + "max_tokens": 16, + "messages": [{"role": "user", "content": text}], + "litellm_metadata": {"tags": [caller_tag]}, + } + case "responses": + return {"model": model, "input": text, "litellm_metadata": {"tags": [caller_tag]}} + + +ROUTE_PATHS: Final[Mapping[Route, str]] = { + "chat": "/v1/chat/completions", + "chat_litellm_metadata": "/v1/chat/completions", + "messages": "/v1/messages", + "responses": "/v1/responses", +} + + +def _guardrail_inputs(rig: Rig) -> tuple[JsonValue, ...]: + return tuple( + JSON_OBJECT.validate_json(request.body)["input_type"] + for request in rig.provider.drain() + if request.target == GUARDRAIL_PATH + ) + + +def _chat( + rig: Rig, + key: str, + baggage: str | None, + *, + model: str | None = None, + fields: Mapping[str, JsonValue] = NO_FIELDS, + **headers: str, +) -> httpx.Response: + return rig.proxy.request( + "POST", + "/v1/chat/completions", + { + "model": model or rig.chat_model, + "messages": [{"role": "user", "content": f"baggage {uuid.uuid4().hex}"}], + "max_tokens": 16, + **fields, + }, + key=key, + headers={**({"baggage": baggage} if baggage is not None else {}), **headers}, + ) + + +@pytest.mark.parametrize("route", ["chat", "chat_litellm_metadata", "messages", "responses"]) +def test_mapped_baggage_tags_reach_spend_and_callbacks_but_not_requester_metadata(mapped: Rig, route: Route) -> None: + tenant: Final = uuid.uuid4().hex + caller_tag: Final = f"caller-{tenant}" + mapped_tags: Final = frozenset({mapped.exact_tag, mapped.tenant_tag(tenant)}) + with mapped.proxy.scenario() as scenario: + key: Final = scenario.key() + mapped.provider.drain() + response: Final = mapped.proxy.request( + "POST", + ROUTE_PATHS[route], + _inference_body(route, mapped.chat_model, f"baggage {tenant}", caller_tag), + key=key, + headers={"baggage": f"entrypoint=a, tenant={tenant}"}, + ) + assert response.status_code == 200, response.text + assert len(mapped.provider.drain()) == 1 + assert _spend_tags(mapped.proxy, key) == mapped_tags | {caller_tag} + event: Final = _logged_event(mapped.callbacks, response.headers["x-litellm-call-id"]) + assert mapped_tags <= _caller_tags(event["request_tags"]), event["request_tags"] + requester: Final = object_value(object_value(event["metadata"]).get("requester_metadata") or {}) + assert not mapped_tags & _caller_tags(requester.get("tags") or []), requester + + +def test_json_string_metadata_gains_exactly_the_mapped_tags(optional: Rig) -> None: + tenant: Final = uuid.uuid4().hex + mapped_tags: Final = frozenset({optional.exact_tag, optional.tenant_tag(tenant)}) + fields: Final[Mapping[str, JsonValue]] = {"metadata": json.dumps({"tags": [f"caller-{tenant}"]})} + with optional.proxy.scenario() as scenario: + plain_key: Final = scenario.key() + key: Final = scenario.key() + optional.provider.drain() + plain: Final = _chat(optional, plain_key, None, fields=fields) + response: Final = _chat(optional, key, f"entrypoint=a,tenant={tenant}", fields=fields) + assert (plain.status_code, response.status_code) == (200, 200), (plain.text, response.text) + assert len(optional.provider.drain()) == 2 + assert _spend_tags(optional.proxy, key) == _spend_tags(optional.proxy, plain_key) | mapped_tags + event: Final = _logged_event(optional.callbacks, response.headers["x-litellm-call-id"]) + assert mapped_tags <= _caller_tags(event["request_tags"]), event["request_tags"] + requester: Final = object_value(object_value(event["metadata"]).get("requester_metadata") or {}) + assert not mapped_tags & _caller_tags(requester.get("tags") or []), requester + + +def test_a2a_spend_row_carries_the_mapped_tags(mapped: Rig) -> None: + marker: Final = "baggagea2a" + uuid.uuid4().hex[:12] + + def card(url: str) -> dict[str, JsonValue]: + return { + "protocolVersion": "0.3", + "name": marker, + "description": "Synthetic peer", + "version": "1.0.0", + "url": url + "/", + "capabilities": {"streaming": True}, + "defaultInputModes": ["text"], + "defaultOutputModes": ["text"], + "skills": [], + } + + def peer(request: Request) -> Reply: + if request.method == "GET": + return Reply(body=json.dumps(card(wire.url)).encode()) + rpc: Final = JSON_OBJECT.validate_json(request.body) + return Reply( + body=json.dumps( + { + "jsonrpc": "2.0", + "id": rpc["id"], + "result": { + "kind": "message", + "role": "agent", + "messageId": f"{marker}-out", + "parts": [{"kind": "text", "text": "done"}], + }, + } + ).encode() + ) + + tenant: Final = uuid.uuid4().hex + with wire_server(peer) as wire, mapped.proxy.scenario() as scenario: + created: Final = mapped.proxy.post( + "/v1/agents", + {"agent_name": marker, "agent_card_params": card(wire.url), "litellm_params": {"cost_per_query": 0.25}}, + ) + agent: Final = str(created["agent_id"]) + scenario.cleanups.callback(lambda: mapped.proxy.request("DELETE", f"/v1/agents/{agent}")) + key: Final = scenario.key() + sent: Final = mapped.proxy.request( + "POST", + f"/a2a/{agent}", + { + "jsonrpc": "2.0", + "id": f"{marker}-send", + "method": "message/send", + "params": { + "message": { + "kind": "message", + "role": "user", + "messageId": f"{marker}-in", + "parts": [{"kind": "text", "text": "ping"}], + } + }, + }, + key=key, + headers={"baggage": f"entrypoint=a,tenant={tenant}"}, + ) + assert sent.status_code == 200, sent.text + assert {mapped.exact_tag, mapped.tenant_tag(tenant)} <= _spend_tags(mapped.proxy, key) + + +@pytest.mark.parametrize( + ("key_fields", "fields", "headers"), + [ + ({}, NO_FIELDS, {}), + ({}, {"metadata": {"tags": "run-123"}}, {"x-litellm-tags": "caller-tag"}), + ({"metadata": {"tags": ["key-tag"]}}, {"metadata": {"tags": "run-123"}}, {}), + ({}, {"metadata": {"tags": ""}}, {}), + ], + ids=["plain", "string_tags_with_header_tag", "string_tags_with_key_tag", "empty_string_tags"], +) +def test_spent_mapped_tag_budget_rejects_the_next_request_before_the_provider( + mapped: Rig, key_fields: Mapping[str, JsonValue], fields: Mapping[str, JsonValue], headers: Mapping[str, str] +) -> None: + entry: Final = uuid.uuid4().hex + tag: Final = mapped.usecase_tag(entry) + baggage: Final = f"entrypoint={entry},tenant={entry}" + with mapped.proxy.scenario() as scenario: + key: Final = scenario.key(**key_fields) + mapped.proxy.post("/tag/new", {"name": tag, "max_budget": 0.0001}) + scenario.cleanups.callback(lambda: mapped.proxy.request("POST", "/tag/delete", {"name": tag})) + first: Final = _chat(mapped, key, baggage) + assert first.status_code == 200, first.text + eventually(lambda: _chat(mapped, key, baggage).status_code, lambda status: status != 200, seconds=70) + mapped.provider.drain() + blocked: Final = _chat(mapped, key, baggage, fields=fields, **headers) + assert blocked.status_code == 422, blocked.text + assert "budget" in blocked.text.lower(), blocked.text + assert mapped.provider.drain() == () + other: Final = _chat(mapped, key, f"entrypoint=other-{entry},tenant={entry}") + assert other.status_code == 200, other.text + + +def test_required_mapping_rejects_a_request_without_its_baggage_value_before_the_provider(mapped: Rig) -> None: + with mapped.proxy.scenario() as scenario: + key: Final = scenario.key() + mapped.provider.drain() + rejected: Final = _chat(mapped, key, "entrypoint=a") + assert rejected.status_code == 400, rejected.text + assert mapped.provider.drain() == () + admitted: Final = _chat(mapped, key, f"entrypoint=a,tenant={uuid.uuid4().hex}") + assert admitted.status_code == 200, admitted.text + + +def test_require_budgeted_tag_rejects_an_unbudgeted_mapped_tag_until_the_tag_has_a_budget(mapped: Rig) -> None: + value: Final = uuid.uuid4().hex + tag: Final = mapped.budgeted_tag(value) + baggage: Final = f"budgeted={value},tenant={value}" + with mapped.proxy.scenario() as scenario: + key: Final = scenario.key() + mapped.provider.drain() + unbudgeted: Final = tuple(_chat(mapped, key, baggage, Connection="close") for _ in range(WORKER_PROBES)) + assert [response.status_code for response in unbudgeted] == [400] * WORKER_PROBES, [ + response.text for response in unbudgeted + ] + assert mapped.provider.drain() == () + mapped.proxy.post("/tag/new", {"name": tag, "max_budget": 100}) + scenario.cleanups.callback(lambda: mapped.proxy.request("POST", "/tag/delete", {"name": tag})) + budgeted: Final = tuple(_chat(mapped, key, baggage, Connection="close") for _ in range(WORKER_PROBES)) + assert [response.status_code for response in budgeted] == [200] * WORKER_PROBES, [ + response.text for response in budgeted + ] + + +def _routing_caller_tag(rig: Rig, caller_tag: CallerTag) -> tuple[Mapping[str, JsonValue], Mapping[str, str]]: + json_metadata: Final[Mapping[str, JsonValue]] = {"metadata": json.dumps({"tags": [rig.exact_tag]})} + match caller_tag: + case "none": + return NO_FIELDS, {} + case "header": + return NO_FIELDS, {"x-litellm-tags": rig.exact_tag} + case "metadata_json_with_header": + return json_metadata, {"x-litellm-tags": f"other-{rig.marker}"} + case "litellm_metadata": + return {"litellm_metadata": {"tags": [rig.exact_tag]}}, {} + + +@pytest.mark.parametrize( + "caller_tag", + ["none", "header", "metadata_json_with_header", "litellm_metadata"], + ids=["baggage_only", "caller_tag", "metadata_json_tag_with_header", "litellm_metadata_tag"], +) +def test_mapped_tags_never_select_a_tag_routed_deployment(mapped: Rig, caller_tag: CallerTag) -> None: + tenant: Final = uuid.uuid4().hex + expected: Final = f"default-{mapped.marker}" if caller_tag == "none" else f"tagged-{mapped.marker}" + fields, headers = _routing_caller_tag(mapped, caller_tag) + with mapped.proxy.scenario() as scenario: + key: Final = scenario.key() + responses: Final = tuple( + _chat( + mapped, + key, + f"entrypoint=a,tenant={tenant}", + model=f"routed-{mapped.marker}", + fields=fields, + **headers, + ) + for _ in range(3) + ) + assert [response.status_code for response in responses] == [200] * 3, [r.text for r in responses] + assert [response.headers.get("x-litellm-model-id") for response in responses] == [expected] * 3 + + +def test_mapped_tags_never_trigger_a_tag_scoped_policy_attachment(mapped: Rig) -> None: + entry: Final = uuid.uuid4().hex + tag: Final = mapped.usecase_tag(entry) + baggage: Final = f"entrypoint={entry},tenant={entry}" + proxy: Final = mapped.proxy + with wire_server(lambda _request: Reply(body=GUARDRAIL_VERDICT)) as guardrail, proxy.scenario() as scenario: + key: Final = scenario.key() + guardrail_created: Final = proxy.post( + "/guardrails", + { + "guardrail": { + "guardrail_name": f"guardrail-{entry}", + "litellm_params": { + "guardrail": "generic_guardrail_api", + "mode": "pre_call", + "default_on": False, + "api_base": f"{guardrail.url}/beta/litellm_basic_guardrail_api", + "api_key": "synthetic-guardrail-key", + }, + } + }, + ) + scenario.cleanups.callback(lambda: proxy.request("DELETE", f"/guardrails/{guardrail_created['guardrail_id']}")) + policy: Final = proxy.post( + "/policies", {"policy_name": f"policy-{entry}", "guardrails_add": [f"guardrail-{entry}"]} + ) + scenario.cleanups.callback(lambda: proxy.request("DELETE", f"/policies/{policy['policy_id']}")) + attachment: Final = proxy.post("/policies/attachments", {"policy_name": f"policy-{entry}", "tags": [tag]}) + scenario.cleanups.callback( + lambda: proxy.request("DELETE", f"/policies/attachments/{attachment['attachment_id']}") + ) + caller_tagged: Final = _chat(mapped, key, baggage, **{"x-litellm-tags": tag}) + assert caller_tagged.status_code == 200, caller_tagged.text + assert caller_tagged.headers.get("x-litellm-applied-policies") == f"policy-{entry}", caller_tagged.text + assert len(guardrail.drain()) == 1 + baggage_only: Final = _chat(mapped, key, baggage) + assert baggage_only.status_code == 200, baggage_only.text + assert baggage_only.headers.get("x-litellm-applied-policies") is None, baggage_only.text + assert guardrail.drain() == () + + +def _handshake_status(url: str, headers: Mapping[str, str]) -> int: + try: + with connect(url, additional_headers=headers, open_timeout=30): + return 101 + except InvalidStatus as refusal: + return refusal.response.status_code + + +def test_required_mapping_refuses_websocket_inference_before_the_upstream(mapped: Rig) -> None: + proxy_url: Final = str(mapped.proxy.client.base_url).rstrip("/") + upstream_api_key: Final = f"Bearer realtime-{mapped.marker}" + with ( + mapped.proxy.scenario() as scenario, + httpx.Client(base_url=mapped.proxy.upstream_url, timeout=5, trust_env=False) as upstream, + ): + key: Final = scenario.key() + upstream.get("/__observations").raise_for_status() + status: Final = _handshake_status( + f"ws{proxy_url.removeprefix('http')}/v1/realtime?model=realtime-{mapped.marker}", + {"Authorization": f"Bearer {key}", "baggage": f"entrypoint=a,tenant={uuid.uuid4().hex}"}, + ) + assert status == 403, status + observed: Final = list_value(JSON_OBJECT.validate_json(upstream.get("/__observations").content)["requests"]) + upgrades: Final = tuple( + request + for request in (object_value(item) for item in observed) + if request["method"] == "WEBSOCKET" and request["authorization"] == upstream_api_key + ) + assert upgrades == (), upgrades + + +@pytest.mark.parametrize( + ("route", "body"), + [ + ("/utils/token_counter", {"messages": [{"role": "user", "content": "count me"}]}), + ("/v1/messages/count_tokens", {"messages": [{"role": "user", "content": "count me"}]}), + ], + ids=["token_counter", "messages_count_tokens"], +) +def test_required_mapping_skips_unbilled_routes(mapped: Rig, route: str, body: dict[str, JsonValue]) -> None: + with mapped.proxy.scenario() as scenario: + key: Final = scenario.key() + response: Final = mapped.proxy.request("POST", route, {"model": mapped.chat_model, **body}, key=key) + assert response.status_code == 200, response.text + + +def test_custom_passthrough_spend_row_carries_the_mapped_tags(mapped: Rig) -> None: + tenant: Final = uuid.uuid4().hex + path: Final = f"/baggage-passthrough-{tenant}" + with mapped.proxy.scenario() as scenario: + created: Final = mapped.proxy.post( + "/config/pass_through_endpoint", {"path": path, "target": f"{mapped.provider.url}/echo", "auth": True} + ) + endpoint_id: Final = str(object_value(list_value(created["endpoints"])[0])["id"]) + scenario.cleanups.callback( + lambda: mapped.proxy.request("DELETE", "/config/pass_through_endpoint", params={"endpoint_id": endpoint_id}) + ) + key: Final = scenario.key(allowed_passthrough_routes=[path]) + response: Final = mapped.proxy.request( + "POST", + path, + {"model": mapped.chat_model, "messages": [{"role": "user", "content": tenant}]}, + key=key, + headers={"baggage": f"entrypoint=a,tenant={tenant}"}, + ) + assert response.status_code == 200, response.text + assert {mapped.exact_tag, mapped.tenant_tag(tenant)} <= _spend_tags(mapped.proxy, key) + + +def _rejected_client_tag(client_tag: RejectedClientTag, tag: str) -> tuple[Mapping[str, JsonValue], Mapping[str, str]]: + match client_tag: + case "none": + return NO_FIELDS, {} + case "header": + return NO_FIELDS, {"x-litellm-tags": tag} + case "body": + return {"metadata": {"tags": [tag]}}, {} + case "metadata_json": + return {"metadata": json.dumps({"tags": [tag]})}, {} + case "litellm_metadata": + return {"litellm_metadata": {"tags": [tag]}}, {} + + +@pytest.mark.parametrize( + ("client_tag", "baggage_maps", "plain_statuses"), + [ + ("none", True, (200,)), + ("header", True, (400,)), + ("body", True, (400,)), + ("metadata_json", True, (200, 400)), + ("metadata_json", False, (200, 400)), + ("litellm_metadata", True, (200,)), + ], + ids=[ + "baggage_only", + "header_tag", + "body_tag", + "metadata_json_tag", + "metadata_json_tag_unmapped_baggage", + "litellm_metadata_tag", + ], +) +def test_reject_clientside_metadata_tags_with_baggage( + rejecting: Rig, client_tag: RejectedClientTag, baggage_maps: bool, plain_statuses: tuple[int, ...] +) -> None: + tenant: Final = uuid.uuid4().hex + fields, headers = _rejected_client_tag(client_tag, f"caller-{tenant}") + baggage: Final = f"entrypoint=a,tenant={tenant}" if baggage_maps else f"region={tenant}" + mapped_tags: Final = frozenset({rejecting.exact_tag, rejecting.tenant_tag(tenant)} if baggage_maps else ()) + with rejecting.proxy.scenario() as scenario: + plain_key: Final = scenario.key() + key: Final = scenario.key() + plain: Final = _chat(rejecting, plain_key, None, fields=fields, **headers) + response: Final = _chat(rejecting, key, baggage, fields=fields, **headers) + assert plain.status_code in plain_statuses, plain.text + assert response.status_code == plain.status_code, (plain.text, response.text) + if response.status_code != 200: + return + assert _spend_tags(rejecting.proxy, key) == _spend_tags(rejecting.proxy, plain_key) | mapped_tags + + +@pytest.mark.parametrize( + ("baggage_keys", "maps_team"), + [(("team+id",), True), (("team%2Bid",), True), (("team+id", "team%2Bid"), False)], + ids=["plus", "percent_encoded", "both_spellings"], +) +def test_plus_sign_rule_key_matches_either_wire_spelling_but_not_both( + optional: Rig, baggage_keys: tuple[str, ...], maps_team: bool +) -> None: + value: Final = uuid.uuid4().hex + baggage: Final = ",".join(("entrypoint=a", *(f"{baggage_key}={value}" for baggage_key in baggage_keys))) + with optional.proxy.scenario() as scenario: + key: Final = scenario.key() + response: Final = _chat(optional, key, baggage) + assert response.status_code == 200, response.text + team_tags: Final = {optional.team_tag(value)} if maps_team else set() + assert _spend_tags(optional.proxy, key) == {optional.exact_tag} | team_tags + + +@pytest.mark.parametrize("deployment_tagged", [True, False], ids=["tagged_deployment", "untagged_deployment"]) +def test_serving_deployment_tag_selects_its_guardrail_mode_when_baggage_maps_the_same_tag( + optional: Rig, deployment_tagged: bool +) -> None: + model: Final = f"guarded-{optional.marker}" if deployment_tagged else optional.chat_model + with optional.proxy.scenario() as scenario: + key: Final = scenario.key() + optional.provider.drain() + response: Final = _chat(optional, key, "guard=on", model=model) + assert response.status_code == 200, response.text + guardrail_inputs: Final = _guardrail_inputs(optional) + assert guardrail_inputs == (("response",) if deployment_tagged else ()), guardrail_inputs + + +def test_baggage_mapping_a_deployment_tag_keeps_the_team_tag_guardrail_mode_first(optional: Rig) -> None: + model: Final = f"guarded-pre-{optional.marker}" + with optional.proxy.scenario() as scenario: + team: Final = scenario.team(metadata={"tags": [optional.policy_tag]}) + key: Final = scenario.key(team_id=team) + optional.provider.drain() + plain: Final = _chat(optional, key, None, model=model) + plain_inputs: Final = _guardrail_inputs(optional) + response: Final = _chat(optional, key, "guard=pre", model=model) + assert (plain.status_code, response.status_code) == (200, 200), (plain.text, response.text) + assert (plain_inputs, _guardrail_inputs(optional)) == (("response",), ("response",)) + + +@pytest.mark.parametrize("metadata", ["run-123", {"tags": "run-123"}], ids=["non_object", "string_tags"]) +@pytest.mark.parametrize( + ("header_tag", "baggage"), + [(True, "region=sg"), (True, "entrypoint=a"), (False, "region=sg"), (False, "entrypoint=a")], + ids=["header_tag_unmapped_baggage", "header_tag_mapped_baggage", "unmapped_baggage", "mapped_baggage"], +) +def test_malformed_metadata_keeps_its_status_under_baggage( + optional: Rig, header_tag: bool, baggage: str, metadata: JsonValue +) -> None: + headers: Final = {"x-litellm-tags": f"caller-{uuid.uuid4().hex}"} if header_tag else {} + fields: Final[Mapping[str, JsonValue]] = {"metadata": metadata} + with optional.proxy.scenario() as scenario: + key: Final = scenario.key() + optional.provider.drain() + plain: Final = _chat(optional, key, None, fields=fields, **headers) + plain_calls: Final = len(optional.provider.drain()) + response: Final = _chat(optional, key, baggage, fields=fields, **headers) + assert plain.status_code < 500, plain.text + assert plain_calls == (1 if plain.status_code == 200 else 0) + assert (response.status_code, len(optional.provider.drain())) == (plain.status_code, plain_calls), ( + plain.text, + response.text, + ) + + +@pytest.mark.parametrize("outage_flag", [False, True], ids=["outage_flag_off", "outage_flag_on"]) +def test_budgeted_mapping_on_a_failed_tag_read(gateway: Gateway, tmp_path: Path, outage_flag: bool) -> None: + marker: Final = uuid.uuid4().hex[:12] + value: Final = uuid.uuid4().hex + tag: Final = f"budgeted-{marker}:{value}" + with ExitStack() as stack: + database_url: Final = stack.enter_context(scratch_database()) + relay, relayed_url = stack.enter_context(dropped_connection_relay(database_url, tag.encode())) + provider: Final = stack.enter_context(wire_server(_provider)) + config: Final = _write_config( + tmp_path, + [_deployment(f"chat-{marker}", provider)], + { + "request_tag_mappings": [ + {"key": "budgeted", "tag_prefix": f"budgeted-{marker}:", "require_budgeted_tag": True} + ], + "allow_requests_on_db_unavailable": outage_flag, + }, + {}, + ) + proxy: Final = stack.enter_context( + owned_proxy( + gateway, + tmp_path, + {"DATABASE_URL": plaintext_database_url(relayed_url), "PRISMA_HEALTH_WATCHDOG_ENABLED": "false"}, + config=config, + remove_environment=("DATABASE_URL_READ_REPLICA",), + ) + ) + scenario: Final = stack.enter_context(proxy.scenario()) + key: Final = scenario.key(max_budget=50) + proxy.post("/tag/new", {"name": tag, "max_budget": 100}) + rig: Final = Rig(proxy, provider, provider, marker) + warmed: Final = _chat(rig, key, None) + assert warmed.status_code == 200, warmed.text + provider.drain() + relay.arm() + response: Final = _chat(rig, key, f"budgeted={value}") + relay.disarm() + assert relay.dropped.is_set(), response.text + if not outage_flag: + assert response.status_code == 503, response.text + assert provider.drain() == () + return + assert response.status_code == 200, response.text + assert response.headers.get("x-litellm-key-max-budget") == "50.0", dict(response.headers) + assert len(provider.drain()) == 1 diff --git a/tests/integration/spend/test_tag_budget_enforcement.py b/tests/integration/spend/test_tag_budget_enforcement.py index 0c0db2943c6..5430f77a3ec 100644 --- a/tests/integration/spend/test_tag_budget_enforcement.py +++ b/tests/integration/spend/test_tag_budget_enforcement.py @@ -105,6 +105,46 @@ def test_spend_over_a_tag_max_budget_rejects_the_next_request(gateway: Gateway) assert control.status_code == 200, control.text +def test_tag_new_counts_spend_from_creation_not_before(gateway: Gateway) -> None: + tag: Final = f"tag-created-late-{uuid.uuid4().hex}" + + def delete_tag() -> None: + gateway.post("/tag/delete", {"name": tag}) + + with gateway.scenario() as scenario: + model: Final = scenario.model(input_cost_per_token=0.01, output_cost_per_token=0.01) + + def tagged(text: str) -> httpx.Response: + return gateway.request( + "POST", + "/v1/chat/completions", + { + "model": model, + "messages": [{"role": "user", "content": f"{text} {uuid.uuid4().hex}"}], + "metadata": {"tags": [tag]}, + }, + ) + + untracked: Final = tagged("spend before the tag exists") + assert untracked.status_code == 200, untracked.text + with Redis(host=os.environ["REDIS_HOST"], port=int(os.environ["REDIS_PORT"])) as cache: + counted: Final = eventually( + lambda: float(cache.get(f"spend:tag:{tag}") or 0), + lambda spend: spend > 0.0001, + seconds=70, + ) + assert counted > 0.0001, counted + gateway.post("/tag/new", {"name": tag, "max_budget": 0.0001}) + scenario.cleanups.callback(delete_tag) + admitted: Final = tagged("first request under the new budget") + assert admitted.status_code == 200, admitted.text + blocked: Final = eventually( + lambda: tagged("budget probe"), lambda response: response.status_code != 200, seconds=70 + ) + assert blocked.status_code == 422, blocked.text + assert "budget" in blocked.text.lower(), blocked.text + + def test_tag_object_rpm_limit_rejects_the_second_request_across_keys(gateway: Gateway) -> None: tag: Final = f"tag-rpm-{uuid.uuid4().hex}" @@ -446,9 +486,7 @@ def test_tag_budget_duration_can_be_cleared_and_limits_are_read_back(gateway: Ga "rpm_limit": budget_info["rpm_limit"], "budget_duration": budget_info["budget_duration"], "budget_reset_at": budget_info["budget_reset_at"], - } == {"max_budget": 100, "rpm_limit": 1, "budget_duration": None, "budget_reset_at": None}, ( - info_response.text - ) + } == {"max_budget": 100, "rpm_limit": 1, "budget_duration": None, "budget_reset_at": None}, info_response.text tag_rows: Final = read_rows('SELECT budget_id FROM "LiteLLM_TagTable" WHERE tag_name = %s', (tag,)) assert len(tag_rows) == 1, tag_rows @@ -457,9 +495,9 @@ def test_tag_budget_duration_can_be_cleared_and_limits_are_read_back(gateway: Ga "WHERE budget_id = %s", (str(tag_rows[0]["budget_id"]),), ) - assert budget_rows == [ - {"max_budget": 100, "rpm_limit": 1, "budget_duration": None, "budget_reset_at": None} - ], budget_rows + assert budget_rows == [{"max_budget": 100, "rpm_limit": 1, "budget_duration": None, "budget_reset_at": None}], ( + budget_rows + ) @pytest.mark.parametrize("tag_source", _TAG_SOURCES, ids=_TAG_SOURCE_IDS) diff --git a/tests/unit/enterprise/integrations/test_custom_guardrail.py b/tests/unit/enterprise/integrations/test_custom_guardrail.py index c6e48061698..84f34144d99 100644 --- a/tests/unit/enterprise/integrations/test_custom_guardrail.py +++ b/tests/unit/enterprise/integrations/test_custom_guardrail.py @@ -1,3 +1,5 @@ +import pytest + from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.types.guardrails import GuardrailEventHooks, Mode @@ -249,3 +251,82 @@ def test_custom_guardrail_with_mode(monkeypatch): ) is False ) + + +@pytest.mark.parametrize("caller_tags,runs_pre_call", [((), True), (("test_tag",), False)]) +def test_custom_guardrail_mode_ignores_baggage_mapped_tags_the_caller_did_not_send( + monkeypatch, caller_tags, runs_pre_call +): + from types import SimpleNamespace + from typing import cast + + from fastapi import Request + + from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup + + monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True) + cg = CustomGuardrail( + guardrail_name="test_guardrail", + supported_event_hooks=[GuardrailEventHooks.pre_call, GuardrailEventHooks.logging_only], + event_hook=Mode(tags={"test_tag": "logging_only"}, default="pre_call"), + default_on=True, + ) + request = cast(Request, SimpleNamespace(state=SimpleNamespace(litellm_baggage_mapped_tags=("test_tag",)))) + data = {"messages": [{"role": "user", "content": "test"}], "metadata": {"tags": []}} + LiteLLMProxyRequestSetup.apply_authenticated_baggage_tags(request, data) + data["metadata"]["caller_tags"] = caller_tags + assert cg.should_run_guardrail(data=data, event_type=GuardrailEventHooks.pre_call) is runs_pre_call + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "team_tags,deployment_tags,runs_post_call", + [ + (["policy:sensitive"], ["usecase:dev"], True), + ([], ["policy:sensitive", "usecase:dev"], True), + (["policy:sensitive", "usecase:dev"], [], True), + ([], ["usecase:dev"], False), + ([], [], True), + ], +) +async def test_baggage_never_changes_the_guardrail_mode_key_team_and_deployment_tags_select( + monkeypatch, team_tags, deployment_tags, runs_post_call +): + from types import SimpleNamespace + from unittest.mock import MagicMock + + import litellm + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request + + monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True) + cg = CustomGuardrail( + guardrail_name="test_guardrail", + supported_event_hooks=[GuardrailEventHooks.post_call, GuardrailEventHooks.logging_only], + event_hook=Mode(tags={"policy:sensitive": "post_call", "usecase:dev": "logging_only"}, default="post_call"), + default_on=True, + ) + router = litellm.Router( + model_list=[{"model_name": "gpt-4", "litellm_params": {"model": "gpt-4o", "tags": deployment_tags}}] + ) + + async def runs_post_call_guardrail(baggage_tags: tuple[str, ...]) -> bool: + request = MagicMock() + request.url.path = "/chat/completions" + request.method = "POST" + request.query_params = {} + request.headers = {"Content-Type": "application/json"} + request.client.host = "127.0.0.1" + request.state = SimpleNamespace(litellm_baggage_mapped_tags=baggage_tags) + data = await add_litellm_data_to_request( + data={"model": "gpt-4", "messages": [{"role": "user", "content": "test"}]}, + request=request, + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key", team_metadata={"tags": team_tags}), + proxy_config=MagicMock(), + general_settings={}, + ) + await router.acompletion(**data, mock_response="hi") + return cg.should_run_guardrail(data=data, event_type=GuardrailEventHooks.post_call) + + assert await runs_post_call_guardrail(()) is runs_post_call + assert await runs_post_call_guardrail(("usecase:dev",)) is runs_post_call diff --git a/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py b/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py index c583407e61e..effdd9725c2 100644 --- a/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py +++ b/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py @@ -3753,6 +3753,85 @@ async def test_tag_max_budget_check_still_enforces_registered_tag_over_budget(): assert [call.kwargs["where"]["tag_name"]["in"] for call in batch_calls] == [["paid-tag"]] +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("client_tags", "key_metadata"), + [ + (b'"run-123"', {"tags": ["engineering"]}), + (b'""', {}), + (b"{}", {}), + (b"false", {}), + (b"0", {}), + ], + ids=["string_replaced_by_key_tags", "empty_string", "empty_object", "false", "zero"], +) +async def test_common_checks_rejects_a_spent_mapped_tag_whatever_non_list_tags_the_client_sends( + client_tags: bytes, key_metadata: dict[str, list[str]] +): + from litellm.proxy.auth.auth_checks import common_checks + from litellm.proxy.common_utils.http_parsing_utils import read_request_body + from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup + from litellm.proxy.utils import ProxyLogging + + async def fake_find_many(**kwargs): + if "where" not in kwargs: + return [_tag_registry_row("usecase:spent")] + return [_tag_db_row(name, max_budget=1.0) for name in kwargs["where"]["tag_name"]["in"]] + + async def spent_tag_spend(counter_key, fallback_spend, max_budget=None, **kwargs): + return 1.5 if counter_key == "spend:tag:usecase:spent" else fallback_spend + + async def receive(): + return { + "type": "http.request", + "body": b'{"model":"gpt-4o","messages":[{"role":"user","content":"hi"}],"metadata":{"tags":' + + client_tags + + b"}}", + } + + mock_prisma: Final = MagicMock() + mock_prisma.db.litellm_tagtable.find_many = AsyncMock(side_effect=fake_find_many) + request: Final = Request( + { + "type": "http", + "method": "POST", + "path": "/v1/chat/completions", + "scheme": "http", + "server": ("localhost", 4000), + "query_string": b"", + "headers": [(b"content-type", b"application/json"), (b"baggage", b"entrypoint=spent")], + }, + receive, + ) + request_body: Final = await read_request_body(request) + await LiteLLMProxyRequestSetup.apply_baggage_tag_mappings_pre_auth( + request, + request_body, + {"request_tag_mappings": [{"key": "entrypoint", "tag_prefix": "usecase:"}]}, + None, + MagicMock(), + ) + with ( + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), + patch("litellm.proxy.proxy_server.get_current_spend", spent_tag_spend), + pytest.raises(litellm.BudgetExceededError) as exc_info, + ): + await common_checks( + request_body=request_body, + team_object=None, + user_object=None, + end_user_object=None, + global_proxy_spend=None, + general_settings={}, + route="/v1/chat/completions", + llm_router=None, + proxy_logging_obj=ProxyLogging(user_api_key_cache=None), + valid_token=UserAPIKeyAuth(token="test-token", metadata=key_metadata), + request=request, + ) + assert exc_info.value.entity_id == "usecase:spent" + + @pytest.mark.asyncio async def test_get_team_object_raises_404_when_not_found(): from unittest.mock import AsyncMock, MagicMock @@ -11541,3 +11620,84 @@ async def test_authoritative_group_grants_propagate_policy_outages( await get_agent_ids_from_access_groups(["group"], check_db_only=True) else: assert await get_agent_ids_from_access_groups(["group"]) == [] + + +def _client_tag_request(client_tags, baggage): + return Request( + { + "type": "http", + "headers": [ + *([(b"baggage", b"entrypoint=a")] if baggage else []), + *([(b"x-litellm-tags", b"usecase:a")] if client_tags == "header" else []), + ], + } + ) + + +def _client_tag_body(client_tags, client_slot): + if client_tags is None or client_tags == "header": + return {} + return {client_slot: client_tags if isinstance(client_tags, str) else {"tags": client_tags}} + + +def _client_tag_rejection(route, data, request): + from litellm.proxy.auth.auth_checks import _reject_clientside_metadata_tags_check + + try: + _reject_clientside_metadata_tags_check({"reject_clientside_metadata_tags": True}, data, route, request) + except ProxyException as error: + return error.code + return None + + +@pytest.mark.asyncio +@pytest.mark.parametrize("route", ["/v1/chat/completions", "/v1/messages", "/v1/responses"]) +@pytest.mark.parametrize("client_tags", [None, [], ["usecase:a"], '{"tags":["usecase:a"]}', "header"]) +@pytest.mark.parametrize("client_slot", ["metadata", "litellm_metadata"]) +@pytest.mark.parametrize("baggage", [True, False]) +async def test_baggage_mapping_rejects_exactly_the_client_tags_rejected_without_it( + route, client_tags, client_slot, baggage +): + from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup + + unmapped_request = _client_tag_request(client_tags, baggage=False) + unmapped = _client_tag_body(client_tags, client_slot) + LiteLLMProxyRequestSetup.pre_seed_litellm_metadata_for_route(unmapped, route) + LiteLLMProxyRequestSetup.apply_client_tag_policy_pre_auth(unmapped_request, unmapped, UserAPIKeyAuth(api_key="qa-key")) + + mapped_request = _client_tag_request(client_tags, baggage=baggage) + mapped = _client_tag_body(client_tags, client_slot) + LiteLLMProxyRequestSetup.pre_seed_litellm_metadata_for_route(mapped, route) + await LiteLLMProxyRequestSetup.apply_baggage_tag_mappings_pre_auth( + mapped_request, + mapped, + {"request_tag_mappings": [{"key": "entrypoint", "tag_prefix": "usecase:"}]}, + None, + MagicMock(), + ) + LiteLLMProxyRequestSetup.apply_client_tag_policy_pre_auth(mapped_request, mapped, UserAPIKeyAuth(api_key="qa-key")) + + assert _client_tag_rejection(route, mapped, mapped_request) == _client_tag_rejection( + route, unmapped, unmapped_request + ) + if client_tags is None and baggage: + assert "usecase:a" in mapped["metadata" if "chat/completions" in route else "litellm_metadata"]["tags"] + + +@pytest.mark.asyncio +async def test_baggage_mapping_still_rejects_client_metadata_tags(): + from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup + + request = _client_tag_request(["usecase:b"], baggage=True) + data = {"metadata": {"tags": ["usecase:b"]}} + await LiteLLMProxyRequestSetup.apply_baggage_tag_mappings_pre_auth( + request, + data, + {"request_tag_mappings": [{"key": "entrypoint", "tag_prefix": "usecase:"}]}, + None, + MagicMock(), + ) + assert _client_tag_rejection("/v1/chat/completions", data, request) == "400" + from litellm.utils import get_request_guardrails + + assert get_request_guardrails(data) == [] diff --git a/tests/unit/proxy/auth/test_user_api_key_auth.py b/tests/unit/proxy/auth/test_user_api_key_auth.py index 3660019cc2c..60a8dd73a34 100644 --- a/tests/unit/proxy/auth/test_user_api_key_auth.py +++ b/tests/unit/proxy/auth/test_user_api_key_auth.py @@ -2102,3 +2102,183 @@ async def test_jwt_path_skips_user_model_budget_only_for_zero_spend_mcp_methods( assert result.user_id == "jwt-user" assert user_model_budget_check.await_count == (1 if expect_user_model_budget_check else 0) + + +def _qa_websocket(messages): + async def receive(): + return {"type": "websocket.connect"} + + async def send(message): + messages.append(message) + + return WebSocket( + { + "type": "websocket", + "path": "/v1/responses", + "scheme": "ws", + "server": ("localhost", 4000), + "query_string": b"model=qa-model", + "headers": [(b"authorization", b"Bearer qa-key")], + }, + receive=receive, + send=send, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("close_on_auth_error", [True, False]) +async def test_websocket_inference_is_denied_when_a_required_baggage_mapping_is_configured(close_on_auth_error): + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth_websocket_for_model + + messages = [] + with ( + patch.dict( + litellm.proxy.proxy_server.general_settings, + { + "request_tag_mappings": [ + {"key": "region", "tag_prefix": "region:"}, + {"key": "entrypoint", "tag_prefix": "usecase:", "required": True}, + ] + }, + ), + patch( + "litellm.proxy.auth.user_api_key_auth.user_api_key_auth", + new=AsyncMock(return_value=UserAPIKeyAuth(api_key="qa-key")), + ), + ): + with pytest.raises(HTTPException, match="WebSocket inference is disabled") as error: + await user_api_key_auth_websocket_for_model( + _qa_websocket(messages), "qa-model", close_on_auth_error=close_on_auth_error + ) + assert error.value.status_code == 403 + assert messages == ([{"type": "websocket.close", "code": 1008, "reason": ""}] if close_on_auth_error else []) + + +@pytest.mark.asyncio +async def test_websocket_invalid_key_is_rejected_before_the_required_baggage_mapping_denial(): + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth_websocket_for_model + + with ( + patch.dict( + litellm.proxy.proxy_server.general_settings, + {"request_tag_mappings": [{"key": "entrypoint", "tag_prefix": "usecase:", "required": True}]}, + ), + patch( + "litellm.proxy.auth.user_api_key_auth.user_api_key_auth", + new=AsyncMock(side_effect=HTTPException(status_code=401, detail="Invalid proxy server token passed")), + ), + ): + with pytest.raises(HTTPException) as error: + await user_api_key_auth_websocket_for_model(_qa_websocket([]), "qa-model", close_on_auth_error=False) + assert error.value.status_code == 401 + + +@pytest.mark.asyncio +async def test_websocket_inference_is_allowed_when_no_baggage_mapping_is_required(): + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth_websocket_for_model + + messages = [] + authorized = UserAPIKeyAuth(api_key="qa-key") + with ( + patch.dict( + litellm.proxy.proxy_server.general_settings, + {"request_tag_mappings": [{"key": "entrypoint", "tag_prefix": "usecase:", "require_budgeted_tag": True}]}, + ), + patch( + "litellm.proxy.auth.user_api_key_auth.user_api_key_auth", new=AsyncMock(return_value=authorized) + ) as http_auth, + ): + result = await user_api_key_auth_websocket_for_model(_qa_websocket(messages), "qa-model") + assert result is authorized + assert http_auth.await_args.kwargs["api_key"] == "Bearer qa-key" + assert messages == [] + + +async def _auth_then_serve_chat(monkeypatch, general_settings, body, headers): + import copy + + from fastapi import Request as FastAPIRequest + from fastapi import Response + + from litellm.proxy.common_utils.http_parsing_utils import read_request_body + from litellm.proxy.proxy_server import chat_completion + + router = litellm.Router( + model_list=[ + { + "model_name": "gpt-4", + "litellm_params": {"model": "gpt-4o", "tags": ["usecase:a"], "mock_response": "hi"}, + "model_info": {"id": "usecase-deployment"}, + }, + { + "model_name": "gpt-4", + "litellm_params": {"model": "gpt-4o-mini", "tags": ["default"], "mock_response": "hi"}, + "model_info": {"id": "default-deployment"}, + }, + ], + enable_tag_filtering=True, + ) + monkeypatch.setattr(litellm.proxy.proxy_server, "master_key", MASTER_KEY) + monkeypatch.setattr(litellm.proxy.proxy_server, "prisma_client", None) + monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", router) + monkeypatch.setattr(litellm.proxy.proxy_server, "general_settings", general_settings) + + async def receive(): + return {"type": "http.request", "body": json.dumps(body).encode()} + + request = FastAPIRequest( + { + "type": "http", + "method": "POST", + "path": "/chat/completions", + "scheme": "http", + "server": ("localhost", 4000), + "query_string": b"", + "headers": [(b"content-type", b"application/json"), *headers], + }, + receive, + ) + auth = await user_api_key_auth(request=request, api_key=f"Bearer {MASTER_KEY}") + cached_body = copy.deepcopy(await read_request_body(request)) + response = Response() + await chat_completion(request=request, fastapi_response=response, user_api_key_dict=auth) + return cached_body, response.headers["x-litellm-model-id"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "general_settings", + [{}, {"request_tag_mappings": [{"key": "region", "tag_prefix": "region:"}]}], + ids=["no_mappings", "mapped_baggage"], +) +async def test_auth_header_tag_merge_keeps_json_string_metadata_tags_routable(monkeypatch, general_settings): + body = { + "model": "gpt-4", + "messages": [{"role": "user", "content": "hi"}], + "metadata": json.dumps({"tags": ["usecase:a"]}), + } + cached_body, served_by = await _auth_then_serve_chat( + monkeypatch, general_settings, body, [(b"x-litellm-tags", b"other"), (b"baggage", b"region=sg")] + ) + assert served_by == "usecase-deployment" + assert cached_body["metadata"] == {"tags": ["usecase:a", "other"]} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("client_metadata", ["run-123", 42, ["x"]]) +@pytest.mark.parametrize( + "baggage", + [[], [(b"baggage", b"region=sg")], [(b"baggage", b"entrypoint=a")]], + ids=["no_baggage", "unmapped_baggage", "mapped_baggage"], +) +async def test_non_object_metadata_is_served_like_unset_mappings_whatever_the_baggage_maps( + monkeypatch, client_metadata, baggage +): + body = {"model": "gpt-4", "messages": [{"role": "user", "content": "hi"}], "metadata": client_metadata} + headers = [(b"x-litellm-tags", b"usecase:a"), *baggage] + unset = await _auth_then_serve_chat(monkeypatch, {}, body, headers) + mapped = await _auth_then_serve_chat( + monkeypatch, {"request_tag_mappings": [{"key": "entrypoint", "tag_prefix": "usecase:"}]}, body, headers + ) + assert mapped == unset + assert mapped[1] == "usecase-deployment" diff --git a/tests/unit/proxy/common_utils/test_request_tag_mappings.py b/tests/unit/proxy/common_utils/test_request_tag_mappings.py new file mode 100644 index 00000000000..e1644544c46 --- /dev/null +++ b/tests/unit/proxy/common_utils/test_request_tag_mappings.py @@ -0,0 +1,473 @@ +import json +from types import SimpleNamespace +from typing import cast +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from fastapi import Request +from pydantic import ValidationError + +from litellm.models.budget import LiteLLM_BudgetTable +from litellm.constants import BAGGAGE_MAPPED_TAGS_METADATA_KEY +from litellm.models.tag import LiteLLM_TagTable +from litellm.proxy._types import ConfigGeneralSettings, ProxyException +from litellm.proxy.common_utils.http_parsing_utils import read_request_body +from litellm.proxy.common_utils.request_tag_mappings import is_baggage_mapped_route, resolve_baggage_tag_mappings +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache, tag_cache_key +from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup, add_litellm_data_to_request +from litellm.types.proxy.request_tag_mappings import RequestTagMapping + + +def test_accounting_enrichment_without_request_state_preserves_legacy_metadata(): + request = cast(Request, SimpleNamespace(headers={})) + data = {"metadata": {"tags": ["existing"]}} + LiteLLMProxyRequestSetup.apply_authenticated_baggage_tags(request, data) + assert data == {"metadata": {"tags": ["existing"]}} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("header", [b"", b"entrypoint=a,entrypoint=b"]) +@pytest.mark.parametrize("client_metadata", [None, '{"requester_metadata":{"guardrails":["qa-guard"]}}']) +async def test_mapping_rejection_keeps_string_metadata_safe_for_auth_failure_logging(header, client_metadata): + from litellm.utils import get_request_guardrails + + request = Request({"type": "http", "headers": [(b"baggage", header)]}) + data = {"metadata": client_metadata} + with pytest.raises(ProxyException) as error: + await LiteLLMProxyRequestSetup.apply_baggage_tag_mappings_pre_auth( + request, + data, + {"request_tag_mappings": [{"key": "entrypoint", "tag_prefix": "usecase:", "required": True}]}, + None, + MagicMock(), + ) + assert error.value.code == "400" + assert get_request_guardrails(data) == (["qa-guard"] if client_metadata is not None else []) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("slot", ["metadata", "litellm_metadata"]) +@pytest.mark.parametrize("invalid_metadata", [[], 42, "not-json", '"not-an-object"']) +async def test_invalid_mapper_metadata_maps_like_empty_metadata_and_keeps_the_client_cache(slot, invalid_metadata): + from litellm.proxy.auth.user_api_key_auth import _read_request_body_deferring_parse_failure + from litellm.utils import get_request_guardrails + + body = {"metadata": '{"tags":["client"]}', "litellm_metadata": '{"tags":["alternate"]}'} + body[slot] = invalid_metadata + + async def receive(): + return {"type": "http.request", "body": json.dumps(body).encode()} + + request = Request( + {"type": "http", "method": "POST", "path": "/v1/responses", "headers": [(b"baggage", b"entrypoint=a")]}, + receive, + ) + data, parse_error = await _read_request_body_deferring_parse_failure(request) + assert parse_error is None + await LiteLLMProxyRequestSetup.apply_baggage_tag_mappings_pre_auth( + request, + data, + {"request_tag_mappings": [{"key": "entrypoint", "tag_prefix": "usecase:"}]}, + None, + MagicMock(), + ) + assert "usecase:a" in data["litellm_metadata"]["tags"] + assert "tags" not in data[slot] or slot == "litellm_metadata" + assert get_request_guardrails(data) == [] + assert await read_request_body(request) == body + + +@pytest.mark.parametrize( + "header,value", + [ + ("entrypoint=a", "a"), + ("entrypoint = a", "a"), + ("region=sg, entrypoint=customer%20support;version=1", "customer support"), + ], +) +def test_variable_mapper_uses_decoded_baggage_value(header, value): + rule = RequestTagMapping(key="entrypoint", tag_prefix="usecase:") + assert resolve_baggage_tag_mappings(header, [rule]) == ((f"usecase:{value}", False),) + + +@pytest.mark.parametrize( + "header,expected", + [("entrypoint=a", "department:support"), ("entrypoint=b", "department:sales"), ("entrypoint=x", "usecase:x")], +) +def test_exact_rules_override_variable_fallback_independent_of_order(header, expected): + rules = [ + RequestTagMapping(key="entrypoint", tag_prefix="usecase:"), + RequestTagMapping(key="entrypoint", value="a", tag="department:support"), + RequestTagMapping(key="entrypoint", value="b", tag="department:sales"), + ] + assert resolve_baggage_tag_mappings(header, rules) == ((expected, False),) + assert resolve_baggage_tag_mappings(header, list(reversed(rules))) == ((expected, False),) + + +@pytest.mark.parametrize("header", ["entrypoint=%00", "entrypoint=a%0Ab", "entrypoint=b,entry%70oint=a"]) +def test_variable_mapper_skips_control_characters_and_percent_encoded_duplicate_keys(header): + rule = RequestTagMapping(key="entrypoint", tag_prefix="usecase:") + assert resolve_baggage_tag_mappings(header, [rule]) == () + + +@pytest.mark.parametrize( + "header,expected", + [ + ("entry+point=a", (("usecase:a", False),)), + ("entry%2Bpoint=a", (("usecase:a", False),)), + ("entry+point=a,entry%2Bpoint=b", ()), + ], +) +def test_key_with_a_plus_matches_both_wire_spellings_and_rejects_both_at_once(header, expected): + rule = RequestTagMapping(key="entry+point", tag_prefix="usecase:") + assert resolve_baggage_tag_mappings(header, [rule]) == expected + + +def test_exact_rules_work_without_a_variable_rule(): + rule = RequestTagMapping(key="entrypoint", value="a", tag="usecase:manual") + assert resolve_baggage_tag_mappings("entrypoint=a", [rule]) == (("usecase:manual", False),) + assert resolve_baggage_tag_mappings("entrypoint=x", [rule]) == () + + +@pytest.mark.parametrize( + "header", ["", "region=sg", "entrypoint=", "entrypoint=a,entrypoint=b", "entrypoint=b,entry%70oint=a", "entrypoint=x"] +) +def test_required_mapping_rejects_missing_ambiguous_or_unmapped_values(header): + rule = RequestTagMapping(key="entrypoint", value="a", tag="usecase:a", required=True) + with pytest.raises(ProxyException) as error: + resolve_baggage_tag_mappings(header, [rule]) + assert error.value.code == "400" + + +@pytest.mark.parametrize( + "rule", + [ + {"key": "entrypoint", "tag": "usecase:a"}, + {"key": "entrypoint", "value": "a", "tag_prefix": "usecase:"}, + {"key": "entrypoint", "value": "a", "tag": "usecase:a", "tag_prefix": "usecase:"}, + {"key": "", "tag_prefix": "usecase:"}, + ], +) +def test_invalid_rule_is_rejected_during_config_validation(rule): + with pytest.raises(ValidationError): + ConfigGeneralSettings.model_validate({"request_tag_mappings": [rule]}) + + +def test_duplicate_fallbacks_are_rejected(): + with pytest.raises(ValidationError, match="Only one mapping"): + ConfigGeneralSettings.model_validate( + { + "request_tag_mappings": [ + {"key": "entrypoint", "tag_prefix": "one:"}, + {"key": "entrypoint", "tag_prefix": "two:"}, + ] + } + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "route,slot", + [("/chat/completions", "metadata"), ("/v1/messages", "litellm_metadata"), ("/v1/responses", "litellm_metadata")], +) +async def test_mapped_tag_merges_with_existing_tags_in_the_budget_metadata_slot(route, slot): + data = {slot: {"tags": ["existing", "usecase:a"]}} + request = Request({"type": "http", "headers": [(b"baggage", b"entrypoint=a")]}) + LiteLLMProxyRequestSetup.pre_seed_litellm_metadata_for_route(data, route) + await LiteLLMProxyRequestSetup.apply_baggage_tag_mappings_pre_auth( + request, + data, + {"request_tag_mappings": [{"key": "entrypoint", "tag_prefix": "usecase:"}]}, + None, + MagicMock(), + ) + assert data[slot]["tags"] == ["existing", "usecase:a"] + + +@pytest.mark.asyncio +async def test_budgeted_mapper_requires_a_database(): + request = Request({"type": "http", "headers": [(b"baggage", b"entrypoint=a")]}) + config = {"request_tag_mappings": [{"key": "entrypoint", "tag_prefix": "usecase:", "require_budgeted_tag": True}]} + with pytest.raises(ProxyException) as error: + await LiteLLMProxyRequestSetup.apply_baggage_tag_mappings_pre_auth( + request, {}, config, None, MagicMock() + ) + assert error.value.code == "503" + + +@pytest.mark.asyncio +async def test_budgeted_mapper_answers_503_when_the_tag_budget_lookup_fails(): + request = Request({"type": "http", "headers": [(b"baggage", b"entrypoint=a")]}) + prisma_client = MagicMock() + prisma_client.db.litellm_tagtable.find_many = AsyncMock(side_effect=Exception("database unreachable")) + config = {"request_tag_mappings": [{"key": "entrypoint", "tag_prefix": "usecase:", "require_budgeted_tag": True}]} + with pytest.raises(ProxyException, match="could not read tag budgets") as error: + await LiteLLMProxyRequestSetup.apply_baggage_tag_mappings_pre_auth( + request, {"metadata": {}}, config, prisma_client, UserApiKeyCache() + ) + assert error.value.code == "503" + + +@pytest.mark.asyncio +async def test_budgeted_mapper_serves_through_a_db_outage_when_requests_are_allowed_on_db_unavailable(): + request = Request({"type": "http", "headers": [(b"baggage", b"entrypoint=a")]}) + prisma_client = MagicMock() + prisma_client.db.litellm_tagtable.find_many = AsyncMock(side_effect=Exception("database unreachable")) + data = {"metadata": {}} + config = {"request_tag_mappings": [{"key": "entrypoint", "tag_prefix": "usecase:", "require_budgeted_tag": True}]} + with patch.dict("litellm.proxy.proxy_server.general_settings", {"allow_requests_on_db_unavailable": True}): + await LiteLLMProxyRequestSetup.apply_baggage_tag_mappings_pre_auth( + request, data, config, prisma_client, UserApiKeyCache() + ) + assert data["metadata"]["tags"] == ["usecase:a"] + + +async def _cache_with_tag(tag: LiteLLM_TagTable) -> UserApiKeyCache: + cache = UserApiKeyCache() + await cache.async_set_cache(key=tag_cache_key(tag.tag_name), value=tag, model_type=LiteLLM_TagTable) + return cache + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "budget_table", + [None, LiteLLM_BudgetTable(budget_id="qa-budget", max_budget=None, tpm_limit=10)], +) +async def test_budgeted_mapper_rejects_a_tag_without_a_max_budget(budget_table): + request = Request({"type": "http", "headers": [(b"baggage", b"entrypoint=a")]}) + data = {"metadata": {}} + cache = await _cache_with_tag(LiteLLM_TagTable(tag_name="usecase:a", litellm_budget_table=budget_table)) + config = {"request_tag_mappings": [{"key": "entrypoint", "tag_prefix": "usecase:", "require_budgeted_tag": True}]} + with pytest.raises(ProxyException, match="usecase:a requires an existing tag budget") as error: + await LiteLLMProxyRequestSetup.apply_baggage_tag_mappings_pre_auth( + request, data, config, MagicMock(), cache + ) + assert error.value.code == "400" + assert data == {"metadata": {}} + + +@pytest.mark.asyncio +async def test_budgeted_mapper_admits_a_tag_with_a_max_budget_and_unbudgeted_rules_need_none(): + request = Request({"type": "http", "headers": [(b"baggage", b"entrypoint=a,region=sg")]}) + data = {"metadata": {}} + cache = await _cache_with_tag( + LiteLLM_TagTable( + tag_name="usecase:a", litellm_budget_table=LiteLLM_BudgetTable(budget_id="qa-budget", max_budget=1.0) + ) + ) + config = { + "request_tag_mappings": [ + {"key": "entrypoint", "tag_prefix": "usecase:", "require_budgeted_tag": True}, + {"key": "region", "tag_prefix": "region:"}, + ] + } + await LiteLLMProxyRequestSetup.apply_baggage_tag_mappings_pre_auth( + request, data, config, MagicMock(), cache + ) + assert data["metadata"]["tags"] == ["usecase:a", "region:sg"] + + +@pytest.mark.parametrize( + "route", + ["/v1/chat/completions", "/chat/completions", "/v1/messages", "/v1/responses", "/a2a/qa-agent/message/send"], +) +def test_inference_routes_take_baggage_mappings(route): + assert is_baggage_mapped_route(route) + + +@pytest.mark.parametrize( + "route", + [ + "/utils/token_counter", + "/utils/transform_request", + "/v1/messages/count_tokens", + "/v1/responses/input_tokens", + "/v1beta/models/gemini-pro:countTokens", + "/mcp", + ], +) +def test_unbilled_and_mcp_routes_skip_baggage_mappings(route): + assert not is_baggage_mapped_route(route) + + +@pytest.mark.parametrize("registered", [True, False]) +def test_custom_pass_through_routes_take_baggage_mappings_once_registered(registered): + from litellm.proxy.pass_through_endpoints.pass_through_endpoints import _registered_pass_through_routes + + endpoints = {"qa-endpoint:exact:/billing-api": {"target": "http://example.com", "headers": {}}} + with patch.dict(_registered_pass_through_routes, endpoints if registered else {}, clear=True): + assert is_baggage_mapped_route("/billing-api") is registered + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "route,slot", + [("/chat/completions", "metadata"), ("/v1/messages", "litellm_metadata"), ("/v1/responses", "litellm_metadata")], +) +@pytest.mark.parametrize("query", [b"api-version=qa-version", b"key=qa-virtual-key"]) +@pytest.mark.parametrize("client_metadata", [None, '{"requester_metadata":{"guardrails":[]}}']) +async def test_mapping_survives_cached_body_read_without_leaking_auth_query_fields(route, slot, query, client_metadata): + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.proxy_server import ProxyConfig + + body = {"model": "qa-model"} + if client_metadata is not None: + body["metadata"] = client_metadata + + async def receive(): + return {"type": "http.request", "body": json.dumps(body).encode()} + + request = Request( + { + "type": "http", + "method": "POST", + "path": route, + "scheme": "http", + "server": ("localhost", 4000), + "query_string": query, + "headers": [(b"content-type", b"application/json"), (b"baggage", b"entrypoint=a")], + }, + receive, + ) + from litellm.proxy.auth.user_api_key_auth import _read_request_body_deferring_parse_failure + + data, parse_error = await _read_request_body_deferring_parse_failure(request) + assert parse_error is None + LiteLLMProxyRequestSetup.pre_seed_litellm_metadata_for_route(data, route) + await LiteLLMProxyRequestSetup.apply_baggage_tag_mappings_pre_auth( + request, + data, + {"request_tag_mappings": [{"key": "entrypoint", "tag_prefix": "usecase:"}]}, + None, + MagicMock(), + ) + downstream = await read_request_body(request) + assert downstream == body + enriched = await add_litellm_data_to_request( + data=downstream, + request=request, + user_api_key_dict=UserAPIKeyAuth(api_key="qa-key"), + proxy_config=ProxyConfig(), + general_settings={}, + ) + assert enriched[slot]["tags"] == ["usecase:a"] + assert "key" not in enriched + assert "api-version" not in enriched + assert "key" not in enriched["proxy_server_request"]["body"] + assert "api-version" not in enriched["proxy_server_request"]["body"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "client_metadata", + [{"trace_name": "qa"}, {"trace_name": "qa", "tags": ["client:x"]}, {"trace_name": "qa", "tags": ["usecase:a"]}], +) +async def test_mapped_tags_stay_out_of_requester_metadata(client_metadata): + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.auth.user_api_key_auth import _read_request_body_deferring_parse_failure + from litellm.proxy.proxy_server import ProxyConfig + + body = {"model": "qa-model", "metadata": client_metadata} + + async def receive(): + return {"type": "http.request", "body": json.dumps(body).encode()} + + request = Request( + { + "type": "http", + "method": "POST", + "path": "/chat/completions", + "scheme": "http", + "server": ("localhost", 4000), + "query_string": b"", + "headers": [(b"content-type", b"application/json"), (b"baggage", b"entrypoint=a")], + }, + receive, + ) + data, parse_error = await _read_request_body_deferring_parse_failure(request) + assert parse_error is None + await LiteLLMProxyRequestSetup.apply_baggage_tag_mappings_pre_auth( + request, data, {"request_tag_mappings": [{"key": "entrypoint", "tag_prefix": "usecase:"}]}, None, MagicMock() + ) + enriched = await add_litellm_data_to_request( + data=await read_request_body(request), + request=request, + user_api_key_dict=UserAPIKeyAuth(api_key="qa-key"), + proxy_config=ProxyConfig(), + general_settings={}, + ) + requester_metadata = enriched["metadata"]["requester_metadata"] + assert "usecase:a" in enriched["metadata"]["tags"] + assert requester_metadata["trace_name"] == "qa" + assert requester_metadata.get("tags") == client_metadata.get("tags") + assert BAGGAGE_MAPPED_TAGS_METADATA_KEY not in requester_metadata + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "headers,body_fields", + [ + ([(b"x-litellm-tags", b"usecase:a")], {}), + ([], {"tags": ["usecase:a"]}), + ([], {"litellm_metadata": {"tags": ["usecase:a"]}}), + ], + ids=["header", "body", "litellm_metadata"], +) +async def test_client_sent_tag_still_routes_when_baggage_maps_to_it(headers, body_fields): + import litellm + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.auth.user_api_key_auth import _read_request_body_deferring_parse_failure + from litellm.proxy.proxy_server import ProxyConfig + + body = {"model": "gpt-4", **body_fields} + + async def receive(): + return {"type": "http.request", "body": json.dumps(body).encode()} + + request = Request( + { + "type": "http", + "method": "POST", + "path": "/chat/completions", + "scheme": "http", + "server": ("localhost", 4000), + "query_string": b"", + "headers": [(b"content-type", b"application/json"), (b"baggage", b"entrypoint=a"), *headers], + }, + receive, + ) + data, parse_error = await _read_request_body_deferring_parse_failure(request) + assert parse_error is None + await LiteLLMProxyRequestSetup.apply_baggage_tag_mappings_pre_auth( + request, data, {"request_tag_mappings": [{"key": "entrypoint", "tag_prefix": "usecase:"}]}, None, MagicMock() + ) + enriched = await add_litellm_data_to_request( + data=await read_request_body(request), + request=request, + user_api_key_dict=UserAPIKeyAuth(api_key="qa-key"), + proxy_config=ProxyConfig(), + general_settings={}, + ) + router = litellm.Router( + model_list=[ + { + "model_name": "gpt-4", + "litellm_params": {"model": "gpt-4o", "tags": ["usecase:a"]}, + "model_info": {"id": "usecase-deployment"}, + }, + { + "model_name": "gpt-4", + "litellm_params": {"model": "gpt-4o-mini", "tags": ["default"]}, + "model_info": {"id": "default-deployment"}, + }, + ], + enable_tag_filtering=True, + ) + for _ in range(5): + response = await router.acompletion( + model="gpt-4", + messages=[{"role": "user", "content": "hi"}], + metadata=dict(enriched["metadata"]), + mock_response="hi", + ) + assert response._hidden_params["model_id"] == "usecase-deployment" diff --git a/tests/unit/proxy/management_endpoints/test_tag_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_tag_management_endpoints.py index 1dfababae55..e5d13a381ae 100644 --- a/tests/unit/proxy/management_endpoints/test_tag_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_tag_management_endpoints.py @@ -283,6 +283,66 @@ async def test_new_tag_persists_a_budget(): assert created_tag.budget_id == "budget-1" +@pytest.mark.asyncio +async def test_new_tag_budget_reads_the_row_spend_not_spend_from_before_the_tag_existed(): + """Spend on an unbudgeted request tag must not count against a budget created later. + + The new row starts at the spend it persists, and a restart reloads only that row, so the + counter the budget check reads has to agree with it from the moment the tag is created. + """ + from litellm.proxy.management_endpoints.tag_management_endpoints import new_tag + from litellm.proxy.proxy_server import get_current_spend, increment_spend_counter_cache + + counter_key: Final = "spend:tag:spent-before-budget-tag" + await increment_spend_counter_cache(counter_key=counter_key, increment=5.0) + + created_tag: Final = SimpleNamespace( + tag_name="spent-before-budget-tag", + description=None, + models=[], + created_at=datetime(2024, 1, 1), + updated_at=datetime(2024, 1, 1), + created_by="admin", + ) + + async def create_tag(data, **_): + created_tag.spend = data["spend"] + return created_tag + + mock_db: Final = Mock() + mock_db.litellm_tagtable.find_unique = AsyncMock(return_value=None) + mock_db.litellm_proxymodeltable.find_many = AsyncMock(return_value=[]) + mock_db.litellm_budgettable.create = AsyncMock(return_value=SimpleNamespace(budget_id="budget-1")) + mock_db.litellm_tagtable.create = create_tag + with ( + patch( # test-quality-ok: endpoint resolves the fake database through proxy_server + "litellm.proxy.proxy_server.prisma_client", + SimpleNamespace(db=mock_db, jsonify_object=lambda data: dict(data)), + ), + patch( # test-quality-ok: endpoint reads the audit actor from proxy_server + "litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin" + ), + patch( # test-quality-ok: endpoint requires a router before the budget write + "litellm.proxy.proxy_server.llm_router", object() + ), + patch( # test-quality-ok: auth object cache invalidation is outside this counter contract + "litellm.proxy.management_endpoints.tag_management_endpoints._evict_tag_cache_keys", new=AsyncMock() + ), + ): + await new_tag( + tag=TagNewRequest(name="spent-before-budget-tag", max_budget=25.0), + user_api_key_dict=UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN), + ) + + budget_check_spend: Final = await get_current_spend( + counter_key=counter_key, + fallback_spend=created_tag.spend, + max_budget=25.0, + fallback_authoritative=True, + ) + assert budget_check_spend == created_tag.spend + + @pytest.mark.asyncio @pytest.mark.parametrize( ("budget_fields", "should_update", "expected_max_budget"), diff --git a/tests/unit/proxy/pass_through_endpoints/test_llm_passthrough_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_llm_passthrough_endpoints.py index e658556bd1f..e654c6d444b 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_llm_passthrough_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_llm_passthrough_endpoints.py @@ -264,6 +264,42 @@ async def test_gemini_pass_through_endpoint_forwards_key_and_body( assert outbound.url.params["key"] == "gemini-test-key" assert json.loads(outbound.content) == json.loads(body) +@pytest.mark.asyncio +async def test_router_passthrough_keeps_mapped_baggage_tags_outside_upstream_body(): + from fastapi import Request + + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.common_utils.http_parsing_utils import read_request_body + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup + from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import get_passthrough_router_request_metadata + from litellm.proxy.utils import ProxyLogging + + async def receive(): + return {"type": "http.request", "body": b'{"model":"qa-model"}'} + + request = Request( + { + "type": "http", + "method": "POST", + "path": "/openai/v1/chat/completions", + "query_string": b"", + "headers": [(b"content-type", b"application/json"), (b"baggage", b"entrypoint=a")], + }, + receive, + ) + admitted = await read_request_body(request) + await LiteLLMProxyRequestSetup.apply_baggage_tag_mappings_pre_auth( + request, + admitted, + {"request_tag_mappings": [{"key": "entrypoint", "tag_prefix": "usecase:"}]}, + None, + UserApiKeyCache(), + ) + metadata = get_passthrough_router_request_metadata(UserAPIKeyAuth(api_key="qa-key"), request) + assert metadata["tags"] == ["usecase:a"] + assert await read_request_body(request) == {"model": "qa-model"} + _VERTEX_DEPLOYMENT_PATH: Final = ( "v1/projects/deployment-project/locations/global/publishers/anthropic/models/claude-sonnet-4-5:rawPredict" diff --git a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py index 4c3c7ae0f0a..f0baceecdaa 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -10337,3 +10337,39 @@ def test_init_kwargs_no_authenticated_identity_field_is_client_settable(mock_req metadata: Final = result["litellm_params"]["metadata"] survived: Final = {key: metadata.get(key) for key in expected if metadata.get(key) != expected[key]} assert survived == {}, f"client-supplied values survived for: {sorted(survived)}" + +@pytest.mark.asyncio +async def test_direct_passthrough_keeps_mapped_baggage_tags_outside_upstream_body(): + from litellm.proxy.common_utils.http_parsing_utils import read_request_body + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup + from litellm.proxy.utils import ProxyLogging + + async def receive(): + return {"type": "http.request", "body": b'{"model":"qa-model"}'} + + request = Request( + { + "type": "http", + "method": "POST", + "path": "/openai/v1/chat/completions", + "query_string": b"", + "headers": [ + (b"content-type", b"application/json"), + (b"baggage", b"entrypoint=a"), + (b"x-litellm-tags", b"existing"), + ], + }, + receive, + ) + admitted = await read_request_body(request) + await LiteLLMProxyRequestSetup.apply_baggage_tag_mappings_pre_auth( + request, + admitted, + {"request_tag_mappings": [{"key": "entrypoint", "tag_prefix": "usecase:"}]}, + None, + UserApiKeyCache(), + ) + metadata = _update_metadata_with_tags_in_header(request, {}) + assert metadata["tags"] == ["existing", "usecase:a"] + assert await read_request_body(request) == {"model": "qa-model"} diff --git a/tests/unit/proxy/test_litellm_pre_call_utils.py b/tests/unit/proxy/test_litellm_pre_call_utils.py index 4647952312a..4a97106b3ac 100644 --- a/tests/unit/proxy/test_litellm_pre_call_utils.py +++ b/tests/unit/proxy/test_litellm_pre_call_utils.py @@ -4756,6 +4756,43 @@ async def test_add_guardrails_from_policy_engine(): attachment_registry._initialized = False +@pytest.mark.asyncio +@pytest.mark.parametrize( + "caller_tags,inherited_tags,expected_guardrails", + [((), (), []), (("usecase:a",), (), ["usecase_guard"]), ((), ("usecase:a",), ["usecase_guard"])], +) +async def test_baggage_mapped_tags_match_tag_policies_only_when_the_caller_or_key_asked_for_them( + caller_tags, inherited_tags, expected_guardrails +): + from typing import cast + + from litellm.proxy.policy_engine.attachment_registry import get_attachment_registry + from litellm.proxy.policy_engine.policy_registry import get_policy_registry + from litellm.types.proxy.policy_engine import Policy, PolicyAttachment, PolicyGuardrails + + request = cast(Request, SimpleNamespace(state=SimpleNamespace(litellm_baggage_mapped_tags=("usecase:a",)))) + data = {"model": "gpt-4", "metadata": {"tags": [], "inherited_tags": inherited_tags}} + LiteLLMProxyRequestSetup.apply_authenticated_baggage_tags(request, data) + data["metadata"]["caller_tags"] = caller_tags + policy_registry = get_policy_registry() + attachment_registry = get_attachment_registry() + policy_registry._policies = {"usecase": Policy(guardrails=PolicyGuardrails(add=["usecase_guard"]))} + policy_registry._initialized = True + attachment_registry._attachments = [PolicyAttachment(policy="usecase", tags=["usecase:a"])] + attachment_registry._initialized = True + try: + await add_guardrails_from_policy_engine( + data=data, metadata_variable_name="metadata", user_api_key_dict=UserAPIKeyAuth(api_key="test-key") + ) + finally: + policy_registry._policies = {} + policy_registry._initialized = False + attachment_registry._attachments = [] + attachment_registry._initialized = False + assert "usecase:a" in data["metadata"]["tags"] + assert data["metadata"].get("guardrails", []) == expected_guardrails + + def test_match_and_track_policies_preserves_attachment_and_request_body_order(): from litellm.proxy.policy_engine.attachment_registry import AttachmentRegistry from litellm.types.proxy.policy_engine import Policy, PolicyMatchContext @@ -10358,3 +10395,112 @@ async def test_per_user_token_resolution_uses_key_identity_without_header_mappin called_user_ids = [call.kwargs["user_id"] for call in load.await_args_list] assert called_user_ids == ["user-a"] assert updated["secret_fields"]["user_provider_credentials"]["copilot-cred"] == "gho_user_a_token" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("route", ["/v1/images/generations", "/v1/rerank"]) +async def test_raw_body_inference_keeps_authenticated_baggage_tags_for_spend(route): + from litellm.proxy.common_utils.http_parsing_utils import read_request_body + from litellm.proxy.proxy_server import ProxyConfig + + async def receive(): + return {"type": "http.request", "body": b'{"model":"qa-model","metadata":{"tags":["existing"]}}'} + + request = Request( + { + "type": "http", + "method": "POST", + "path": route, + "scheme": "http", + "server": ("localhost", 4000), + "query_string": b"", + "headers": [(b"content-type", b"application/json"), (b"baggage", b"entrypoint=a")], + }, + receive, + ) + admitted = await read_request_body(request) + await LiteLLMProxyRequestSetup.apply_baggage_tag_mappings_pre_auth( + request, + admitted, + {"request_tag_mappings": [{"key": "entrypoint", "tag_prefix": "usecase:"}]}, + None, + MagicMock(), + ) + reparsed = json.loads(await request.body()) + enriched = await add_litellm_data_to_request( + data=reparsed, + request=request, + user_api_key_dict=UserAPIKeyAuth(api_key="qa-key"), + proxy_config=ProxyConfig(), + general_settings={}, + ) + assert "usecase:a" in enriched["metadata"]["tags"] + assert "existing" in enriched["metadata"]["tags"] + assert "usecase:a" not in enriched["metadata"]["inherited_tags"] + assert enriched["metadata"]["tags"].count("usecase:a") == 1 + + +async def _string_tags_outcome(baggage_headers: tuple[tuple[bytes, bytes], ...]) -> str: + from litellm.proxy.common_utils.http_parsing_utils import read_request_body + from litellm.proxy.proxy_server import ProxyConfig + + async def receive(): + return {"type": "http.request", "body": b'{"model":"qa-model","metadata":{"tags":"client-tag"}}'} + + request: Final = Request( + { + "type": "http", + "method": "POST", + "path": "/v1/chat/completions", + "scheme": "http", + "server": ("localhost", 4000), + "query_string": b"", + "headers": [(b"content-type", b"application/json"), *baggage_headers], + }, + receive, + ) + admitted: Final = await read_request_body(request) + await LiteLLMProxyRequestSetup.apply_baggage_tag_mappings_pre_auth( + request, + admitted, + {"request_tag_mappings": [{"key": "tenant", "tag_prefix": "tenant:"}]}, + None, + MagicMock(), + ) + try: + enriched: Final = await add_litellm_data_to_request( + data=json.loads(await request.body()), + request=request, + user_api_key_dict=UserAPIKeyAuth(api_key="qa-key"), + proxy_config=ProxyConfig(), + general_settings={}, + ) + except Exception as error: + return type(error).__name__ + return repr(enriched["metadata"].get("tags")) + + +@pytest.mark.asyncio +async def test_string_metadata_tags_get_the_same_outcome_with_or_without_mapped_baggage(): + without_baggage: Final = await _string_tags_outcome(()) + with_baggage: Final = await _string_tags_outcome(((b"baggage", b"tenant=a"),)) + assert with_baggage == without_baggage + + +def test_applying_mapped_baggage_tags_again_leaves_the_metadata_unchanged(): + request: Final = Request({"type": "http", "headers": []}) + request.state.litellm_baggage_mapped_tags = ("tenant:a",) + once: Final[dict[str, object]] = {"metadata": {"tags": ["client-tag"]}} + twice: Final[dict[str, object]] = copy.deepcopy(once) + LiteLLMProxyRequestSetup.apply_authenticated_baggage_tags(request, once) + LiteLLMProxyRequestSetup.apply_authenticated_baggage_tags(request, twice) + LiteLLMProxyRequestSetup.apply_authenticated_baggage_tags(request, twice) + assert twice == once + + +def test_mapped_baggage_tags_keep_non_string_client_tags_as_sent(): + request: Final = Request({"type": "http", "headers": []}) + request.state.litellm_baggage_mapped_tags = ("tenant:a",) + data: Final[dict[str, object]] = {"metadata": {"tags": [1, {"team": "x"}, "client-tag", "client-tag"]}} + LiteLLMProxyRequestSetup.apply_authenticated_baggage_tags(request, data) + assert data["metadata"]["tags"] == [1, {"team": "x"}, "client-tag", "client-tag", "tenant:a"] diff --git a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py index 0dbf3451c3b..7eb3f557971 100644 --- a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py +++ b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py @@ -8311,6 +8311,37 @@ async def test_update_general_settings_ignores_an_invalid_user_api_key_cache_max assert cache.get_cache(key="key-0", local_only=True) == {"token": "key-0"} +@pytest.mark.asyncio +async def test_update_general_settings_serves_an_invalid_db_request_tag_mappings_like_unset_until_fixed(monkeypatch): + from fastapi import Request + + from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup + from litellm.proxy.proxy_server import ProxyConfig + + proxy_config = ProxyConfig() + monkeypatch.setattr(proxy_server_module, "general_settings", proxy_config.settings) + + async def mapped_tags() -> list[str]: + data: Final[dict[str, object]] = {} + await LiteLLMProxyRequestSetup.apply_baggage_tag_mappings_pre_auth( + Request({"type": "http", "headers": [(b"baggage", b"tenant=a")]}), + data, + proxy_server_module.general_settings, + None, + MagicMock(), + ) + metadata: Final = data.get("metadata") + return list(metadata.get("tags", [])) if isinstance(metadata, dict) else [] + + await proxy_config._update_general_settings(db_general_settings={"request_tag_mappings": [{"key": "tenant"}]}) + assert await mapped_tags() == [] + + await proxy_config._update_general_settings( + db_general_settings={"request_tag_mappings": [{"key": "tenant", "tag_prefix": "tenant:", "required": True}]} + ) + assert await mapped_tags() == ["tenant:a"] + + @pytest.mark.asyncio async def test_update_general_settings_user_api_key_cache_max_size_yaml_wins(monkeypatch): """A DB value must not silently override an explicit YAML capacity on reload.""" diff --git a/tests/unit/router_strategy/test_router_tag_routing.py b/tests/unit/router_strategy/test_router_tag_routing.py index f0e85208933..097a34cf2e6 100644 --- a/tests/unit/router_strategy/test_router_tag_routing.py +++ b/tests/unit/router_strategy/test_router_tag_routing.py @@ -3175,6 +3175,112 @@ async def test_chat_request_carrying_litellm_metadata_still_routes_on_proxy_merg assert deployment["model_info"]["id"] == "team-a-deployment" +def _baggage_mapped_metadata(client_tags: list[str], inherited_tags: list[str]) -> dict: + from types import SimpleNamespace + from typing import cast + + from fastapi import Request + + from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup + + request = cast(Request, SimpleNamespace(state=SimpleNamespace(litellm_baggage_mapped_tags=("usecase:a",)))) + data = {"metadata": {"tags": client_tags, "inherited_tags": inherited_tags}} + LiteLLMProxyRequestSetup.apply_authenticated_baggage_tags(request, data) + return data["metadata"] + + +@pytest.mark.asyncio() +@pytest.mark.parametrize( + "client_tags,inherited_tags,expected_model_id", + [ + ([], [], "default-deployment"), + (["usecase:a"], [], "usecase-deployment"), + ([], ["usecase:a"], "usecase-deployment"), + ], +) +async def test_baggage_mapped_tags_classify_spend_without_selecting_deployments( + client_tags, inherited_tags, expected_model_id +): + router = litellm.Router( + model_list=[ + { + "model_name": "gpt-4", + "litellm_params": {"model": "gpt-4o", "tags": ["usecase:a"]}, + "model_info": {"id": "usecase-deployment"}, + }, + { + "model_name": "gpt-4", + "litellm_params": {"model": "gpt-4o-mini", "tags": ["default"]}, + "model_info": {"id": "default-deployment"}, + }, + ], + enable_tag_filtering=True, + ) + metadata = _baggage_mapped_metadata(client_tags, inherited_tags) + assert "usecase:a" in metadata["tags"] + for _ in range(5): + response = await router.acompletion( + model="gpt-4", + messages=[{"role": "user", "content": "hi"}], + metadata=dict(metadata), + mock_response="hi", + ) + assert response._hidden_params["model_id"] == expected_model_id + + +@pytest.mark.asyncio() +@pytest.mark.parametrize( + "client_tags,inherited_tags,expected_model_id", + [ + ([], [], "tier-b"), + (["usecase:a"], [], "tier-a"), + ([], ["usecase:a"], "tier-a"), + ], +) +async def test_baggage_mapped_tags_never_select_a_tagged_pre_routing_strategy( + client_tags, inherited_tags, expected_model_id +): + from litellm.types.router import TaggedPreRoutingStrategy + + router = litellm.Router( + model_list=[ + { + "model_name": "smart", + "litellm_params": {"model": "openai/gpt-4o"}, + "model_info": {"id": "plain-smart"}, + }, + { + "model_name": "tier-a", + "litellm_params": {"model": "openai/gpt-4o"}, + "model_info": {"id": "tier-a"}, + }, + { + "model_name": "tier-b", + "litellm_params": {"model": "openai/gpt-4o-mini"}, + "model_info": {"id": "tier-b"}, + }, + ], + num_retries=0, + ) + router.auto_routers = { + "smart": [ + TaggedPreRoutingStrategy(tags=("usecase:b",), strategy=_RewriteToTierStrategy("tier-b")), + TaggedPreRoutingStrategy(tags=("usecase:a",), strategy=_RewriteToTierStrategy("tier-a")), + ] + } + metadata = _baggage_mapped_metadata(client_tags, inherited_tags) + assert "usecase:a" in metadata["tags"] + + response = await router.acompletion( + model="smart", + messages=[{"role": "user", "content": "hi"}], + metadata=dict(metadata), + mock_response="hi", + ) + + assert response._hidden_params["model_id"] == expected_model_id + + def _tag_router(model_list: list[dict[str, object]], tag_routing_prefix: str = "route:") -> litellm.Router: return litellm.Router(model_list=model_list, enable_tag_filtering=True, tag_routing_prefix=tag_routing_prefix) diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index cc25d73b83a..4a9dd199c08 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -9673,6 +9673,26 @@ def test_update_kwargs_with_deployment_propagates_model_tags(): assert "production" in kwargs["metadata"]["tags"] +def test_update_kwargs_with_deployment_merges_many_client_tags_in_linear_time(): + router: Final = litellm.Router( + model_list=[ + { + "model_name": "gpt-4o-mini", + "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "fake-key", "tags": ["deployment-tag"]}, + }, + ], + ) + client_tags: Final = [f"client-{index}" for index in range(100_000)] + kwargs: Final = {"metadata": {"tags": client_tags}} + deployment: Final = router.get_deployment_by_model_group_name(model_group_name="gpt-4o-mini") + + started: Final = time.perf_counter() + router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) + + assert time.perf_counter() - started < 2 + assert kwargs["metadata"]["tags"] == [*client_tags, "deployment-tag"] + + def test_update_kwargs_with_deployment_merges_tags_without_duplicates(): """ Test that when both request-level and deployment-level tags exist, diff --git a/ui/litellm-dashboard/src/app/(dashboard)/router-settings/_components/RequestTagMappings.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/router-settings/_components/RequestTagMappings.test.tsx new file mode 100644 index 00000000000..b7ecd6bbb64 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/router-settings/_components/RequestTagMappings.test.tsx @@ -0,0 +1,145 @@ +import { fireEvent, render, screen } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; +import { useState } from "react"; +import { describe, expect, it, vi } from "vitest"; + +import { RequestTagMapping, RequestTagMappingsEditor, requestTagMappingsAreValid } from "./RequestTagMappings"; + +const exact: RequestTagMapping = { header: "baggage", key: "entrypoint", value: "a", tag: "support" }; +const variable: RequestTagMapping = { header: "baggage", key: "entrypoint", tag_prefix: "usecase:" }; + +function Editor({ + initial, + changed, + saved = vi.fn(), +}: { + initial: RequestTagMapping[]; + changed: (rules: RequestTagMapping[]) => void; + saved?: (rules: RequestTagMapping[]) => void; +}) { + const [rules, setRules] = useState(initial); + return ( + <> + { + setRules(next); + changed(next); + }} + /> + + + ); +} + +describe("RequestTagMappingsEditor", () => { + it("shows YAML-owned rules without allowing edits", () => { + render(); + expect(screen.getByLabelText("Mapping 1 match value")).toBeDisabled(); + expect(screen.getByRole("button", { name: "Add mapping" })).toBeDisabled(); + expect(screen.getByText(/Managed in config.yaml/)).toBeInTheDocument(); + }); + it("shows one-by-one mappings alongside variable prefix mappings", () => { + render(); + expect(screen.getByRole("combobox", { name: "Mapping 1 mode" })).toHaveTextContent("Exact value"); + expect(screen.getByRole("combobox", { name: "Mapping 2 mode" })).toHaveTextContent("Any value"); + expect(screen.getByLabelText("Mapping 1 match value")).toHaveValue("a"); + expect(screen.getByLabelText("Mapping 1 tag")).toHaveValue("support"); + expect(screen.getByLabelText("Mapping 2 tag prefix")).toHaveValue("usecase:"); + expect(screen.getByText("entrypoint=a → support")).toBeInTheDocument(); + expect(screen.getByText("entrypoint=x → usecase:x")).toBeInTheDocument(); + }); + + it("adds and edits a new exact value mapping", async () => { + const changed = vi.fn(); + render(); + await userEvent.click(screen.getByRole("button", { name: "Add mapping" })); + fireEvent.change(screen.getByLabelText("Mapping 1 baggage key"), { target: { value: "entrypoint" } }); + fireEvent.change(screen.getByLabelText("Mapping 1 match value"), { target: { value: "b" } }); + fireEvent.change(screen.getByLabelText("Mapping 1 tag"), { target: { value: "sales" } }); + expect(screen.getByText("entrypoint=b → sales")).toBeInTheDocument(); + expect(changed).toHaveBeenLastCalledWith([{ header: "baggage", key: "entrypoint", value: "b", tag: "sales" }]); + }); + + it("changes an exact mapping to any-value mode and removes its exact match fields", async () => { + const changed = vi.fn(); + render(); + await userEvent.click(screen.getByRole("combobox", { name: "Mapping 1 mode" })); + await userEvent.click(await screen.findByRole("option", { name: "Any value" })); + expect(screen.queryByLabelText("Mapping 1 match value")).not.toBeInTheDocument(); + expect(screen.getByLabelText("Mapping 1 tag prefix")).toHaveValue(""); + expect(screen.getByLabelText("Mapping 1 tag prefix")).toHaveAttribute("placeholder", "usecase:"); + expect(changed).toHaveBeenLastCalledWith([ + expect.objectContaining({ value: undefined, tag: undefined, tag_prefix: "" }), + ]); + }); + + it("blocks saving a new mapping until its match value and tag are filled in", async () => { + const saved = vi.fn(); + render(); + await userEvent.click(screen.getByRole("button", { name: "Add mapping" })); + expect(screen.getByLabelText("Mapping 1 baggage key")).toHaveValue(""); + expect(screen.getByLabelText("Mapping 1 baggage key")).toHaveAttribute("placeholder", "entrypoint"); + expect(screen.getByText("Enter a baggage key")).toBeInTheDocument(); + fireEvent.change(screen.getByLabelText("Mapping 1 baggage key"), { target: { value: "entrypoint" } }); + expect(screen.getByText("Enter a match value")).toBeInTheDocument(); + await userEvent.click(screen.getByRole("button", { name: "Update" })); + fireEvent.change(screen.getByLabelText("Mapping 1 match value"), { target: { value: "b" } }); + expect(screen.getByText("Enter a tag")).toBeInTheDocument(); + await userEvent.click(screen.getByRole("button", { name: "Update" })); + expect(saved).not.toHaveBeenCalled(); + fireEvent.change(screen.getByLabelText("Mapping 1 tag"), { target: { value: "sales" } }); + expect(screen.queryByText(/^Enter a/)).not.toBeInTheDocument(); + await userEvent.click(screen.getByRole("button", { name: "Update" })); + expect(saved).toHaveBeenCalledWith([{ header: "baggage", key: "entrypoint", value: "b", tag: "sales" }]); + }); + + it("blocks saving an any-value mapping until the admin types a tag prefix", async () => { + const saved = vi.fn(); + render(); + await userEvent.click(screen.getByRole("combobox", { name: "Mapping 1 mode" })); + await userEvent.click(await screen.findByRole("option", { name: "Any value" })); + expect(screen.getByText("Enter a tag prefix")).toBeInTheDocument(); + expect(screen.getByRole("button", { name: "Update" })).toBeDisabled(); + fireEvent.change(screen.getByLabelText("Mapping 1 tag prefix"), { target: { value: "team:" } }); + await userEvent.click(screen.getByRole("button", { name: "Update" })); + expect(saved).toHaveBeenCalledWith([expect.objectContaining({ key: "entrypoint", tag_prefix: "team:" })]); + }); + + it("blocks saving a baggage key that is empty or not a header token", async () => { + render(); + fireEvent.change(screen.getByLabelText("Mapping 1 baggage key"), { target: { value: "" } }); + expect(screen.getByText("Enter a baggage key")).toBeInTheDocument(); + expect(screen.getByRole("button", { name: "Update" })).toBeDisabled(); + fireEvent.change(screen.getByLabelText("Mapping 1 baggage key"), { target: { value: "entry point" } }); + expect(screen.getByText(/Use only letters, digits/)).toBeInTheDocument(); + expect(screen.getByRole("button", { name: "Update" })).toBeDisabled(); + fireEvent.change(screen.getByLabelText("Mapping 1 baggage key"), { target: { value: "entry.point" } }); + expect(screen.getByRole("button", { name: "Update" })).toBeEnabled(); + }); + + it("blocks saving a second mapping for the same key and value until it differs", () => { + render(); + expect(screen.getByText("Another mapping already matches this key and value")).toBeInTheDocument(); + expect(screen.getByRole("button", { name: "Update" })).toBeDisabled(); + fireEvent.change(screen.getByLabelText("Mapping 2 match value"), { target: { value: "b" } }); + expect(screen.queryByText("Another mapping already matches this key and value")).not.toBeInTheDocument(); + expect(screen.getByRole("button", { name: "Update" })).toBeEnabled(); + }); + + it("blocks saving a second any-value mapping for the same key", () => { + render(); + expect(screen.getByText("Another mapping already matches this key and value")).toBeInTheDocument(); + expect(screen.getByRole("button", { name: "Update" })).toBeDisabled(); + }); + + it("deletes one rule while keeping the other mode intact", async () => { + const changed = vi.fn(); + render(); + await userEvent.click(screen.getByRole("button", { name: "Delete mapping 1" })); + expect(screen.getByLabelText("Mapping 1 tag prefix")).toHaveValue("usecase:"); + expect(changed).toHaveBeenLastCalledWith([variable]); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/router-settings/_components/RequestTagMappings.tsx b/ui/litellm-dashboard/src/app/(dashboard)/router-settings/_components/RequestTagMappings.tsx new file mode 100644 index 00000000000..d1101eaaf27 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/router-settings/_components/RequestTagMappings.tsx @@ -0,0 +1,189 @@ +import { Button } from "@/components/ui/button"; +import { Input } from "@/components/ui/input"; +import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select"; +import { Switch } from "@/components/ui/switch"; +import { Trash2 } from "lucide-react"; + +export interface RequestTagMapping { + header: "baggage"; + key: string; + value?: string | null; + tag?: string | null; + tag_prefix?: string | null; + required?: boolean; + require_budgeted_tag?: boolean; +} + +interface Props { + value: unknown; + disabled?: boolean; + onChange: (rules: RequestTagMapping[]) => void; +} + +const isMapping = (value: unknown): value is RequestTagMapping => { + if (typeof value !== "object" || value === null) return false; + return "key" in value && typeof value.key === "string"; +}; + +const MODES = [ + { value: "exact", label: "Exact value" }, + { value: "any", label: "Any value" }, +]; + +const BAGGAGE_KEY_PATTERN = /^[!#$%&'*+\-.^_`|~0-9A-Za-z]+$/; + +const isExact = (rule: RequestTagMapping) => rule.value !== undefined && rule.value !== null; + +const mappingProblem = (rule: RequestTagMapping): string | null => { + if (!rule.key) return "Enter a baggage key"; + if (!BAGGAGE_KEY_PATTERN.test(rule.key)) return "Use only letters, digits, and !#$%&'*+-.^_`|~ in the baggage key"; + if (isExact(rule)) { + if (!rule.value) return "Enter a match value"; + return rule.tag ? null : "Enter a tag"; + } + return rule.tag_prefix ? null : "Enter a tag prefix"; +}; + +const matchIdentity = (rule: RequestTagMapping) => JSON.stringify([rule.key, isExact(rule) ? rule.value : null]); + +const problemAt = (rules: RequestTagMapping[], index: number): string | null => { + const identity = matchIdentity(rules[index]); + const duplicated = rules.slice(0, index).some((earlier) => matchIdentity(earlier) === identity); + return mappingProblem(rules[index]) ?? (duplicated ? "Another mapping already matches this key and value" : null); +}; + +export const requestTagMappingsAreValid = (value: unknown) => + Array.isArray(value) && value.every(isMapping) && value.every((_, index) => problemAt(value, index) === null); + +export function RequestTagMappingsEditor({ value, onChange, disabled = false }: Props) { + const rules = Array.isArray(value) ? value.filter(isMapping) : []; + const update = (index: number, patch: Partial) => + onChange(rules.map((rule, position) => (position === index ? { ...rule, ...patch } : rule))); + + return ( +
+

+ Exact value mappings take priority over any-value rules for the same key. +

+ {rules.map((rule, index) => { + const exact = isExact(rule); + const problem = problemAt(rules, index); + const label = `Mapping ${index + 1}`; + return ( +
+ {label} +
+ + +
+ + {exact ? ( + <> + + +

+ {rule.key}={rule.value || "value"} → {rule.tag || "tag"} +

+ + ) : ( + <> + +

+ {rule.key}=x → {rule.tag_prefix}x +

+ + )} + {problem ?

{problem}

: null} + + +
+ ); + })} + {disabled ? ( +

Managed in config.yaml. Edit the YAML to change these rules.

+ ) : null} + +
+ ); +} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/router-settings/_components/general_settings.integration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/router-settings/_components/general_settings.integration.test.tsx index cf5cad564ea..35434830d98 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/router-settings/_components/general_settings.integration.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/router-settings/_components/general_settings.integration.test.tsx @@ -253,3 +253,33 @@ it("should delete only the Default setting and retain explicit false and zero", ["token", "synthetic_count", 0], ]); }); + +it("keeps a new baggage mapping unsaved until every field the backend requires is filled in", async () => { + vi.mocked(getGeneralSettingsCall).mockResolvedValue([ + { + field_name: "request_tag_mappings", + field_type: "List", + field_value: [], + field_description: "baggage mappings", + stored_in_db: true, + editable: true, + }, + ]); + vi.mocked(updateConfigFieldSetting).mockClear(); + const user = userEvent.setup(); + renderWithProviders(); + await user.click(screen.getByRole("tab", { name: "General" })); + await user.click(await screen.findByRole("button", { name: "Add mapping" })); + const update = within(screen.getByRole("row", { name: /request_tag_mappings/ })).getByRole("button", { + name: "Update", + }); + expect(update).toBeDisabled(); + fireEvent.change(screen.getByLabelText("Mapping 1 baggage key"), { target: { value: "entrypoint" } }); + fireEvent.change(screen.getByLabelText("Mapping 1 match value"), { target: { value: "a" } }); + expect(update).toBeDisabled(); + fireEvent.change(screen.getByLabelText("Mapping 1 tag"), { target: { value: "usecase:support" } }); + await user.click(update); + expect(vi.mocked(updateConfigFieldSetting).mock.calls).toEqual([ + ["token", "request_tag_mappings", [{ header: "baggage", key: "entrypoint", value: "a", tag: "usecase:support" }]], + ]); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/router-settings/_components/general_settings.tsx b/ui/litellm-dashboard/src/app/(dashboard)/router-settings/_components/general_settings.tsx index 3d0a2e438ac..5d28adb7ef3 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/router-settings/_components/general_settings.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/router-settings/_components/general_settings.tsx @@ -14,6 +14,7 @@ import { StatusBadge } from "@/components/shared/table_cells"; import RouterSettings from "@/components/router_settings"; import Fallbacks from "@/components/Settings/RouterSettings/Fallbacks/Fallbacks"; import RoutingGroups from "@/components/routing_groups"; +import { RequestTagMappingsEditor, requestTagMappingsAreValid } from "./RequestTagMappings"; const PROMPT_CACHING_TAB = "prompt_caching"; const ENABLE_ANTHROPIC_PROMPT_CACHING = "enable_anthropic_prompt_caching"; @@ -34,6 +35,7 @@ export interface generalSettingsItem { field_value: any; field_description: string; stored_in_db: boolean | null; + editable?: boolean; field_options?: string[] | null; field_tab?: string | null; field_default_value?: any; @@ -57,6 +59,15 @@ const SettingValueEditor: React.FC<{ setting: generalSettingsItem; onChange: (fieldName: string, newValue: any) => void; }> = ({ setting, onChange }) => { + if (setting.field_name === "request_tag_mappings") { + return ( + onChange(setting.field_name, rules)} + /> + ); + } if (setting.field_type === "Integer") { return ( = ({ accessToken, user setGeneralSettings(updatedSettings); }; - const handleUpdateField = (fieldName: string) => { + const handleUpdateField = async (fieldName: string) => { if (!accessToken) { return; } @@ -253,7 +264,7 @@ const GeneralSettings: React.FC = ({ accessToken, user return; } try { - updateConfigFieldSetting(accessToken, fieldName, fieldValue); + await updateConfigFieldSetting(accessToken, fieldName, fieldValue); // update value in state const updatedSettings = generalSettings.map((setting) => @@ -354,13 +365,23 @@ const GeneralSettings: React.FC = ({ accessToken, user )} - - handleResetField(value.field_name)} - className="inline-flex shrink-0 cursor-pointer items-center justify-center px-1.5 py-1.5 text-destructive" + + {value.field_name === "request_tag_mappings" && value.editable === false ? null : ( + handleResetField(value.field_name)} + className="inline-flex shrink-0 cursor-pointer items-center justify-center px-1.5 py-1.5 text-destructive" + > + + + )} ))} diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index bf2fe4300a3..992a7d91b0d 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -30199,6 +30199,11 @@ export interface components { * @description When set to True, rejects requests that contain client-side 'metadata.tags' to prevent users from influencing budgets by sending different tags. Tags can only be inherited from the API key metadata. */ reject_clientside_metadata_tags?: boolean | null; + /** + * Request Tag Mappings + * @description Map W3C baggage values to request tags. Exact value/tag rules take precedence over tag_prefix fallbacks. + */ + request_tag_mappings?: components["schemas"]["RequestTagMapping"][] | null; /** * Responses Websocket Session Limit Seconds * @description Maximum lifetime in seconds of a Responses API WebSocket session, measured from connection accept and covering the idle wait for the first response.create frame. Defaults to 3600, matching OpenAI's documented 60-minute WebSocket connection limit. Must be between 60 and 7200 seconds. @@ -42655,6 +42660,36 @@ export interface components { /** Status Code */ status_code: number; }; + /** + * RequestTagMapping + * @description An exact baggage value mapping, or a prefix for any value of a key. + */ + RequestTagMapping: { + /** + * Header + * @default baggage + * @constant + */ + header: "baggage"; + /** Key */ + key: string; + /** + * Require Budgeted Tag + * @default false + */ + require_budgeted_tag: boolean; + /** + * Required + * @default false + */ + required: boolean; + /** Tag */ + tag?: string | null; + /** Tag Prefix */ + tag_prefix?: string | null; + /** Value */ + value?: string | null; + }; /** * RequestType * @description Fixed v0 taxonomy. User-extensible types come in v1.