Merge pull request #31477 from BerriAI/litellm_internal_staging
Some checks are pending
CodeQL / Analyze (actions) (push) Waiting to run
CodeQL / Analyze (javascript-typescript) (push) Waiting to run
CodeQL / Analyze (python) (push) Waiting to run
CodSpeed Benchmarks / benchmarks (push) Waiting to run
Helm unit test / unit-test (push) Waiting to run
Scorecard supply-chain security / Scorecard analysis (push) Waiting to run
GitHub Actions Security Analysis / zizmor (push) Waiting to run

chore(ci): promote internal staging to main
This commit is contained in:
yuneng-jiang 2026-06-26 14:20:27 -07:00 committed by GitHub
commit a311f11b00
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
25 changed files with 1381 additions and 149 deletions

View file

@ -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"],

View file

@ -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",

View file

@ -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.",

View file

@ -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

View file

@ -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,

View file

@ -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(

View file

@ -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(

View file

@ -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(

View file

@ -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:

View file

@ -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

View file

@ -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(

View file

@ -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

View file

@ -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",

View file

@ -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"

View file

@ -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")

View file

@ -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()

View file

@ -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)

View file

@ -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

View file

@ -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",
)

View file

@ -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,
)

View file

@ -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():
"""

View file

@ -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):

View file

@ -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

View file

@ -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"
# ---------------------------------------------------------------------------

View file

@ -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