diff --git a/litellm/llms/bedrock/base_aws_llm.py b/litellm/llms/bedrock/base_aws_llm.py index c31462a735b..b71f37023e8 100644 --- a/litellm/llms/bedrock/base_aws_llm.py +++ b/litellm/llms/bedrock/base_aws_llm.py @@ -1,3 +1,4 @@ +import base64 import hashlib import json import os @@ -19,7 +20,7 @@ from typing import ( ) import httpx -from pydantic import BaseModel +from pydantic import BaseModel, ValidationError from litellm._logging import verbose_logger from litellm.caching.caching import DualCache @@ -56,6 +57,11 @@ class Boto3CredentialsInfo(BaseModel): aws_bedrock_runtime_endpoint: Optional[str] +class _WebIdentityTokenClaims(BaseModel): + aud: Optional[Union[str, list[str]]] = None + iss: Optional[str] = None + + class AwsAuthError(Exception): def __init__(self, status_code, message): self.status_code = status_code @@ -817,6 +823,25 @@ class BaseAWSLLM: return False + @staticmethod + def _unverified_web_identity_audience(oidc_token: str) -> Optional[str]: + """Return the public ``aud``/``iss`` claims of a web identity JWT + without verifying its signature, so a rejected-token error can name + the audience LiteLLM actually sent. The signature is never read, so no + secret is exposed.""" + segments = oidc_token.split(".") + if len(segments) != 3: + return None + payload = segments[1] + try: + decoded = base64.urlsafe_b64decode(payload + "=" * (-len(payload) % 4)) + claims = _WebIdentityTokenClaims.model_validate_json(decoded) + except (ValueError, ValidationError): + return None + if claims.aud is None and claims.iss is None: + return None + return f"aud={claims.aud!r}, iss={claims.iss!r}" + @tracer.wrap() def _auth_with_web_identity_token( self, @@ -925,7 +950,21 @@ class BaseAWSLLM: if aws_external_id is not None: assume_role_params["ExternalId"] = aws_external_id - sts_response = sts_client.assume_role_with_web_identity(**assume_role_params) + try: + sts_response = sts_client.assume_role_with_web_identity( + **assume_role_params + ) + except sts_client.exceptions.InvalidIdentityTokenException as e: + audience = ( + self._unverified_web_identity_audience(oidc_token) + if isinstance(oidc_token, str) + else None + ) + detail = f" Token {audience}" if audience else "" + raise AwsAuthError( + status_code=401, + message=f"AWS STS rejected the web identity token: {e}.{detail}", + ) from e iam_creds_dict = { "aws_access_key_id": sts_response["Credentials"]["AccessKeyId"], diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 5cf5af8a22a..6a65977e909 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -25610,6 +25610,16 @@ ], "source": "https://mistral.ai/pricing#api-pricing" }, + "mistral/mistral-ocr-2512": { + "litellm_provider": "mistral", + "ocr_cost_per_page": 0.002, + "annotation_cost_per_page": 0.003, + "mode": "ocr", + "supported_endpoints": [ + "/v1/ocr" + ], + "source": "https://mistral.ai/pricing#api-pricing" + }, "mistral/magistral-medium-latest": { "input_cost_per_token": 2e-06, "litellm_provider": "mistral", diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index ed030c31b19..559720eb64e 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -3658,6 +3658,10 @@ class TeamMemberAddRequest(MemberAddRequest): default=None, description="Maximum budget allocated to this user within the team. If not set, user has unlimited budget within team limits", ) + budget_duration: Optional[str] = Field( + default=None, + description="Duration after which this team member's budget resets (e.g. '1h', '24h', '7d', '30d'). If not set, the budget never resets.", + ) allowed_models: Optional[List[str]] = Field( default=None, description="List of models this team member can access. If not set, inherits the team's default_team_member_models or all team models.", diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 2dcc7bdbd19..184d91fb103 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -1991,13 +1991,16 @@ async def _user_api_key_auth_builder( else: valid_token.team_object_permission = None - # Only cache when the key is a real team_id (non-team keys must not use key=None). + # Cache under the canonical "team_id:{id}" key so get_team_object and + # _update_team_cache serve this write from the L2 cache. The guard keeps a + # non-team (personal) key, whose team_id is None, from reaching the cache + # layer, which Redis rejects with a NoneType key error. if valid_token.team_id is not None and _team_obj is not None: await user_api_key_cache.async_set_cache( - key=valid_token.team_id, + key=f"team_id:{valid_token.team_id}", value=_team_obj, model_type=LiteLLM_TeamTableCachedObj, - ) # save team table in cache - used for tpm/rpm limiting - tpm_rpm_limiter.py + ) # Fetch project object if key belongs to a project _project_obj = None diff --git a/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/__init__.py index bb0d0a99b31..777c45bcd83 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/__init__.py @@ -21,6 +21,7 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail" unreachable_fallback=getattr( litellm_params, "unreachable_fallback", "fail_closed" ), + fail_on_error=getattr(litellm_params, "fail_on_error", True), extra_headers=getattr(litellm_params, "extra_headers", None), guardrail_name=guardrail.get("guardrail_name", ""), event_hook=litellm_params.mode, diff --git a/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py b/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py index 790ee31f2e0..8d3abe2bdc3 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py +++ b/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py @@ -27,6 +27,7 @@ from litellm.types.proxy.guardrails.guardrail_hooks.generic_guardrail_api import GenericGuardrailAPIMetadata, GenericGuardrailAPIRequest, GenericGuardrailAPIResponse, + GuardrailToolParam, ) from litellm.types.utils import GenericGuardrailAPIInputs @@ -187,6 +188,7 @@ class GenericGuardrailAPI(CustomGuardrail): api_key: Optional[str] = None, additional_provider_specific_params: Optional[Dict[str, Any]] = None, unreachable_fallback: Literal["fail_closed", "fail_open"] = "fail_closed", + fail_on_error: Optional[bool] = True, extra_headers: Optional[list] = None, **kwargs, ): @@ -223,6 +225,8 @@ class GenericGuardrailAPI(CustomGuardrail): unreachable_fallback ) + self.fail_on_error: bool = True if fail_on_error is None else fail_on_error + # Set supported event hooks if "supported_event_hooks" not in kwargs: kwargs["supported_event_hooks"] = [ @@ -299,7 +303,7 @@ class GenericGuardrailAPI(CustomGuardrail): f" http_status_code={http_status_code}" if http_status_code else "" ) verbose_proxy_logger.critical( - "Generic Guardrail API unreachable (fail-open). Proceeding without guardrail.%s " + "Generic Guardrail API error (fail-open). Proceeding without guardrail.%s " "guardrail_name=%s api_base=%s input_type=%s litellm_call_id=%s litellm_trace_id=%s", status_suffix, getattr(self, "guardrail_name", None), @@ -351,7 +355,10 @@ class GenericGuardrailAPI(CustomGuardrail): logging_obj: Optional["LiteLLMLoggingObj"], is_unreachable: bool = True, ) -> GenericGuardrailAPIInputs: - if is_unreachable and self.unreachable_fallback == "fail_open": + unreachable_fail_open = ( + is_unreachable and self.unreachable_fallback == "fail_open" + ) + if unreachable_fail_open or not self.fail_on_error: http_status_code = getattr( getattr(error, "response", None), "status_code", None ) @@ -432,26 +439,30 @@ class GenericGuardrailAPI(CustomGuardrail): extra_allowlist=extra_allowlist, ) - # Create request payload - guardrail_request = GenericGuardrailAPIRequest( - litellm_call_id=logging_obj.litellm_call_id if logging_obj else None, - litellm_trace_id=logging_obj.litellm_trace_id if logging_obj else None, - texts=texts, - request_data=user_metadata, - request_headers=inbound_headers, - litellm_version=litellm_version, - images=images, - tools=tools, - structured_messages=structured_messages, - tool_calls=tool_calls, - additional_provider_specific_params=additional_params, - input_type=input_type, - model=model, - ) - - headers = self._build_request_headers() - try: + # Create request payload + guardrail_request = GenericGuardrailAPIRequest( + litellm_call_id=logging_obj.litellm_call_id if logging_obj else None, + litellm_trace_id=logging_obj.litellm_trace_id if logging_obj else None, + texts=texts, + request_data=user_metadata, + request_headers=inbound_headers, + litellm_version=litellm_version, + images=images, + tools=( + [GuardrailToolParam.model_validate(t) for t in tools] + if tools + else None + ), + structured_messages=structured_messages, + tool_calls=tool_calls, + additional_provider_specific_params=additional_params, + input_type=input_type, + model=model, + ) + + headers = self._build_request_headers() + # Make the API request # Use mode="json" to ensure all iterables are converted to lists response = await self.async_handler.post( diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 249cfdae4eb..f533774ce18 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -2225,6 +2225,26 @@ async def _validate_team_member_add_permissions( }, ) + # Available-team self-join grants only the ability to join; per-member + # budget and model controls stay admin-only. Reject them here so a + # self-joining non-admin cannot set their own cap, reset window, or model + # scope via the bypass. + if ( + data.max_budget_in_team is not None + or data.budget_duration is not None + or data.allowed_models is not None + ): + raise HTTPException( + status_code=403, + detail={ + "error": ( + "Available-team self-join cannot set per-member budget or " + "model controls (max_budget_in_team, budget_duration, " + "allowed_models); these are admin-only." + ) + }, + ) + # Available-team self-join: caller may add only themselves, only as a # standard user. Enforce that here so the bypass cannot be used as a # privilege-escalation or cross-user-injection primitive. @@ -2290,6 +2310,7 @@ async def _process_team_members( team_id=data.team_id, default_team_budget_id=default_team_budget_id, allowed_models=member_allowed_models, + budget_duration=data.budget_duration, ) except Exception as e: raise HTTPException( @@ -2315,6 +2336,7 @@ async def _process_team_members( team_id=data.team_id, default_team_budget_id=default_team_budget_id, allowed_models=member_allowed_models, + budget_duration=data.budget_duration, ) except Exception as e: raise HTTPException( @@ -2578,6 +2600,8 @@ async def team_member_add( except HTTPException as e: raise e + _validate_budget_duration(data.budget_duration) + prisma_client = cast(PrismaClient, prisma_client) existing_team_row = await get_team_object( diff --git a/litellm/proxy/management_helpers/utils.py b/litellm/proxy/management_helpers/utils.py index 4ce8633ed57..0c5439e26e2 100644 --- a/litellm/proxy/management_helpers/utils.py +++ b/litellm/proxy/management_helpers/utils.py @@ -167,6 +167,7 @@ async def _clone_team_default_budget_for_member( default_team_budget_id: str, user_api_key_dict: UserAPIKeyAuth, litellm_proxy_admin_name: str, + budget_duration_override: Optional[str] = None, ) -> Optional[str]: """ Create a new budget row that copies the values from the team's default @@ -176,6 +177,10 @@ async def _clone_team_default_budget_for_member( Used when adding a new team member without an explicit per-member budget, so the member starts with the team default's values but gets their own private budget row (which can be edited independently). + + ``budget_duration_override`` replaces the default's reset window for this + member while keeping the default's other limits, so an admin can set a + member's reset cadence without discarding the team default's max_budget. """ default_budget = await BudgetRepository(prisma_client).table.find_unique( where={"budget_id": default_team_budget_id} @@ -198,6 +203,9 @@ async def _clone_team_default_budget_for_member( continue cloned_data[field] = value + if budget_duration_override is not None: + cloned_data["budget_duration"] = budget_duration_override + # Start the member's budget window at clone time, not the pool's reset # timestamp — otherwise a member joining mid-cycle inherits a stale reset. if cloned_data.get("budget_duration"): @@ -209,6 +217,55 @@ async def _clone_team_default_budget_for_member( return new_budget.budget_id +async def _resolve_member_budget_id( + prisma_client: PrismaClient, + user_api_key_dict: UserAPIKeyAuth, + litellm_proxy_admin_name: str, + max_budget_in_team: Optional[float], + allowed_models: Optional[list[str]], + budget_duration: Optional[str], + default_team_budget_id: Optional[str], +) -> Optional[str]: + """ + Resolve the budget a new team member should be linked to. + + Explicit per-member limits create a fresh budget. Otherwise the team's + default member budget is cloned (with ``budget_duration`` overriding its + reset window while keeping its other limits). A lone ``budget_duration`` + with no team default creates a window-only budget. With nothing set the + member gets no budget. + """ + has_explicit_limit = max_budget_in_team is not None or allowed_models is not None + + if not has_explicit_limit and default_team_budget_id is not None: + return await _clone_team_default_budget_for_member( + prisma_client=prisma_client, + default_team_budget_id=default_team_budget_id, + user_api_key_dict=user_api_key_dict, + litellm_proxy_admin_name=litellm_proxy_admin_name, + budget_duration_override=budget_duration, + ) + + if not has_explicit_limit and budget_duration is None: + return None + + budget_data: dict = { + "created_by": user_api_key_dict.user_id or litellm_proxy_admin_name, + "updated_by": user_api_key_dict.user_id or litellm_proxy_admin_name, + } + if max_budget_in_team is not None: + budget_data["max_budget"] = max_budget_in_team + if allowed_models is not None: + budget_data["allowed_models"] = allowed_models + if budget_duration is not None: + budget_data["budget_duration"] = budget_duration + budget_data["budget_reset_at"] = get_budget_reset_time( + budget_duration=budget_duration + ) + response = await BudgetRepository(prisma_client).table.create(data=budget_data) + return response.budget_id + + async def add_new_member( new_member: Member, max_budget_in_team: Optional[float], @@ -218,6 +275,7 @@ async def add_new_member( litellm_proxy_admin_name: str, default_team_budget_id: Optional[str] = None, allowed_models: Optional[List[str]] = None, + budget_duration: Optional[str] = None, ) -> Tuple[LiteLLM_UserTable, Optional[LiteLLM_TeamMembership]]: """ Add a new member to a team @@ -278,34 +336,15 @@ async def add_new_member( }, ) - # Check if trying to set a budget or model scope for team member - if max_budget_in_team is not None or allowed_models is not None: - # create a new budget item for this member - budget_data: dict = { - "created_by": user_api_key_dict.user_id or litellm_proxy_admin_name, - "updated_by": user_api_key_dict.user_id or litellm_proxy_admin_name, - } - if max_budget_in_team is not None: - budget_data["max_budget"] = max_budget_in_team - if allowed_models is not None: - budget_data["allowed_models"] = allowed_models - response = await BudgetRepository(prisma_client).table.create(data=budget_data) - - _budget_id = response.budget_id - elif default_team_budget_id is not None: - # No per-member budget was provided, but the team has a default member - # budget. Clone the default budget into a new row for this user so that - # later edits to one member's budget do not bleed into other members. - # If the default no longer exists in the DB, fall back to no budget. - _budget_id = await _clone_team_default_budget_for_member( - prisma_client=prisma_client, - default_team_budget_id=default_team_budget_id, - user_api_key_dict=user_api_key_dict, - litellm_proxy_admin_name=litellm_proxy_admin_name, - ) - else: - # No per-member budget and no team default → member gets no budget. - _budget_id = None + _budget_id = await _resolve_member_budget_id( + prisma_client=prisma_client, + user_api_key_dict=user_api_key_dict, + litellm_proxy_admin_name=litellm_proxy_admin_name, + max_budget_in_team=max_budget_in_team, + allowed_models=allowed_models, + budget_duration=budget_duration, + default_team_budget_id=default_team_budget_id, + ) if _budget_id and returned_user is not None and returned_user.user_id is not None: _returned_team_membership = await TeamMembershipRepository( diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index 8ca048a0caa..3374877572b 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -2018,11 +2018,6 @@ async def ui_view_spend_logs( order_column = sort_by order_direction = (sort_order or "desc").lower() - # Get total count of records - total_records = await SpendLogsRepository(prisma_client).table.count( - where=where_conditions, - ) - # Build raw SQL to fetch paginated data WITHOUT heavy columns # (messages, response, proxy_server_request can be hundreds of KB per row). # These are only needed in the detail endpoint /spend/logs/ui/{request_id}. @@ -2136,7 +2131,8 @@ async def ui_view_spend_logs( cache_hit, cache_key, request_tags, team_id, organization_id, end_user, requester_ip_address, session_id, status, mcp_namespaced_tool_name, agent_id, - COALESCE(request_duration_ms, (EXTRACT(EPOCH FROM ("endTime" - "startTime")) * 1000)::INTEGER) AS request_duration_ms + COALESCE(request_duration_ms, (EXTRACT(EPOCH FROM ("endTime" - "startTime")) * 1000)::INTEGER) AS request_duration_ms, + COUNT(*) OVER () AS total_count FROM "LiteLLM_SpendLogs" WHERE {" AND ".join(sql_conditions)} ORDER BY {_order_expr} {_sql_dir}{_nulls_clause} @@ -2146,13 +2142,34 @@ async def ui_view_spend_logs( data = await prisma_client.db.query_raw(sql_query, *sql_params) + # `COUNT(*) OVER ()` folds the total-match count into the same scan as the + # page data; a standalone `COUNT(*)` is a distributed RPC on sharded + # engines like YugabyteDB that contacts every tablet and times out + # regardless of row count (LIT-4027). The hot path (page 1 and in-range + # pages) always carries the count on its rows, so the count round trip is + # gone there. Only an out-of-range page overshoots the last row and comes + # back empty; fall back to a direct count there so total/total_pages stay + # accurate rather than collapsing to zero. + if data: + total_records = int(data[0]["total_count"]) + elif page > 1: + total_records = int( + await SpendLogsRepository(prisma_client).table.count( + where=where_conditions, + ) + ) + else: + total_records = 0 + # query_raw returns the JSONB `metadata` column as a string (the Prisma # serialiser bypasses the model-layer JSON hydration we get on the ORM # path). The UI reads `metadata.status` / `metadata.error_information` # as object fields, so failure rows looked like successes (#29674). - # Re-hydrate to dict here. + # Re-hydrate to dict here. Also drop the window-function `total_count` + # helper column so it does not leak into the serialised rows. for row in data: if isinstance(row, dict): + row.pop("total_count", None) md = row.get("metadata") if isinstance(md, str): try: diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 781f0e9f301..f6d9d8975be 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -100,6 +100,7 @@ from litellm.proxy.hooks.sensitive_data_routing import ( _PROXY_SensitiveDataRoutingHandler, ) from litellm.integrations.custom_logger import CustomLogger +from litellm.integrations.prometheus import PrometheusLogger from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting from litellm.integrations.SlackAlerting.utils import _add_langfuse_trace_id_to_alert from litellm.litellm_core_utils.litellm_logging import Logging @@ -1124,7 +1125,6 @@ class ProxyLogging: Returns: Updated data dictionary if guardrail passes, None if guardrail should be skipped """ - from litellm.integrations.prometheus import PrometheusLogger from litellm.types.guardrails import GuardrailEventHooks # Determine the event type based on call type @@ -1197,17 +1197,13 @@ class ProxyLogging: or "unknown" ) - # Find PrometheusLogger in callbacks and record metrics - for prom_callback in litellm.callbacks: - if isinstance(prom_callback, PrometheusLogger): - prom_callback._record_guardrail_metrics( - guardrail_name=metrics_guardrail_name, - latency_seconds=latency_seconds, - status=status, - error_type=error_type, - hook_type="pre_call", - ) - break + self._emit_guardrail_metrics( + guardrail_name=metrics_guardrail_name, + latency_seconds=latency_seconds, + status=status, + error_type=error_type, + hook_type="pre_call", + ) return data @@ -1625,19 +1621,58 @@ class ProxyLogging: return data @staticmethod - async def _run_guardrail_task_with_enrichment( - callback: Any, coro: Awaitable[Any] + def _emit_guardrail_metrics( + guardrail_name: str, + latency_seconds: float, + status: str, + error_type: Optional[str], + hook_type: str, + ) -> None: + for prom_callback in litellm.callbacks: + if isinstance(prom_callback, PrometheusLogger): + prom_callback._record_guardrail_metrics( + guardrail_name=guardrail_name, + latency_seconds=latency_seconds, + status=status, + error_type=error_type, + hook_type=hook_type, + ) + break + + @staticmethod + async def _run_guardrail_with_metrics( + callback: Any, coro: Awaitable[Any], hook_type: str ) -> Any: """ - Await `coro`; if it raises an HTTPException with dict detail, - enrich the detail with the originating callback's `guardrail_name` - and `guardrail_mode` before re-raising. + Await `coro`, recording its latency and status to the + `litellm_guardrail_latency_seconds` metric under `hook_type`, and + enriching any raised HTTPException with the originating callback's + `guardrail_name`/`guardrail_mode` before re-raising. """ + guardrail_name = ( + getattr(callback, "guardrail_name", None) or type(callback).__name__ + ) + start_time = time.perf_counter() + status = "success" + error_type: Optional[str] = None try: return await coro + except SensitiveDataRouteException: + status = "intervened" + raise except Exception as e: + status = "error" + error_type = type(e).__name__ _enrich_http_exception_with_guardrail_context(e, callback) raise + finally: + ProxyLogging._emit_guardrail_metrics( + guardrail_name=guardrail_name, + latency_seconds=time.perf_counter() - start_time, + status=status, + error_type=error_type, + hook_type=hook_type, + ) @staticmethod async def _wrap_streaming_iterator_with_enrichment( @@ -1853,22 +1888,24 @@ class ProxyLogging: and not getattr(callback, "use_native_during_call_hook", False) ): data["guardrail_to_apply"] = callback - guardrail_task = self._run_guardrail_task_with_enrichment( + guardrail_task = self._run_guardrail_with_metrics( callback, unified_guardrail.async_moderation_hook( user_api_key_dict=user_api_key_dict, data=data, call_type=call_type, ), + "during_call", ) else: - guardrail_task = self._run_guardrail_task_with_enrichment( + guardrail_task = self._run_guardrail_with_metrics( callback, callback.async_moderation_hook( data=data, user_api_key_dict=user_api_key_auth_dict, # type: ignore call_type=call_type, # type: ignore ), + "during_call", ) guardrail_tasks.append(guardrail_task) @@ -2394,29 +2431,25 @@ class ProxyLogging: if "apply_guardrail" in type(callback).__dict__: data["guardrail_to_apply"] = callback - try: - guardrail_response = ( - await unified_guardrail.async_post_call_success_hook( - user_api_key_dict=user_api_key_dict, - data=data, - response=response, - ) - ) - except Exception as e: - _enrich_http_exception_with_guardrail_context(e, callback) - raise + guardrail_response = await self._run_guardrail_with_metrics( + callback, + unified_guardrail.async_post_call_success_hook( + user_api_key_dict=user_api_key_dict, + data=data, + response=response, + ), + "post_call", + ) else: - try: - guardrail_response = ( - await callback.async_post_call_success_hook( - user_api_key_dict=user_api_key_dict, - data=data, - response=response, - ) - ) - except Exception as e: - _enrich_http_exception_with_guardrail_context(e, callback) - raise + guardrail_response = await self._run_guardrail_with_metrics( + callback, + callback.async_post_call_success_hook( + user_api_key_dict=user_api_key_dict, + data=data, + response=response, + ), + "post_call", + ) if guardrail_response is not None: response = guardrail_response diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py index c9623d8595a..6a76dd337fa 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -750,7 +750,12 @@ class BaseLitellmParams( ) fail_on_error: Optional[bool] = Field( default=True, - description="Whether to fail the request if Model Armor encounters an error", + description=( + "Whether to fail the request if the guardrail encounters an error. " + "Implemented by guardrail='model_armor' and 'generic_guardrail_api'. " + "True (default) raises the error. False logs a critical error and lets the request proceed, " + "so only a valid guardrail response can block or modify it." + ), ) additional_provider_specific_params: Optional[Dict[str, Any]] = Field( diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/generic_guardrail_api.py b/litellm/types/proxy/guardrails/guardrail_hooks/generic_guardrail_api.py index 94f219a5fc6..2fe976463c4 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/generic_guardrail_api.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/generic_guardrail_api.py @@ -1,17 +1,28 @@ from typing import Any, Dict, List, Literal, Optional, Union -from pydantic import BaseModel, Field +from pydantic import BaseModel, ConfigDict, Field from typing_extensions import TYPE_CHECKING, TypedDict from litellm.types.llms.openai import ( AllMessageValues, ChatCompletionToolCallChunk, - ChatCompletionToolParam, ) from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel from litellm.types.utils import ChatCompletionMessageToolCall +class GuardrailToolParam(BaseModel): + """A tool forwarded verbatim to the guardrail for inspection. + + Built-in tools (code_interpreter, file_search, ...) have no ``function`` block + and stash their config in tool-specific keys, so only ``type`` is required and + ``extra="allow"`` preserves the rest instead of stripping it. + """ + + model_config = ConfigDict(extra="allow") + type: str + + class GenericGuardrailAPIMetadata(TypedDict, total=False): user_api_key_hash: Optional[str] user_api_key_alias: Optional[str] @@ -39,6 +50,16 @@ class GenericGuardrailAPIOptionalParams(BaseModel): ), ) + fail_on_error: Optional[bool] = Field( + default=True, + description=( + "Behavior on any guardrail error, not just unreachability. " + "True (default) raises and blocks the request on error. " + "False logs a critical error and allows the request to proceed, so only a valid " + "guardrail response can block or modify it; broader than unreachable_fallback." + ), + ) + class GenericGuardrailAPIConfigModel( GuardrailConfigModel[GenericGuardrailAPIOptionalParams], @@ -65,7 +86,7 @@ class GenericGuardrailAPIRequest(BaseModel): ) structured_messages: Optional[List[AllMessageValues]] = None images: Optional[List[str]] = None - tools: Optional[List[ChatCompletionToolParam]] = None + tools: Optional[List[GuardrailToolParam]] = None texts: Optional[List[str]] = None request_data: GenericGuardrailAPIMetadata request_headers: Optional[Dict[str, str]] = Field( @@ -88,7 +109,7 @@ class GenericGuardrailAPIResponse: texts: Optional[List[str]] images: Optional[List[str]] - tools: Optional[List[ChatCompletionToolParam]] + tools: Optional[List[GuardrailToolParam]] action: str blocked_reason: Optional[str] @@ -98,7 +119,7 @@ class GenericGuardrailAPIResponse: texts: Optional[List[str]] = None, blocked_reason: Optional[str] = None, images: Optional[List[str]] = None, - tools: Optional[List[ChatCompletionToolParam]] = None, + tools: Optional[List[GuardrailToolParam]] = None, ): self.action = action self.blocked_reason = blocked_reason diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index ed1b2fb7db9..d5af7391439 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -25774,6 +25774,16 @@ ], "source": "https://mistral.ai/pricing#api-pricing" }, + "mistral/mistral-ocr-2512": { + "litellm_provider": "mistral", + "ocr_cost_per_page": 0.002, + "annotation_cost_per_page": 0.003, + "mode": "ocr", + "supported_endpoints": [ + "/v1/ocr" + ], + "source": "https://mistral.ai/pricing#api-pricing" + }, "mistral/magistral-medium-latest": { "input_cost_per_token": 2e-06, "litellm_provider": "mistral", diff --git a/pyproject.toml b/pyproject.toml index 6e99d81f8f3..23f02cb5762 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -234,24 +234,8 @@ healthcheck = [ ] [build-system] -requires = ["maturin==1.9.4"] -build-backend = "maturin" - -[tool.maturin] -manifest-path = "litellm-rust/crates/python-bridge/Cargo.toml" -module-name = "litellm.rust_bridge._native" -python-source = "." -bindings = "pyo3" -exclude = [ - "litellm/proxy/enterprise", - "litellm/proxy/enterprise/**", - "**/__pycache__", - "**/__pycache__/**", - "**/.pytest_cache", - "**/.pytest_cache/**", - "**/.ruff_cache", - "**/.ruff_cache/**", -] +requires = ["uv_build==0.11.8"] +build-backend = "uv_build" [tool.uv] constraint-dependencies = [ @@ -269,6 +253,18 @@ litellm-enterprise = { workspace = true } [tool.uv.workspace] members = ["enterprise", "litellm-proxy-extras"] +[tool.uv.build-backend] +module-root = "" +source-exclude = [ + "litellm/proxy/enterprise", + "**/__pycache__", + "**/__pycache__/**", + "**/.pytest_cache", + "**/.pytest_cache/**", + "**/.ruff_cache", + "**/.ruff_cache/**", +] + [tool.isort] profile = "black" diff --git a/tests/local_testing/test_pass_through_endpoints.py b/tests/local_testing/test_pass_through_endpoints.py index 68ba62bcbab..eeb29dea531 100644 --- a/tests/local_testing/test_pass_through_endpoints.py +++ b/tests/local_testing/test_pass_through_endpoints.py @@ -402,7 +402,9 @@ async def test_aaapass_through_endpoint_pass_through_keys_langfuse( mock_api_key = "sk-my-test-key" cache_value = UserAPIKeyAuth( - token=hash_token(mock_api_key), rpm_limit=rpm_limit + token=hash_token(mock_api_key), + rpm_limit=rpm_limit, + metadata={"allowed_passthrough_routes": ["/api/public/ingestion"]}, ) _cohere_api_key = os.environ.get("COHERE_API_KEY") diff --git a/tests/test_litellm/llms/bedrock/test_web_identity_session_policy.py b/tests/test_litellm/llms/bedrock/test_web_identity_session_policy.py index 2cf1fa16e91..7e9c8a273ae 100644 --- a/tests/test_litellm/llms/bedrock/test_web_identity_session_policy.py +++ b/tests/test_litellm/llms/bedrock/test_web_identity_session_policy.py @@ -29,6 +29,7 @@ claude_platform statement are present and cover every documented action. """ +import base64 import json from datetime import datetime, timedelta, timezone from unittest.mock import MagicMock, patch @@ -157,6 +158,78 @@ class TestClaudePlatformActionsCovered: ) +def _make_jwt(payload: dict) -> str: + def _segment(data: dict) -> str: + return base64.urlsafe_b64encode(json.dumps(data).encode()).rstrip(b"=").decode() + + return f"{_segment({'alg': 'RS256', 'typ': 'JWT'})}.{_segment(payload)}.signature" + + +class TestInvalidIdentityTokenSurfacesAudience: + """LIT-4026: when STS rejects the web identity token with + ``InvalidIdentityToken`` (the "Incorrect token audience" case), the raised + error must name the ``aud``/``iss`` the token actually carries so an + operator can diagnose the mismatch without enabling LITELLM_LOG=DEBUG on a + prod instance.""" + + _AUD = "https://guidepoint.litellm-prod.ai" + _ISS = "https://accounts.google.com" + _STS_MESSAGE = ( + "An error occurred (InvalidIdentityToken) when calling the " + "AssumeRoleWithWebIdentity operation: Incorrect token audience" + ) + + def _raise_invalid_identity_token(self) -> Exception: + from litellm.llms.bedrock.base_aws_llm import AwsAuthError, BaseAWSLLM + + token = _make_jwt({"aud": self._AUD, "iss": self._ISS, "sub": "svc-account"}) + + mock_sts = MagicMock() + + class _InvalidIdentityTokenException(Exception): + pass + + mock_sts.exceptions.InvalidIdentityTokenException = ( + _InvalidIdentityTokenException + ) + mock_sts.assume_role_with_web_identity.side_effect = ( + _InvalidIdentityTokenException(self._STS_MESSAGE) + ) + + with ( + patch("boto3.client", return_value=mock_sts), + patch( + "litellm.llms.bedrock.base_aws_llm.get_secret", + return_value=token, + ), + pytest.raises(AwsAuthError) as exc_info, + ): + BaseAWSLLM()._auth_with_web_identity_token( + aws_web_identity_token="oidc/google/" + self._AUD, + aws_role_name="arn:aws:iam::123456789012:role/litellm-bedrock-role", + aws_session_name="test-session", + aws_region_name="us-east-1", + aws_sts_endpoint=None, + ) + return exc_info.value + + def test_error_names_token_audience(self): + err = self._raise_invalid_identity_token() + assert self._AUD in str(err) + + def test_error_names_token_issuer(self): + err = self._raise_invalid_identity_token() + assert self._ISS in str(err) + + def test_error_preserves_original_sts_reason(self): + err = self._raise_invalid_identity_token() + assert "Incorrect token audience" in str(err) + + def test_error_is_401(self): + err = self._raise_invalid_identity_token() + assert err.status_code == 401 + + class TestPolicyTransportConditions: def test_bedrock_statement_keeps_secure_transport_condition(self): policy = _captured_policy() diff --git a/tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py b/tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py index b98b0c22e50..c7f959826fe 100644 --- a/tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py +++ b/tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py @@ -5,6 +5,9 @@ for mistral-ocr-4-0 and mistral-ocr-latest, which now both resolve to OCR 4 at $4 / 1000 pages. """ +import json +from pathlib import Path + import pytest import litellm @@ -13,6 +16,14 @@ from litellm.llms.base_llm.ocr.transformation import OCRPage, OCRResponse, OCRUs OCR4_COST_PER_PAGE = 0.004 +REPO_ROOT = Path(__file__).parents[5] +MAIN_COST_MAP = REPO_ROOT / "model_prices_and_context_window.json" +BACKUP_COST_MAP = REPO_ROOT / "litellm" / "model_prices_and_context_window_backup.json" + +OCR3_MODEL = "mistral/mistral-ocr-2512" +OCR3_COST_PER_PAGE = 0.002 +OCR3_ANNOTATION_COST_PER_PAGE = 0.003 + def _ocr_response(model: str, pages_processed: int) -> OCRResponse: return OCRResponse( @@ -38,3 +49,43 @@ def test_ocr4_cost_scales_with_pages(model: str, pages_processed: int) -> None: call_type="ocr", ) assert cost == pytest.approx(OCR4_COST_PER_PAGE * pages_processed) + + +@pytest.fixture +def local_model_cost_map(monkeypatch): + """Force get_model_info to resolve against the in-repo cost map instead of the + remote one fetched at import time, which does not yet carry OCR 3 pricing.""" + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) + litellm.get_model_info.cache_clear() + yield + litellm.get_model_info.cache_clear() + + +@pytest.mark.parametrize("cost_map_path", [MAIN_COST_MAP, BACKUP_COST_MAP]) +def test_ocr3_pricing_entry(cost_map_path: Path) -> None: + with open(cost_map_path) as f: + info = json.load(f).get(OCR3_MODEL) + + assert info is not None, f"{OCR3_MODEL} missing from {cost_map_path.name}" + assert info["litellm_provider"] == "mistral" + assert info["mode"] == "ocr" + assert info["supported_endpoints"] == ["/v1/ocr"] + assert info["ocr_cost_per_page"] == OCR3_COST_PER_PAGE + assert info["annotation_cost_per_page"] == OCR3_ANNOTATION_COST_PER_PAGE + + +def test_ocr3_model_info_price(local_model_cost_map) -> None: + info = litellm.get_model_info(model=OCR3_MODEL, custom_llm_provider="mistral") + assert info["ocr_cost_per_page"] == OCR3_COST_PER_PAGE + + +@pytest.mark.parametrize("pages_processed", [1, 3, 10]) +def test_ocr3_cost_scales_with_pages(local_model_cost_map, pages_processed: int) -> None: + cost = completion_cost( + completion_response=_ocr_response("mistral-ocr-2512", pages_processed), + model=OCR3_MODEL, + custom_llm_provider="mistral", + call_type="ocr", + ) + assert cost == pytest.approx(OCR3_COST_PER_PAGE * pages_processed) diff --git a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py index f80a8b28cc3..4141b1f20c3 100644 --- a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py +++ b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py @@ -3989,3 +3989,98 @@ async def test_non_admin_cli_session_token_reaches_production_auth_path(monkeypa assert call_kwargs["valid_token_dict"]["is_session_token"] is True assert call_kwargs["valid_token_dict"]["user_role"] == LitellmUserRoles.INTERNAL_USER assert result.is_session_token is True + + +@pytest.mark.asyncio +async def test_auth_path_caches_team_object_under_canonical_team_id_key(): + """Regression for LIT-4000: the auth builder must cache the team object under + the canonical ``team_id:{id}`` key that ``get_team_object`` and + ``_update_team_cache`` read, never under the raw ``team_id`` (and never under + a ``None`` key, which Redis rejects with a NoneType key error). A raw or None + key is silently dropped by Redis / never served back, so every request + re-hits Postgres for the team object instead of the L2 cache. + + Drives the real builder for a team-scoped key against a real in-memory + ``UserApiKeyCache`` and reads the team object back. Mutating the cache key at + the write site to the raw ``valid_token.team_id`` (or ``None``) makes the + canonical-key read miss and fails this test. + """ + from fastapi import Request + from starlette.datastructures import URL + + import litellm.proxy.proxy_server as _proxy_server_mod + from litellm.proxy._types import LiteLLM_TeamTableCachedObj + from litellm.proxy.auth.user_api_key_auth import _user_api_key_auth_builder + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + from litellm.proxy.proxy_server import hash_token + + team_id = "team-lit-4000" + api_key = "sk-lit-4000-team-key" + cache = UserApiKeyCache() + + team_token = UserAPIKeyAuth(token=hash_token(api_key), team_id=team_id) + team_obj = LiteLLM_TeamTableCachedObj(team_id=team_id) + + proxy_logging_obj = MagicMock() + proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None) + attrs = { + "prisma_client": MagicMock(), + "user_api_key_cache": cache, + "proxy_logging_obj": proxy_logging_obj, + "master_key": "sk-test-master", + "general_settings": {"allow_requests_on_db_unavailable": False}, + "llm_model_list": [], + "llm_router": None, + "open_telemetry_logger": None, + "model_max_budget_limiter": MagicMock(), + "user_custom_auth": None, + "jwt_handler": None, + "litellm_proxy_admin_name": "admin", + } + originals = {a: getattr(_proxy_server_mod, a, None) for a in attrs} + try: + for k, v in attrs.items(): + setattr(_proxy_server_mod, k, v) + request = Request(scope={"type": "http"}) + request._url = URL(url="/chat/completions") + with ( + patch( + "litellm.proxy.auth.resolvers.store.IdentityStore._resolve_key", + AsyncMock(return_value=team_token), + ), + patch( + "litellm.proxy.auth.user_api_key_auth.get_team_object", + AsyncMock(return_value=team_obj), + ), + patch( + "litellm.proxy.auth.user_api_key_auth._enforce_key_and_fallback_model_access", + new_callable=AsyncMock, + ), + patch( + "litellm.proxy.auth.user_api_key_auth._return_user_api_key_auth_obj", + new_callable=AsyncMock, + return_value=team_token, + ), + patch( + "litellm.proxy.auth.auth_exception_handler.seed_request_identity", + ), + ): + await _user_api_key_auth_builder( + request=request, + api_key=f"Bearer {api_key}", + azure_api_key_header="", + anthropic_api_key_header=None, + google_ai_studio_api_key_header=None, + azure_apim_header=None, + request_data={}, + ) + finally: + for k, v in originals.items(): + setattr(_proxy_server_mod, k, v) + + served = cache.get_cache( + key=f"team_id:{team_id}", model_type=LiteLLM_TeamTableCachedObj + ) + assert served is not None and served.team_id == team_id + assert cache.get_cache(key=team_id) is None + assert cache.get_cache(key=None) is None diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py index 6ec793a1bb0..399442a5f71 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py @@ -1021,3 +1021,171 @@ class TestMultimodalSupport: call_args = mock_post.call_args json_payload = call_args.kwargs["json"] assert isinstance(json_payload["structured_messages"], list) + + +class TestToolSupport: + """Test tool handling in guardrail requests""" + + @pytest.mark.asyncio + async def test_builtin_tools_without_function_block_do_not_crash( + self, generic_guardrail + ): + """Built-in tools (code_interpreter, file_search) have no `function` block. + + Regression for a 500 where serializing them raised a Pydantic + ValidationError because the tool schema required `function`. The full + tool, including built-in tool config, must reach the guardrail intact. + """ + tools = [ + {"type": "function", "function": {"name": "get_weather", "parameters": {}}}, + {"type": "code_interpreter"}, + { + "type": "file_search", + "vector_store_ids": ["vs_1"], + "max_num_results": 5, + }, + ] + + mock_response = MagicMock() + mock_response.json.return_value = {"action": "NONE", "texts": ["hi"]} + mock_response.raise_for_status = MagicMock() + + with patch.object( + generic_guardrail.async_handler, "post", return_value=mock_response + ) as mock_post: + await generic_guardrail.apply_guardrail( + inputs={"texts": ["hi"], "tools": tools}, + request_data={}, + input_type="request", + ) + + forwarded_tools = mock_post.call_args.kwargs["json"]["tools"] + + assert forwarded_tools == tools + + +class TestFailOnError: + """Test fail_on_error: complete fail-open on any guardrail error""" + + @pytest.fixture + def fail_open_guardrail(self): + return GenericGuardrailAPI( + api_base="https://api.test.guardrail.com", + guardrail_name="test-fail-open-guardrail", + event_hook="pre_call", + default_on=True, + fail_on_error=False, + ) + + @pytest.mark.asyncio + async def test_endpoint_error_continues_when_fail_on_error_false( + self, fail_open_guardrail + ): + """A non-unreachable endpoint error (HTTP 400) is swallowed and the request proceeds unchanged.""" + error = httpx.HTTPStatusError( + "bad request", request=MagicMock(), response=MagicMock(status_code=400) + ) + with patch.object( + fail_open_guardrail.async_handler, "post", side_effect=error + ): + result = await fail_open_guardrail.apply_guardrail( + inputs={"texts": ["hi"]}, + request_data={}, + input_type="request", + ) + + assert result == {"texts": ["hi"]} + + @pytest.mark.asyncio + async def test_internal_error_continues_without_calling_endpoint( + self, fail_open_guardrail + ): + """An error while building the request (here: invalid input_type) fails open too. + + Proves the request construction runs inside the protected block: the + endpoint is never called, yet the request still proceeds unchanged. + """ + with patch.object(fail_open_guardrail.async_handler, "post") as mock_post: + result = await fail_open_guardrail.apply_guardrail( + inputs={"texts": ["hi"]}, + request_data={}, + input_type="bogus", # type: ignore[arg-type] + ) + + mock_post.assert_not_called() + assert result == {"texts": ["hi"]} + + @pytest.mark.asyncio + async def test_valid_block_still_blocks_when_fail_on_error_false( + self, fail_open_guardrail + ): + """Only a valid response acts: a BLOCKED decision still raises even with fail_on_error=False.""" + mock_response = MagicMock() + mock_response.json.return_value = { + "action": "BLOCKED", + "blocked_reason": "policy violation", + } + mock_response.raise_for_status = MagicMock() + + with patch.object( + fail_open_guardrail.async_handler, "post", return_value=mock_response + ): + with pytest.raises(GuardrailRaisedException): + await fail_open_guardrail.apply_guardrail( + inputs={"texts": ["hi"]}, + request_data={}, + input_type="request", + ) + + @pytest.mark.asyncio + async def test_endpoint_error_raises_by_default(self, generic_guardrail): + """Default fail_on_error=True keeps blocking on a non-unreachable endpoint error.""" + error = httpx.HTTPStatusError( + "bad request", request=MagicMock(), response=MagicMock(status_code=400) + ) + with patch.object(generic_guardrail.async_handler, "post", side_effect=error): + with pytest.raises(Exception, match="Generic Guardrail API failed"): + await generic_guardrail.apply_guardrail( + inputs={"texts": ["hi"]}, + request_data={}, + input_type="request", + ) + + @pytest.mark.asyncio + async def test_response_path_continues_when_fail_on_error_false( + self, fail_open_guardrail + ): + """fail_on_error governs the response path identically to the request path.""" + error = httpx.HTTPStatusError( + "bad request", request=MagicMock(), response=MagicMock(status_code=400) + ) + with patch.object( + fail_open_guardrail.async_handler, "post", side_effect=error + ): + result = await fail_open_guardrail.apply_guardrail( + inputs={"texts": ["model output"]}, + request_data={}, + input_type="response", + ) + + assert result == {"texts": ["model output"]} + + @pytest.mark.asyncio + async def test_response_path_valid_block_still_blocks(self, fail_open_guardrail): + """On the response path too, a valid BLOCKED decision raises despite fail_on_error=False.""" + mock_response = MagicMock() + mock_response.json.return_value = { + "action": "BLOCKED", + "blocked_reason": "policy violation", + } + mock_response.raise_for_status = MagicMock() + + with patch.object( + fail_open_guardrail.async_handler, "post", return_value=mock_response + ): + with pytest.raises(GuardrailRaisedException): + await fail_open_guardrail.apply_guardrail( + inputs={"texts": ["model output"]}, + request_data={}, + input_type="response", + ) diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py index a649bc7225e..ecce92e7324 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -1270,6 +1270,95 @@ async def test_available_team_self_join_blocks_admin_role_in_member_list(): assert exc_info.value.status_code == 403 +@pytest.mark.asyncio +@pytest.mark.parametrize( + "budget_control", + [ + {"max_budget_in_team": 1000.0}, + {"budget_duration": "1h"}, + {"allowed_models": ["gpt-4o"]}, + ], +) +async def test_available_team_self_join_blocks_member_budget_controls(budget_control): + """A self-joining non-admin must not be able to set their own per-member + budget or model controls via the available-team bypass; only proxy/team/org + admins may. Without this guard a self-joiner could shorten their budget + reset window or widen their cap/model scope past the team default.""" + from litellm.proxy._types import Member, TeamMemberAddRequest + from litellm.proxy.management_endpoints.team_endpoints import ( + _validate_team_member_add_permissions, + ) + + user = UserAPIKeyAuth(user_id="alice", user_role=LitellmUserRoles.INTERNAL_USER) + team = MagicMock(spec=LiteLLM_TeamTable) + team.team_id = "public-team" + team.members_with_roles = [] + team.organization_id = None + + data = TeamMemberAddRequest( + team_id="public-team", + member=Member(role="user", user_id="alice"), + **budget_control, + ) + + with ( + patch( + "litellm.proxy.management_endpoints.team_endpoints._is_user_team_admin", + return_value=False, + ), + patch( + "litellm.proxy.management_endpoints.team_endpoints._is_available_team", + return_value=True, + ), + pytest.raises(HTTPException) as exc_info, + ): + await _validate_team_member_add_permissions( + user_api_key_dict=user, + complete_team_data=team, + data=data, + ) + + assert exc_info.value.status_code == 403 + assert "admin-only" in str(exc_info.value.detail).lower() + + +@pytest.mark.asyncio +async def test_available_team_self_join_allows_no_budget_controls(): + """The clean self-join (no per-member budget/model controls) must still be + permitted, so the new guard does not break the legitimate join path.""" + from litellm.proxy._types import Member, TeamMemberAddRequest + from litellm.proxy.management_endpoints.team_endpoints import ( + _validate_team_member_add_permissions, + ) + + user = UserAPIKeyAuth(user_id="alice", user_role=LitellmUserRoles.INTERNAL_USER) + team = MagicMock(spec=LiteLLM_TeamTable) + team.team_id = "public-team" + team.members_with_roles = [] + team.organization_id = None + + data = TeamMemberAddRequest( + team_id="public-team", + member=Member(role="user", user_id="alice"), + ) + + with ( + patch( + "litellm.proxy.management_endpoints.team_endpoints._is_user_team_admin", + return_value=False, + ), + patch( + "litellm.proxy.management_endpoints.team_endpoints._is_available_team", + return_value=True, + ), + ): + await _validate_team_member_add_permissions( + user_api_key_dict=user, + complete_team_data=team, + data=data, + ) + + @pytest.mark.asyncio async def test_update_team_member_permissions_blocks_non_admin_via_available_team( mock_db_client, @@ -1385,6 +1474,7 @@ async def test_process_team_members_single_member(): team_id="test-team-123", default_team_budget_id="budget-123", allowed_models=None, + budget_duration=None, ) diff --git a/tests/test_litellm/proxy/management_helpers/test_management_helpers_utils.py b/tests/test_litellm/proxy/management_helpers/test_management_helpers_utils.py index 463aba6f744..dbbdca65cc6 100644 --- a/tests/test_litellm/proxy/management_helpers/test_management_helpers_utils.py +++ b/tests/test_litellm/proxy/management_helpers/test_management_helpers_utils.py @@ -1,6 +1,7 @@ import json import os import sys +from datetime import datetime, timezone from litellm._uuid import uuid from unittest.mock import AsyncMock, MagicMock @@ -283,6 +284,81 @@ async def test_add_new_member_clones_default_team_budget_id(): assert create_data["budget_id"] == test_cloned_budget_id +@pytest.mark.asyncio +async def test_add_new_member_budget_duration_only_clones_default_max_budget(): + """When only a budget_duration is given and the team has a default member + budget, the member must clone the default (keeping its max_budget) and just + override the reset window. Creating a fresh duration-only row instead would + silently drop the team default's cap, leaving the member uncapped.""" + from litellm.proxy._types import LitellmUserRoles + + new_member = Member(user_id="dur-clone-user", role="user") + user_api_key_dict = UserAPIKeyAuth( + user_id="admin_user", user_role=LitellmUserRoles.PROXY_ADMIN + ) + + mock_prisma_client = AsyncMock() + mock_user_response = MagicMock() + mock_user_response.model_dump.return_value = { + "user_id": "dur-clone-user", + "user_email": None, + "teams": ["team-dc"], + "user_role": "internal_user", + } + mock_prisma_client.db.litellm_usertable.upsert = AsyncMock( + return_value=mock_user_response + ) + mock_default_budget_row = MagicMock() + mock_default_budget_row.model_dump.return_value = { + "budget_id": "default-dc", + "max_budget": 100.0, + "soft_budget": None, + "max_parallel_requests": None, + "tpm_limit": 1000, + "rpm_limit": None, + "model_max_budget": None, + "budget_duration": "1d", + "allowed_models": [], + } + mock_prisma_client.db.litellm_budgettable.find_unique = AsyncMock( + return_value=mock_default_budget_row + ) + mock_cloned_budget_row = MagicMock() + mock_cloned_budget_row.budget_id = "cloned-dc" + mock_prisma_client.db.litellm_budgettable.create = AsyncMock( + return_value=mock_cloned_budget_row + ) + mock_team_membership_response = MagicMock() + mock_team_membership_response.model_dump.return_value = { + "team_id": "team-dc", + "user_id": "dur-clone-user", + "budget_id": "cloned-dc", + "litellm_budget_table": None, + } + mock_prisma_client.db.litellm_teammembership.create = AsyncMock( + return_value=mock_team_membership_response + ) + + await add_new_member( + new_member=new_member, + max_budget_in_team=None, + prisma_client=mock_prisma_client, + team_id="team-dc", + user_api_key_dict=user_api_key_dict, + litellm_proxy_admin_name="test_admin", + default_team_budget_id="default-dc", + budget_duration="7d", + ) + + mock_prisma_client.db.litellm_budgettable.create.assert_called_once() + cloned_create_data = ( + mock_prisma_client.db.litellm_budgettable.create.call_args.kwargs["data"] + ) + assert cloned_create_data["max_budget"] == 100.0 # kept from the team default + assert cloned_create_data["budget_duration"] == "7d" # overridden by the caller + assert cloned_create_data["budget_reset_at"] > datetime.now(timezone.utc) + + @pytest.mark.asyncio async def test_add_new_member_no_budget_when_no_default_and_no_max_budget(): """ @@ -434,6 +510,130 @@ async def test_add_new_member_creates_new_budget_when_max_budget_provided(): assert create_data["budget_id"] == test_new_budget_id +@pytest.mark.asyncio +async def test_add_new_member_persists_budget_duration(): + """Regression for the member_add half of the recurring-member-budget gap: + a budget_duration passed to add_new_member must be written to the new + member budget along with a future budget_reset_at, so the per-member budget + recurs instead of acting as a lifetime cap.""" + from litellm.proxy._types import LitellmUserRoles + + new_member = Member(user_id="user-dur", role="user") + user_api_key_dict = UserAPIKeyAuth( + user_id="admin_user", user_role=LitellmUserRoles.PROXY_ADMIN + ) + + mock_prisma_client = AsyncMock() + mock_user_response = MagicMock() + mock_user_response.model_dump.return_value = { + "user_id": "user-dur", + "user_email": None, + "teams": ["team-dur"], + "user_role": "internal_user", + } + mock_prisma_client.db.litellm_usertable.upsert = AsyncMock( + return_value=mock_user_response + ) + mock_budget_response = MagicMock() + mock_budget_response.budget_id = "budget-dur" + mock_prisma_client.db.litellm_budgettable.create = AsyncMock( + return_value=mock_budget_response + ) + mock_team_membership_response = MagicMock() + mock_team_membership_response.model_dump.return_value = { + "team_id": "team-dur", + "user_id": "user-dur", + "budget_id": "budget-dur", + "litellm_budget_table": None, + } + mock_prisma_client.db.litellm_teammembership.create = AsyncMock( + return_value=mock_team_membership_response + ) + + await add_new_member( + new_member=new_member, + max_budget_in_team=10.0, + prisma_client=mock_prisma_client, + team_id="team-dur", + user_api_key_dict=user_api_key_dict, + litellm_proxy_admin_name="test_admin", + default_team_budget_id=None, + allowed_models=["gpt-4o-mini"], + budget_duration="30d", + ) + + mock_prisma_client.db.litellm_budgettable.create.assert_called_once() + budget_data = mock_prisma_client.db.litellm_budgettable.create.call_args.kwargs[ + "data" + ] + assert budget_data["max_budget"] == 10.0 + assert budget_data["allowed_models"] == ["gpt-4o-mini"] + assert budget_data["budget_duration"] == "30d" + reset_at = budget_data["budget_reset_at"] + assert isinstance(reset_at, datetime) + assert reset_at.tzinfo is not None + assert reset_at > datetime.now(timezone.utc) + + +@pytest.mark.asyncio +async def test_add_new_member_persists_budget_duration_without_max_budget(): + """budget_duration alone must still create a member budget; otherwise an + explicit recurring window passed without a cap would be silently dropped.""" + from litellm.proxy._types import LitellmUserRoles + + new_member = Member(user_id="user-dur2", role="user") + user_api_key_dict = UserAPIKeyAuth( + user_id="admin_user", user_role=LitellmUserRoles.PROXY_ADMIN + ) + + mock_prisma_client = AsyncMock() + mock_user_response = MagicMock() + mock_user_response.model_dump.return_value = { + "user_id": "user-dur2", + "user_email": None, + "teams": ["team-dur2"], + "user_role": "internal_user", + } + mock_prisma_client.db.litellm_usertable.upsert = AsyncMock( + return_value=mock_user_response + ) + mock_budget_response = MagicMock() + mock_budget_response.budget_id = "budget-dur2" + mock_prisma_client.db.litellm_budgettable.create = AsyncMock( + return_value=mock_budget_response + ) + mock_team_membership_response = MagicMock() + mock_team_membership_response.model_dump.return_value = { + "team_id": "team-dur2", + "user_id": "user-dur2", + "budget_id": "budget-dur2", + "litellm_budget_table": None, + } + mock_prisma_client.db.litellm_teammembership.create = AsyncMock( + return_value=mock_team_membership_response + ) + + _, result_team_membership = await add_new_member( + new_member=new_member, + max_budget_in_team=None, + prisma_client=mock_prisma_client, + team_id="team-dur2", + user_api_key_dict=user_api_key_dict, + litellm_proxy_admin_name="test_admin", + default_team_budget_id=None, + budget_duration="7d", + ) + + mock_prisma_client.db.litellm_budgettable.create.assert_called_once() + budget_data = mock_prisma_client.db.litellm_budgettable.create.call_args.kwargs[ + "data" + ] + assert budget_data["budget_duration"] == "7d" + assert budget_data["budget_reset_at"] > datetime.now(timezone.utc) + assert result_team_membership is not None + assert result_team_membership.budget_id == "budget-dur2" + + @pytest.mark.asyncio async def test_add_new_member_with_user_email_clones_default_budget(): """ diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py index de5250fc5f3..b37bb18744d 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py @@ -2,6 +2,7 @@ import asyncio import datetime import json import os +import re import sys from datetime import timezone @@ -60,31 +61,115 @@ def _filter_logs_by_date_range(logs, where): return filtered +def _reconstruct_ui_where_from_sql(sql_query, params): + """ + Rebuild the Prisma-style ``where`` dict the filter_fns below expect from the + raw SQL + params the endpoint emits. + + ``ui_view_spend_logs`` folds the total into the page query via + ``COUNT(*) OVER ()`` and no longer issues a separate ``count(where=...)`` + call, so the mock derives the active filter from the one query it sees + instead of from the (now absent) count call. + """ + where: dict = {} + clause = re.search(r"WHERE (.*) ORDER BY", sql_query, re.DOTALL) + if clause is None: + return where + + def _iso(value): + return value.isoformat() if hasattr(value, "isoformat") else str(value) + + eq_cols = { + "team_id": "team_id", + '"user"': "user", + "api_key": "api_key", + "request_id": "request_id", + "model": "model", + "model_id": "model_id", + "model_group": "model_group", + "end_user": "end_user", + } + date_bounds: dict = {} + metadata_conds: list = [] + for cond in (c.strip() for c in clause.group(1).split(" AND ")): + gte = re.search(r'"startTime" >= \(\$(\d+)', cond) + lte = re.search(r'"startTime" <= \(\$(\d+)', cond) + alias = re.search(r"user_api_key_alias' LIKE \$(\d+)", cond) + code = re.search(r"error_code' = \$(\d+)", cond) + msg = re.search(r"error_message' LIKE \$(\d+)", cond) + status = re.fullmatch(r"status = \$(\d+)", cond) + if gte: + date_bounds["gte"] = _iso(params[int(gte.group(1)) - 1]) + elif lte: + date_bounds["lte"] = _iso(params[int(lte.group(1)) - 1]) + elif "OR team_id = ANY" in cond: + where["OR"] = where.get("OR", []) + [{"multi_team": True}] + elif "status = 'success'" in cond: + where["OR"] = where.get("OR", []) + [{"status": "success"}] + elif status: + where["status"] = {"equals": params[int(status.group(1)) - 1]} + elif alias: + metadata_conds.append( + { + "path": ["user_api_key_alias"], + "string_contains": str(params[int(alias.group(1)) - 1]).strip("%"), + } + ) + elif code: + metadata_conds.append( + { + "path": ["error_information", "error_code"], + "equals": params[int(code.group(1)) - 1], + } + ) + elif msg: + metadata_conds.append( + { + "path": ["error_information", "error_message"], + "string_contains": str(params[int(msg.group(1)) - 1]).strip("%"), + } + ) + else: + for sql_col, key in eq_cols.items(): + eq = re.fullmatch(rf"{re.escape(sql_col)} = \$(\d+)", cond) + if eq: + where[key] = params[int(eq.group(1)) - 1] + break + + if date_bounds: + where["startTime"] = date_bounds + if len(metadata_conds) == 1: + where["metadata"] = metadata_conds[0] + elif len(metadata_conds) > 1: + where["AND"] = [{"metadata": cond} for cond in metadata_conds] + return where + + def make_ui_spend_logs_mock_prisma(mock_spend_logs, filter_fn, team_lookup_fn=None): """ Create a MockPrismaClient for /spend/logs/ui endpoint tests. Args: mock_spend_logs: List of mock spend log dicts. - filter_fn: Callable[[dict], list] - receives where_conditions from count(), - returns the filtered list of logs for that query. + filter_fn: Callable[[dict], list] - receives the reconstructed + where_conditions, returns the filtered list of logs. team_lookup_fn: Optional async callable for team RBAC (find_unique). If provided, adds litellm_teamtable to db. """ - filtered_holder = [] class MockDB: async def count(self, *args, **kwargs): - where = kwargs.get("where", {}) - filtered = filter_fn(where) - filtered_holder.clear() - filtered_holder.extend(filtered) - return len(filtered) + return len(filter_fn(kwargs.get("where", {}))) async def query_raw(self, sql_query, *params): + filtered = filter_fn(_reconstruct_ui_where_from_sql(sql_query, params)) page_size = params[-2] if len(params) >= 2 else 50 skip = params[-1] if len(params) >= 1 else 0 - return filtered_holder[skip : skip + page_size] + total = len(filtered) + return [ + {**row, "total_count": total} + for row in filtered[skip : skip + page_size] + ] class MockPrismaClient: def __init__(self): @@ -608,7 +693,10 @@ async def test_ui_view_spend_logs_sort_by_and_sort_order( sorted_logs = _sort_logs(base_logs, order) page_size = params[-2] if len(params) >= 2 else 50 skip = params[-1] if len(params) >= 1 else 0 - return sorted_logs[skip : skip + page_size] + return [ + {**row, "total_count": len(base_logs)} + for row in sorted_logs[skip : skip + page_size] + ] class MockPrismaClient: def __init__(self): @@ -748,7 +836,10 @@ async def test_ui_view_spend_logs_sort_by_request_duration_ms(client, monkeypatc ) page_size = params[-2] if len(params) >= 2 else 50 skip = params[-1] if len(params) >= 1 else 0 - return sorted_logs[skip : skip + page_size] + return [ + {**row, "total_count": len(base_logs)} + for row in sorted_logs[skip : skip + page_size] + ] class MockPrismaClient: def __init__(self): @@ -846,7 +937,10 @@ async def test_ui_view_spend_logs_sort_by_model( ) page_size = params[-2] if len(params) >= 2 else 50 skip = params[-1] if len(params) >= 1 else 0 - return sorted_logs[skip : skip + page_size] + return [ + {**row, "total_count": len(base_logs)} + for row in sorted_logs[skip : skip + page_size] + ] class MockPrismaClient: def __init__(self): @@ -957,7 +1051,7 @@ async def test_ui_view_spend_logs_sort_by_ttft_ms(client, monkeypatch): page_size = params[-2] if len(params) >= 2 else 50 skip = params[-1] if len(params) >= 1 else 0 return [ - {k: v for k, v in row.items() if k != "_ttft_ms"} + {**{k: v for k, v in row.items() if k != "_ttft_ms"}, "total_count": len(base_logs)} for row in sorted_logs[skip : skip + page_size] ] @@ -3668,7 +3762,7 @@ async def test_ui_view_spend_logs_rehydrates_metadata_jsonb_text(client, monkeyp return 1 async def mock_query_raw(sql_query, *params): - return [raw_row] + return [{**raw_row, "total_count": 1}] class MockPrismaClient: def __init__(self): @@ -3754,7 +3848,7 @@ async def test_ui_view_spend_logs_metadata_invalid_json_falls_back_to_empty_dict return 1 async def mock_query_raw(sql_query, *params): - return [raw_row] + return [{**raw_row, "total_count": 1}] class MockPrismaClient: def __init__(self): diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_query_optimization.py b/tests/test_litellm/proxy/spend_tracking/test_spend_query_optimization.py index fbac71e6372..5b793adbbb4 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_query_optimization.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_query_optimization.py @@ -222,3 +222,154 @@ async def test_spend_logs_ui_wraps_params_in_at_time_zone_utc(monkeypatch): "/spend/logs/ui must wrap both `startTime` bounds with " f"`AT TIME ZONE 'UTC'`. SQL was:\n{sql}" ) + + +@pytest.mark.asyncio +async def test_spend_logs_ui_folds_count_into_window_function(monkeypatch): + """ + /spend/logs/ui must not issue a separate `COUNT(*)` round trip to compute + the total. On sharded engines like YugabyteDB a standalone `COUNT(*)` is a + distributed RPC that contacts every tablet and times out regardless of row + count, so the logs tab 500s (LIT-4027). The total is folded into the page + query via `COUNT(*) OVER ()` and read off the returned rows instead. + """ + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + from litellm.proxy.spend_tracking.spend_management_endpoints import ( + ui_view_spend_logs, + ) + + rows = [ + {"request_id": "req-1", "metadata": "{}", "session_id": None, "total_count": 137}, + {"request_id": "req-2", "metadata": "{}", "session_id": None, "total_count": 137}, + ] + + mock_prisma = MagicMock() + mock_prisma.db = MagicMock() + mock_prisma.db.query_raw = AsyncMock(return_value=rows) + mock_prisma.db.litellm_spendlogs = MagicMock() + mock_prisma.db.litellm_spendlogs.count = AsyncMock(return_value=0) + + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) + + auth = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin") + mock_request = MagicMock() + mock_request.url.path = "/spend/logs/ui" + + response = await ui_view_spend_logs( + request=mock_request, + api_key=None, + user_id=None, + request_id=None, + start_date="2026-02-16 00:00:00", + end_date="2026-02-16 23:59:59", + page=1, + page_size=50, + sort_by="startTime", + sort_order="desc", + user_api_key_dict=auth, + ) + + mock_prisma.db.litellm_spendlogs.count.assert_not_called() + + sql = mock_prisma.db.query_raw.call_args[0][0] + assert "COUNT(*) OVER ()" in sql, ( + "the page query must carry a window-function count so a separate " + f"distributed COUNT(*) is avoided. SQL was:\n{sql}" + ) + + assert response["total"] == 137 + assert response["total_pages"] == (137 + 50 - 1) // 50 + + for row in response["data"]: + assert "total_count" not in row, ( + "the window-function helper column must be stripped before " + "serialising rows" + ) + + +@pytest.mark.asyncio +async def test_spend_logs_ui_empty_page_reports_zero_total(monkeypatch): + """ + When a page matches no rows the window-function count row is absent, so the + total must fall back to zero without issuing a separate `COUNT(*)`. + """ + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + from litellm.proxy.spend_tracking.spend_management_endpoints import ( + ui_view_spend_logs, + ) + + mock_prisma = MagicMock() + mock_prisma.db = MagicMock() + mock_prisma.db.query_raw = AsyncMock(return_value=[]) + mock_prisma.db.litellm_spendlogs = MagicMock() + mock_prisma.db.litellm_spendlogs.count = AsyncMock(return_value=0) + + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) + + auth = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin") + mock_request = MagicMock() + mock_request.url.path = "/spend/logs/ui" + + response = await ui_view_spend_logs( + request=mock_request, + api_key=None, + user_id=None, + request_id=None, + start_date="2026-02-16 00:00:00", + end_date="2026-02-16 23:59:59", + page=1, + page_size=50, + sort_by="startTime", + sort_order="desc", + user_api_key_dict=auth, + ) + + mock_prisma.db.litellm_spendlogs.count.assert_not_called() + assert response["total"] == 0 + assert response["total_pages"] == 0 + assert response["data"] == [] + + +@pytest.mark.asyncio +async def test_spend_logs_ui_out_of_range_page_falls_back_to_count(monkeypatch): + """ + An out-of-range page (offset past the last matching row) returns no rows, so + the window-function count is unavailable. The total must not collapse to zero + there; it falls back to a direct count so total/total_pages stay accurate. + This fallback only fires off the hot path (page > 1 with an empty result), so + the YugabyteDB timeout the fix removes from page 1 stays removed. + """ + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + from litellm.proxy.spend_tracking.spend_management_endpoints import ( + ui_view_spend_logs, + ) + + mock_prisma = MagicMock() + mock_prisma.db = MagicMock() + mock_prisma.db.query_raw = AsyncMock(return_value=[]) + mock_prisma.db.litellm_spendlogs = MagicMock() + mock_prisma.db.litellm_spendlogs.count = AsyncMock(return_value=7) + + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) + + auth = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin") + mock_request = MagicMock() + mock_request.url.path = "/spend/logs/ui" + + response = await ui_view_spend_logs( + request=mock_request, + api_key=None, + user_id=None, + request_id=None, + start_date="2026-02-16 00:00:00", + end_date="2026-02-16 23:59:59", + page=99, + page_size=2, + sort_by="startTime", + sort_order="desc", + user_api_key_dict=auth, + ) + + mock_prisma.db.litellm_spendlogs.count.assert_called_once() + assert response["total"] == 7 + assert response["total_pages"] == (7 + 2 - 1) // 2 diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py b/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py index 1ff9fbf8d83..b9f209e5251 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py +++ b/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py @@ -4,7 +4,7 @@ Covers ``_should_use_guardrail_load_balancing``, ``_execute_guardrail_hook``, ``_execute_guardrail_with_load_balancing``, ``_process_guardrail_callback``, ``_process_prompt_template``, ``_process_guardrail_metadata``, ``_maybe_execute_pipelines``, ``_handle_pipeline_result``, -``_run_guardrail_task_with_enrichment``. +``_run_guardrail_with_metrics``, ``_emit_guardrail_metrics``. """ from __future__ import annotations @@ -21,6 +21,7 @@ from litellm.integrations.custom_guardrail import ( CustomGuardrail, ModifyResponseException, ) +from litellm.integrations.prometheus import PrometheusLogger from litellm.proxy.utils import ProxyLogging from litellm.types.guardrails import GuardrailEventHooks @@ -425,23 +426,42 @@ def test_handle_pipeline_result_unknown_action_returns_data(): # --------------------------------------------------------------------------- -# _run_guardrail_task_with_enrichment +# _run_guardrail_with_metrics # --------------------------------------------------------------------------- +def _prometheus_callback() -> MagicMock: + """Stand-in PrometheusLogger that records ``_record_guardrail_metrics`` calls. + + ``MagicMock(spec=PrometheusLogger)`` passes the ``isinstance`` check inside + ``_emit_guardrail_metrics`` while letting us capture the recorded labels. + """ + return MagicMock(spec=PrometheusLogger) + + @pytest.mark.asyncio -async def test_run_guardrail_task_with_enrichment_passes_result(): +async def test_run_guardrail_with_metrics_passes_result_and_records_success(monkeypatch): async def task(): return {"a": 1, "b": 2, "c": 3} - out = await ProxyLogging._run_guardrail_task_with_enrichment( - callback=MagicMock(guardrail_name="g"), coro=task() + prom = _prometheus_callback() + monkeypatch.setattr(litellm, "callbacks", [prom]) + + out = await ProxyLogging._run_guardrail_with_metrics( + callback=MagicMock(guardrail_name="g"), coro=task(), hook_type="during_call" ) + assert out == {"a": 1, "b": 2, "c": 3} + recorded = prom._record_guardrail_metrics.call_args.kwargs + assert recorded["guardrail_name"] == "g" + assert recorded["status"] == "success" + assert recorded["error_type"] is None + assert recorded["hook_type"] == "during_call" + assert recorded["latency_seconds"] >= 0 @pytest.mark.asyncio -async def test_run_guardrail_task_with_enrichment_enriches_http_exception_raises(): +async def test_run_guardrail_with_metrics_records_error_and_enriches(monkeypatch): detail = {"error": "blocked"} async def task(): @@ -450,9 +470,79 @@ async def test_run_guardrail_task_with_enrichment_enriches_http_exception_raises cb = MagicMock() cb.guardrail_name = "presidio" cb.event_hook = "pre_call" + prom = _prometheus_callback() + monkeypatch.setattr(litellm, "callbacks", [prom]) + with pytest.raises(HTTPException): - await ProxyLogging._run_guardrail_task_with_enrichment(callback=cb, coro=task()) + await ProxyLogging._run_guardrail_with_metrics( + callback=cb, coro=task(), hook_type="post_call" + ) + assert detail["guardrail_name"] == "presidio" + recorded = prom._record_guardrail_metrics.call_args.kwargs + assert recorded["status"] == "error" + assert recorded["error_type"] == "HTTPException" + assert recorded["hook_type"] == "post_call" + + +# --------------------------------------------------------------------------- +# during_call / post_call phases emit the latency metric (LIT-3999 regression) +# --------------------------------------------------------------------------- + + +def _moderation_guardrail() -> MagicMock: + cb = MagicMock(spec=CustomGuardrail) + cb.__class__ = CustomGuardrail + cb.guardrail_name = "g" + cb.event_hook = GuardrailEventHooks.during_call + cb.use_native_during_call_hook = False + cb.should_run_guardrail = MagicMock(return_value=True) + cb.async_moderation_hook = AsyncMock(return_value=None) + cb.async_post_call_success_hook = AsyncMock(return_value=None) + return cb + + +@pytest.mark.asyncio +async def test_during_call_hook_records_latency_metric( + proxy_logging, make_user_api_key_auth, monkeypatch +): + cb = _moderation_guardrail() + prom = _prometheus_callback() + monkeypatch.setattr(litellm, "callbacks", [prom, cb]) + + await proxy_logging.during_call_hook( + data={"model": "m"}, + user_api_key_dict=make_user_api_key_auth(), + call_type="completion", + ) + + cb.async_moderation_hook.assert_awaited_once() + recorded = prom._record_guardrail_metrics.call_args.kwargs + assert recorded["hook_type"] == "during_call" + assert recorded["guardrail_name"] == "g" + assert recorded["status"] == "success" + + +@pytest.mark.asyncio +async def test_post_call_success_hook_records_latency_metric( + proxy_logging, make_user_api_key_auth, monkeypatch +): + cb = _moderation_guardrail() + prom = _prometheus_callback() + monkeypatch.setattr(litellm, "callbacks", [prom, cb]) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + + await proxy_logging.post_call_success_hook( + data={"model": "m"}, + response=litellm.ModelResponse(), + user_api_key_dict=make_user_api_key_auth(), + ) + + cb.async_post_call_success_hook.assert_awaited_once() + recorded = prom._record_guardrail_metrics.call_args.kwargs + assert recorded["hook_type"] == "post_call" + assert recorded["guardrail_name"] == "g" + assert recorded["status"] == "success" # --------------------------------------------------------------------------- diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index cdc91371dfe..43be9ad6e39 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -30842,6 +30842,11 @@ export interface components { * @description List of models this team member can access. If not set, inherits the team's default_team_member_models or all team models. */ allowed_models?: string[] | null; + /** + * Budget Duration + * @description Duration after which this team member's budget resets (e.g. '1h', '24h', '7d', '30d'). If not set, the budget never resets. + */ + budget_duration?: string | null; /** * Max Budget In Team * @description Maximum budget allocated to this user within the team. If not set, user has unlimited budget within team limits