mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-27 06:13:07 +00:00
feat(proxy): map baggage values to budget tags (#47358)
* feat(proxy): map baggage values to budget tags * fix(proxy): preserve mapped tags in the parsed request cache * fix(proxy): preserve mapper accounting and tag policy * fix(proxy): preserve baggage tags across native transports * fix(proxy): keep mapper rejection and passthrough logging safe * fix(proxy): scope baggage tag mappings to billed inference and keep them out of tag routing Mapped tags now apply only to billed LLM API routes, skipping token counting, model listing, MCP and A2A. They classify spend and budgets but never select deployments unless a key, team or client set the same tag. WebSocket inference returns 403 only while a required rule exists. A tag created with /tag/new resets its spend counter so its budget counts from creation. The General settings editor keeps a new mapping unsaved until every required field is filled in. * fix(proxy): keep client-tag rejection identical with baggage, map A2A calls, skip request previews * fix(proxy): keep mapped tags out of requester metadata, answer 503 on tag lookup errors, map custom pass-through routes, check the key before the WebSocket 403 * fix(proxy): keep caller-sent tags routable, keep mapped tags out of policy and guardrail matching, honor allow_requests_on_db_unavailable, reject control characters in baggage values * test(proxy): cover baggage tag mappings end to end Adds integration cells for request_tag_mappings across chat, messages, responses, A2A and custom pass-through spend, required and budgeted rules, the WebSocket refusal, unbilled routes, reject_clientside_metadata_tags, tag routing and policy isolation, a dropped tag read with and without allow_requests_on_db_unavailable, and the /tag/new spend counter reset * fix(proxy): keep base tag merges for requests that map nothing, accept JSON-string metadata tags, count litellm_metadata caller tags Requests with no mapped baggage now reuse the cached body exactly as before, so header and key tag merges still normalize metadata for routing. The client body copy is made only when a tag maps, and the base merges are mirrored onto it. A request that maps nothing no longer gets a 400 for non-object metadata. caller_tags now counts tags sent in both metadata and litellm_metadata. Optional mappings no longer reject JSON-string metadata tags under reject_clientside_metadata_tags. * fix(router): keep baggage-mapped tags from selecting tag-scoped pre-routing strategies Auto, complexity, adaptive and quality router selection, the consumed-tags stamp and compression policy lookup now read request tags without baggage-mapped tags, unless the caller sent the tag or the key, team or project carries it, matching deployment tag filtering. * feat(proxy): parse baggage headers into request tag mappings * fix(proxy): match plus-sign baggage keys and keep deployment tags in guardrail modes * fix(proxy): keep baggage-mapped tags behind key, team and deployment tags and serve non-object or invalid DB mapping settings like unset Mapped tags are added after the key, team, deployment and caller tags, so a tag-mode guardrail picks the same mode with or without baggage. A non-object metadata field maps like empty metadata, and an invalid request_tag_mappings value loaded from the DB is ignored with an ERROR log instead of failing every request * fix(router): merge deployment tags in linear time and keep malformed client tags as sent under mapped baggage Merge serving deployment tags into the request tags in one pass, so a large client tag list no longer blocks the event loop. Leave a non-list metadata.tags untouched when baggage maps a tag, so the request gets the same answer as without baggage. Add the C17 integration cell for the team tag's guardrail mode under a mapped deployment tag. * test(spend): cover string metadata.tags under mapped baggage * fix(proxy): check mapped tag budgets when metadata.tags is a non-list or holds non-strings Mapped baggage tags are applied again after the header and key tag merges, so a header or key tag that replaces a string metadata.tags no longer hides the mapped tag from the tag budget check. An empty or false metadata.tags is read like a missing one, and a tags list with non-string items is kept as sent with the mapped tags added after it, instead of failing auth with a 401
This commit is contained in:
parent
0ca29c7d84
commit
dfac770e2b
38 changed files with 3267 additions and 131 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
67
litellm/proxy/common_utils/request_tag_mappings.py
Normal file
67
litellm/proxy/common_utils/request_tag_mappings.py
Normal file
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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),
|
||||
),
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -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 ()
|
||||
)
|
||||
|
|
|
|||
27
litellm/types/proxy/request_tag_mappings.py
Normal file
27
litellm/types/proxy/request_tag_mappings.py
Normal file
|
|
@ -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
|
||||
|
|
@ -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."""
|
||||
|
|
|
|||
847
tests/integration/spend/test_baggage_tag_mappings.py
Normal file
847
tests/integration/spend/test_baggage_tag_mappings.py
Normal file
|
|
@ -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
|
||||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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) == []
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
473
tests/unit/proxy/common_utils/test_request_tag_mappings.py
Normal file
473
tests/unit/proxy/common_utils/test_request_tag_mappings.py
Normal file
|
|
@ -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"
|
||||
|
|
@ -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"),
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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"}
|
||||
|
|
|
|||
|
|
@ -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"]
|
||||
|
|
|
|||
|
|
@ -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."""
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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 (
|
||||
<>
|
||||
<RequestTagMappingsEditor
|
||||
value={rules}
|
||||
onChange={(next) => {
|
||||
setRules(next);
|
||||
changed(next);
|
||||
}}
|
||||
/>
|
||||
<button type="button" disabled={!requestTagMappingsAreValid(rules)} onClick={() => saved(rules)}>
|
||||
Update
|
||||
</button>
|
||||
</>
|
||||
);
|
||||
}
|
||||
|
||||
describe("RequestTagMappingsEditor", () => {
|
||||
it("shows YAML-owned rules without allowing edits", () => {
|
||||
render(<RequestTagMappingsEditor value={[exact]} onChange={vi.fn()} disabled />);
|
||||
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(<Editor initial={[exact, variable]} changed={vi.fn()} />);
|
||||
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(<Editor initial={[]} changed={changed} />);
|
||||
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(<Editor initial={[exact]} changed={changed} />);
|
||||
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(<Editor initial={[]} changed={vi.fn()} saved={saved} />);
|
||||
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(<Editor initial={[exact]} changed={vi.fn()} saved={saved} />);
|
||||
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(<Editor initial={[exact]} changed={vi.fn()} />);
|
||||
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(<Editor initial={[exact, { ...exact, tag: "sales" }, variable]} changed={vi.fn()} />);
|
||||
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(<Editor initial={[variable, { ...variable, tag_prefix: "team:" }]} changed={vi.fn()} />);
|
||||
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(<Editor initial={[exact, variable]} changed={changed} />);
|
||||
await userEvent.click(screen.getByRole("button", { name: "Delete mapping 1" }));
|
||||
expect(screen.getByLabelText("Mapping 1 tag prefix")).toHaveValue("usecase:");
|
||||
expect(changed).toHaveBeenLastCalledWith([variable]);
|
||||
});
|
||||
});
|
||||
|
|
@ -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<RequestTagMapping>) =>
|
||||
onChange(rules.map((rule, position) => (position === index ? { ...rule, ...patch } : rule)));
|
||||
|
||||
return (
|
||||
<div className="min-w-80 space-y-4 whitespace-normal">
|
||||
<p className="text-xs text-muted-foreground">
|
||||
Exact value mappings take priority over any-value rules for the same key.
|
||||
</p>
|
||||
{rules.map((rule, index) => {
|
||||
const exact = isExact(rule);
|
||||
const problem = problemAt(rules, index);
|
||||
const label = `Mapping ${index + 1}`;
|
||||
return (
|
||||
<fieldset key={index} disabled={disabled} className="space-y-3 rounded-lg border p-3 disabled:opacity-60">
|
||||
<legend className="px-1 text-xs font-medium">{label}</legend>
|
||||
<div className="flex items-center gap-2">
|
||||
<Select
|
||||
items={MODES}
|
||||
value={exact ? "exact" : "any"}
|
||||
onValueChange={(mode) =>
|
||||
update(
|
||||
index,
|
||||
mode === "exact"
|
||||
? { value: "", tag: "", tag_prefix: undefined }
|
||||
: { value: undefined, tag: undefined, tag_prefix: "" },
|
||||
)
|
||||
}
|
||||
>
|
||||
<SelectTrigger aria-label={`${label} mode`}>
|
||||
<SelectValue />
|
||||
</SelectTrigger>
|
||||
<SelectContent>
|
||||
{MODES.map((mode) => (
|
||||
<SelectItem key={mode.value} value={mode.value}>
|
||||
{mode.label}
|
||||
</SelectItem>
|
||||
))}
|
||||
</SelectContent>
|
||||
</Select>
|
||||
<Button
|
||||
type="button"
|
||||
variant="ghost"
|
||||
size="icon"
|
||||
aria-label={`Delete ${label.toLowerCase()}`}
|
||||
onClick={() => onChange(rules.filter((_, position) => position !== index))}
|
||||
>
|
||||
<Trash2 className="h-4 w-4" />
|
||||
</Button>
|
||||
</div>
|
||||
<label className="flex flex-col gap-1 text-xs">
|
||||
Baggage key
|
||||
<Input
|
||||
aria-label={`${label} baggage key`}
|
||||
placeholder="entrypoint"
|
||||
value={rule.key}
|
||||
onChange={(event) => update(index, { key: event.target.value })}
|
||||
/>
|
||||
</label>
|
||||
{exact ? (
|
||||
<>
|
||||
<label className="flex flex-col gap-1 text-xs">
|
||||
Match value
|
||||
<Input
|
||||
aria-label={`${label} match value`}
|
||||
value={rule.value ?? ""}
|
||||
onChange={(event) => update(index, { value: event.target.value })}
|
||||
/>
|
||||
</label>
|
||||
<label className="flex flex-col gap-1 text-xs">
|
||||
Tag
|
||||
<Input
|
||||
aria-label={`${label} tag`}
|
||||
value={rule.tag ?? ""}
|
||||
onChange={(event) => update(index, { tag: event.target.value })}
|
||||
/>
|
||||
</label>
|
||||
<p className="text-xs text-muted-foreground">
|
||||
{rule.key}={rule.value || "value"} → {rule.tag || "tag"}
|
||||
</p>
|
||||
</>
|
||||
) : (
|
||||
<>
|
||||
<label className="flex flex-col gap-1 text-xs">
|
||||
Tag prefix
|
||||
<Input
|
||||
aria-label={`${label} tag prefix`}
|
||||
placeholder="usecase:"
|
||||
value={rule.tag_prefix ?? ""}
|
||||
onChange={(event) => update(index, { tag_prefix: event.target.value })}
|
||||
/>
|
||||
</label>
|
||||
<p className="text-xs text-muted-foreground">
|
||||
{rule.key}=x → {rule.tag_prefix}x
|
||||
</p>
|
||||
</>
|
||||
)}
|
||||
{problem ? <p className="text-xs text-destructive">{problem}</p> : null}
|
||||
<label className="flex items-center justify-between gap-3 text-xs">
|
||||
Require a mapped value
|
||||
<Switch
|
||||
aria-label={`${label} required`}
|
||||
checked={rule.required ?? false}
|
||||
onCheckedChange={(checked) => update(index, { required: checked })}
|
||||
/>
|
||||
</label>
|
||||
<label className="flex items-center justify-between gap-3 text-xs">
|
||||
Require an existing tag budget
|
||||
<Switch
|
||||
aria-label={`${label} require budget`}
|
||||
checked={rule.require_budgeted_tag ?? false}
|
||||
onCheckedChange={(checked) => update(index, { require_budgeted_tag: checked })}
|
||||
/>
|
||||
</label>
|
||||
</fieldset>
|
||||
);
|
||||
})}
|
||||
{disabled ? (
|
||||
<p className="text-xs text-muted-foreground">Managed in config.yaml. Edit the YAML to change these rules.</p>
|
||||
) : null}
|
||||
<Button
|
||||
disabled={disabled}
|
||||
type="button"
|
||||
variant="outline"
|
||||
onClick={() => onChange([...rules, { header: "baggage", key: "", value: "", tag: "" }])}
|
||||
>
|
||||
Add mapping
|
||||
</Button>
|
||||
</div>
|
||||
);
|
||||
}
|
||||
|
|
@ -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(<GeneralSettings accessToken="token" userRole="Admin" userID="user" />);
|
||||
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" }]],
|
||||
]);
|
||||
});
|
||||
|
|
|
|||
|
|
@ -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 (
|
||||
<RequestTagMappingsEditor
|
||||
value={setting.field_value}
|
||||
disabled={setting.editable === false}
|
||||
onChange={(rules) => onChange(setting.field_name, rules)}
|
||||
/>
|
||||
);
|
||||
}
|
||||
if (setting.field_type === "Integer") {
|
||||
return (
|
||||
<Input
|
||||
|
|
@ -240,7 +251,7 @@ const GeneralSettings: React.FC<GeneralSettingsPageProps> = ({ accessToken, user
|
|||
setGeneralSettings(updatedSettings);
|
||||
};
|
||||
|
||||
const handleUpdateField = (fieldName: string) => {
|
||||
const handleUpdateField = async (fieldName: string) => {
|
||||
if (!accessToken) {
|
||||
return;
|
||||
}
|
||||
|
|
@ -253,7 +264,7 @@ const GeneralSettings: React.FC<GeneralSettingsPageProps> = ({ 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<GeneralSettingsPageProps> = ({ accessToken, user
|
|||
)}
|
||||
</TableCell>
|
||||
<TableCell>
|
||||
<Button onClick={() => handleUpdateField(value.field_name)}>Update</Button>
|
||||
<span
|
||||
onClick={() => handleResetField(value.field_name)}
|
||||
className="inline-flex shrink-0 cursor-pointer items-center justify-center px-1.5 py-1.5 text-destructive"
|
||||
<Button
|
||||
disabled={
|
||||
value.field_name === "request_tag_mappings" &&
|
||||
(value.editable === false || !requestTagMappingsAreValid(value.field_value))
|
||||
}
|
||||
onClick={() => handleUpdateField(value.field_name)}
|
||||
>
|
||||
<Trash2 className="h-5 w-5 shrink-0" />
|
||||
</span>
|
||||
Update
|
||||
</Button>
|
||||
{value.field_name === "request_tag_mappings" && value.editable === false ? null : (
|
||||
<span
|
||||
onClick={() => handleResetField(value.field_name)}
|
||||
className="inline-flex shrink-0 cursor-pointer items-center justify-center px-1.5 py-1.5 text-destructive"
|
||||
>
|
||||
<Trash2 className="h-5 w-5 shrink-0" />
|
||||
</span>
|
||||
)}
|
||||
</TableCell>
|
||||
</TableRow>
|
||||
))}
|
||||
|
|
|
|||
35
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
35
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
|
|
@ -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.
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue