scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-08-28 05:25:59 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 36

Merge remote-tracking branch 'origin/litellm_internal_staging' into litellm_reasoning_effort_capability_v2

Browse source 
# Conflicts:
#	tests/test_litellm/test_router.py
This commit is contained in:
mateo-berri 2026-08-25 12:07:05 -07:00
commit 5b2d1874b0
61 changed files with 2591 additions and 133 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

8
basedpyright-code-budget.json
View file

@ -1,6 +1,6 @@
{
"reportAny": {
"limit": 19955
"limit": 19949
},
"reportArgumentType": {
"limit": 2566
@ -54,7 +54,7 @@
"limit": 0
},
"reportMissingParameterType": {
"limit": 5663
"limit": 5661
},
"reportMissingTypeArgument": {
"limit": 15555
@ -105,10 +105,10 @@
"limit": 109
},
"reportUnknownMemberType": {
"limit": 39011
"limit": 39009
},
"reportUnknownParameterType": {
"limit": 19885
"limit": 19883
},
"reportUnknownVariableType": {
"limit": 30569

3
litellm/__init__.py
View file

@ -1801,6 +1801,9 @@ if TYPE_CHECKING:
from .llms.gemini.interactions.transformation import (
GoogleAIStudioInteractionsConfig as GoogleAIStudioInteractionsConfig,
)
from .llms.vertex_ai.interactions.transformation import (
VertexAIInteractionsConfig as VertexAIInteractionsConfig,
)
from .llms.openai.chat.o_series_transformation import (
OpenAIOSeriesConfig as OpenAIOSeriesConfig,
OpenAIOSeriesConfig as OpenAIO1Config,

5
litellm/_lazy_imports_registry.py
View file

@ -242,6 +242,7 @@ LLM_CONFIG_NAMES: Final = (
"OpenRouterResponsesAPIConfig",
"BedrockMantleResponsesAPIConfig",
"GoogleAIStudioInteractionsConfig",
"VertexAIInteractionsConfig",
"OpenAIOSeriesConfig",
"AnthropicSkillsConfig",
"BaseSkillsAPIConfig",
@ -977,6 +978,10 @@ _LLM_CONFIGS_IMPORT_MAP: Final = {
".llms.gemini.interactions.transformation",
"GoogleAIStudioInteractionsConfig",
),
"VertexAIInteractionsConfig": (
".llms.vertex_ai.interactions.transformation",
"VertexAIInteractionsConfig",
),
"OpenAIOSeriesConfig": (
".llms.openai.chat.o_series_transformation",
"OpenAIOSeriesConfig",

1
litellm/constants.py
View file

@ -750,6 +750,7 @@ openai_compatible_endpoints: Final[list] = [
"api.groq.com/openai/v1",
"https://integrate.api.nvidia.com/v1",
"api.deepseek.com/v1",
"api.together.ai/v1",
"api.together.xyz/v1",
"app.empower.dev/api/v1",
"https://api.friendli.ai/serverless/v1",

7
litellm/interactions/utils.py
View file

@ -47,6 +47,13 @@ def get_provider_interactions_api_config(
return GoogleAIStudioInteractionsConfig()
if provider in (LlmProviders.VERTEX_AI.value, LlmProviders.VERTEX_AI_BETA.value):
from litellm.llms.vertex_ai.interactions.transformation import (
VertexAIInteractionsConfig,
)
return VertexAIInteractionsConfig()
return None

10
litellm/litellm_core_utils/get_llm_provider_logic.py
View file

@ -272,6 +272,14 @@ def get_llm_provider(
elif endpoint == "api.deepseek.com/v1":
custom_llm_provider = "deepseek"
dynamic_api_key = get_secret_str("DEEPSEEK_API_KEY")
elif endpoint == "api.together.ai/v1" or endpoint == "api.together.xyz/v1":
custom_llm_provider = "together_ai"
dynamic_api_key = api_key or (
get_secret_str("TOGETHER_API_KEY")
or get_secret_str("TOGETHER_AI_API_KEY")
or get_secret_str("TOGETHERAI_API_KEY")
or get_secret_str("TOGETHER_AI_TOKEN")
)
elif endpoint == "ollama.com":
custom_llm_provider = "ollama"
dynamic_api_key = get_secret_str("OLLAMA_API_KEY")
@ -707,7 +715,7 @@ def _get_openai_compatible_provider_info(
dynamic_api_key,
) = litellm.ZAIChatConfig()._get_openai_compatible_provider_info(api_base, api_key)
elif custom_llm_provider == "together_ai":
api_base = api_base or get_secret_str("TOGETHER_AI_API_BASE") or "https://api.together.xyz/v1"
api_base = api_base or get_secret_str("TOGETHER_AI_API_BASE") or "https://api.together.ai/v1"
dynamic_api_key = api_key or (
get_secret_str("TOGETHER_API_KEY")
or get_secret_str("TOGETHER_AI_API_KEY")

43
litellm/litellm_core_utils/prompt_templates/factory.py
View file

@ -643,49 +643,6 @@ def claude_2_1_pt(
return prompt
### TOGETHER AI
def get_model_info(token, model):
try:
headers: Final = {"Authorization": f"Bearer {token}"}
client: Final = HTTPHandler(concurrent_limit=1)
response: Final = client.get("https://api.together.xyz/models/info", headers=headers)
if response.status_code == 200:
model_info: Final = response.json()
for m in model_info:
if m["name"].lower().strip() == model.strip():
return m["config"].get("prompt_format", None), m["config"].get("chat_template", None)
return None, None
else:
return None, None
except Exception: # safely fail a prompt template request
return None, None
## OLD TOGETHER AI FLOW
# def format_prompt_togetherai(messages, prompt_format, chat_template):
# if prompt_format is None:
# return default_pt(messages)
# human_prompt, assistant_prompt = prompt_format.split("{prompt}")
# if chat_template is not None:
# prompt = hf_chat_template(
# model=None, messages=messages, chat_template=chat_template
# )
# elif prompt_format is not None:
# prompt = custom_prompt(
# role_dict={},
# messages=messages,
# initial_prompt_value=human_prompt,
# final_prompt_value=assistant_prompt,
# )
# else:
# prompt = default_pt(messages)
# return prompt
### IBM Granite

9
litellm/llms/anthropic/chat/transformation.py
View file

@ -1215,8 +1215,11 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
if reasoning_effort is None or reasoning_effort == "none":
return None
if AnthropicConfig._is_adaptive_thinking_model(model, custom_llm_provider):
# without display, Anthropic defaults adaptive thinking to
# display="omitted" and returns a blank thinking block
return AnthropicThinkingParam(
type="adaptive",
display="summarized",
)
elif reasoning_effort == "low":
return AnthropicThinkingParam(
@ -2144,7 +2147,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
)
@staticmethod
def _thinking_tokens_from_usage(usage_object: Mapping[str, object]) -> int | None:
def thinking_tokens_from_usage(usage_object: Mapping[str, object]) -> int | None:
details: Final = usage_object.get("output_tokens_details")
if not isinstance(details, Mapping):
return None
@ -2176,7 +2179,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
reported_thinking_tokens: Final = (
iteration_thinking_tokens
if iteration_thinking_tokens is not None
else self._thinking_tokens_from_usage(usage_object)
else self.thinking_tokens_from_usage(usage_object)
)
if reported_thinking_tokens is not None:
capped_reported: Final = min(max(0, reported_thinking_tokens), completion_tokens)
@ -2199,7 +2202,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
def _sum_iteration_thinking_tokens(self, iterations: Sequence[object]) -> int | None:
per_iteration: Final = tuple(
self._thinking_tokens_from_usage(iteration) if isinstance(iteration, Mapping) else None
self.thinking_tokens_from_usage(iteration) if isinstance(iteration, Mapping) else None
for iteration in iterations
)
reported: Final = tuple(tokens for tokens in per_iteration if tokens is not None)

24
litellm/llms/bedrock/chat/converse_transformation.py
View file

@ -1617,6 +1617,8 @@ class AmazonConverseConfig(BaseConfig):
}
if additional_request_params:
data["additionalModelRequestFields"] = additional_request_params
if "thinking" in additional_request_params:
data["additionalModelResponseFieldPaths"] = ("/usage/output_tokens_details",)
if system_content_blocks:
data["system"] = system_content_blocks
@ -1801,6 +1803,17 @@ class AmazonConverseConfig(BaseConfig):
thinking_blocks_list.append(_redacted_block)
return thinking_blocks_list
@staticmethod
def thinking_tokens_from_additional_fields(additional_fields: object) -> int | None:
"""Converse omits thinking tokens from its usage block; they only arrive under
``additionalModelResponseFields`` when ``/usage/output_tokens_details`` is requested."""
if not isinstance(additional_fields, Mapping):
return None
usage: Final = additional_fields.get("usage")
if not isinstance(usage, Mapping):
return None
return AnthropicConfig.thinking_tokens_from_usage(usage)
@staticmethod
def is_converse_usage_shape(usage_object: Mapping[str, object]) -> bool:
"""Converse-family models report camelCase token counts, not Anthropic's snake_case."""
@ -1842,6 +1855,7 @@ class AmazonConverseConfig(BaseConfig):
usage: ConverseTokenUsageBlock,
reasoning_content: str | None = None,
thinking_ran: bool = False,
provider_reasoning_tokens: int | None = None,
) -> Usage:
input_tokens = usage["inputTokens"]
output_tokens: Final = usage["outputTokens"]
@ -1862,9 +1876,14 @@ class AmazonConverseConfig(BaseConfig):
cache_creation_tokens=cache_creation_input_tokens,
text_tokens=raw_input_tokens,
)
reasoning_tokens: Final = (
estimated_reasoning_tokens: Final = (
token_counter(text=reasoning_content, count_response_tokens=True) if reasoning_content else 0
)
reasoning_tokens: Final = (
min(max(0, provider_reasoning_tokens), output_tokens)
if provider_reasoning_tokens is not None
else estimated_reasoning_tokens
)
completion_tokens_details: Final = (
CompletionTokensDetailsWrapper(
reasoning_tokens=reasoning_tokens,
@ -2272,6 +2291,9 @@ class AmazonConverseConfig(BaseConfig):
completion_response["usage"],
reasoning_content=chat_completion_message.get("reasoning_content"),
thinking_ran=reasoningContentBlocks is not None,
provider_reasoning_tokens=self.thinking_tokens_from_additional_fields(
completion_response.get("additionalModelResponseFields")
),
)
## HANDLE TOOL CALLS

5
litellm/llms/bedrock/chat/invoke_handler.py
View file

@ -331,6 +331,7 @@ class AWSEventStreamDecoder:
self.json_mode = json_mode
self._current_tool_name: str | None = None
self._thinking_ran = False
self._provider_reasoning_tokens: int | None = None
def check_empty_tool_call_args(self) -> bool:
"""
@ -559,10 +560,14 @@ class AWSEventStreamDecoder:
tool_use = self._handle_converse_stop_event(content_block_index)
elif "stopReason" in chunk_data:
finish_reason = map_finish_reason(chunk_data.get("stopReason", "stop"))
self._provider_reasoning_tokens = AmazonConverseConfig.thinking_tokens_from_additional_fields(
chunk_data.get("additionalModelResponseFields")
)
elif "usage" in chunk_data:
usage = converse_config.transform_usage(
chunk_data.get("usage", {}),
thinking_ran=self._thinking_ran,
provider_reasoning_tokens=self._provider_reasoning_tokens,
)
if thinking_blocks:
self._thinking_ran = True

5
litellm/llms/bedrock/passthrough/transformation.py
View file

@ -1,4 +1,5 @@
import json
from collections.abc import Mapping
from typing import TYPE_CHECKING, Final, Optional, cast
from httpx import Response
@ -93,6 +94,9 @@ class BedrockPassthroughConfig(BaseAWSLLM, BedrockModelInfo, BedrockEventStreamD
endpoint_url,
)
def get_bedrock_bearer_token(self, litellm_params: Mapping[str, object]) -> str | None:
return None
def sign_request(
self,
headers: dict,
@ -109,6 +113,7 @@ class BedrockPassthroughConfig(BaseAWSLLM, BedrockModelInfo, BedrockEventStreamD
request_data=request_data or {},
api_base=api_base,
model=model,
api_key=self.get_bedrock_bearer_token(optional_params),
)
def logging_non_streaming_response(

42
litellm/llms/bedrock_mantle/common_utils.py
View file

@ -13,6 +13,7 @@ global state.
"""
import re
from collections.abc import Mapping
from typing import Final
from botocore.exceptions import (
@ -31,30 +32,39 @@ BEDROCK_MANTLE_DEFAULT_REGION: Final = "us-east-1"
MANTLE_HOST_RE: Final = re.compile(r"^https?://bedrock-mantle\.([^/.]+)\.api\.aws", re.IGNORECASE)
def resolve_mantle_bearer_token(api_key: str | None) -> str | None:
return api_key or get_secret_str("BEDROCK_MANTLE_API_KEY") or get_secret_str("AWS_BEARER_TOKEN_BEDROCK")
def resolve_mantle_region(params: Mapping[str, object]) -> str:
region: Final = params.get("aws_region_name")
if isinstance(region, str) and region:
BaseAWSLLM._validate_aws_region_name(region)
return region
api_base: Final = params.get("api_base")
base: Final = (api_base if isinstance(api_base, str) else None) or get_secret_str("BEDROCK_MANTLE_API_BASE")
if base:
match: Final = MANTLE_HOST_RE.match(base.rstrip("/"))
if match:
return match.group(1)
return (
get_secret_str("BEDROCK_MANTLE_REGION")
or get_secret_str("AWS_REGION_NAME")
or get_secret_str("AWS_REGION")
or BEDROCK_MANTLE_DEFAULT_REGION
)
class BedrockMantleAuthMixin:
_aws_signer: BaseAWSLLM
@staticmethod
def _resolve_bearer_token(api_key: str | None) -> str | None:
return api_key or get_secret_str("BEDROCK_MANTLE_API_KEY") or get_secret_str("AWS_BEARER_TOKEN_BEDROCK")
return resolve_mantle_bearer_token(api_key)
@staticmethod
def _resolve_region(params: dict) -> str:
region: Final = params.get("aws_region_name")
if region:
BaseAWSLLM._validate_aws_region_name(region)
return region
base: Final = params.get("api_base") or get_secret_str("BEDROCK_MANTLE_API_BASE")
if base:
match: Final = MANTLE_HOST_RE.match(base.rstrip("/"))
if match:
return match.group(1)
return (
get_secret_str("BEDROCK_MANTLE_REGION")
or get_secret_str("AWS_REGION_NAME")
or get_secret_str("AWS_REGION")
or BEDROCK_MANTLE_DEFAULT_REGION
)
return resolve_mantle_region(params)
def sign_request(
self,

71
litellm/llms/bedrock_mantle/passthrough/transformation.py Normal file
View file

@ -0,0 +1,71 @@
from collections.abc import Mapping
from typing import TYPE_CHECKING, Final, Literal, Optional
from httpx import Response
from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.llms.bedrock.passthrough.transformation import BedrockPassthroughConfig
from litellm.llms.bedrock_mantle.common_utils import (
MANTLE_HOST_RE,
resolve_mantle_bearer_token,
resolve_mantle_region,
)
from litellm.types.utils import LlmProviders
if TYPE_CHECKING:
from litellm.types.utils import CostResponseTypes
class BedrockMantlePassthroughConfig(BedrockPassthroughConfig):
"""Native Bedrock runtime passthrough (InvokeModel, Converse) for deployments declared as bedrock_mantle.
The Mantle host only serves the OpenAI-compatible surface, so a Mantle api_base lends its region and the
request itself goes to bedrock-runtime, signed with the deployment's Bearer token or SigV4 credentials.
"""
def _get_aws_region_name(
self,
optional_params: Mapping[str, object],
model: str | None = None,
model_id: str | None = None,
) -> str:
return resolve_mantle_region(optional_params)
def get_runtime_endpoint(
self,
api_base: str | None,
aws_bedrock_runtime_endpoint: str | None,
aws_region_name: str,
endpoint_type: Literal["runtime", "agent", "agentcore"] | None = "runtime",
) -> tuple[str, str]:
is_mantle_host: Final = api_base is not None and MANTLE_HOST_RE.match(api_base.rstrip("/")) is not None
return super().get_runtime_endpoint(
api_base=None if is_mantle_host else api_base,
aws_bedrock_runtime_endpoint=aws_bedrock_runtime_endpoint,
aws_region_name=aws_region_name,
endpoint_type=endpoint_type,
)
def get_bedrock_bearer_token(self, litellm_params: Mapping[str, object]) -> str | None:
api_key: Final = litellm_params.get("api_key")
return resolve_mantle_bearer_token(api_key if isinstance(api_key, str) else None)
def logging_non_streaming_response(
self,
model: str,
custom_llm_provider: str,
httpx_response: Response,
request_data: dict, # mutable-ok: mirrors the inherited BedrockPassthroughConfig signature
logging_obj: Logging,
endpoint: str,
) -> Optional["CostResponseTypes"]:
is_converse: Final = "invoke" not in endpoint and "converse" in endpoint
shape_provider: Final = LlmProviders.BEDROCK.value if is_converse else custom_llm_provider
return super().logging_non_streaming_response(
model=model,
custom_llm_provider=shape_provider,
httpx_response=httpx_response,
request_data=request_data,
logging_obj=logging_obj,
endpoint=endpoint,
)

12
litellm/llms/together_ai/rerank/handler.py
View file

@ -16,11 +16,16 @@ from litellm.llms.together_ai.rerank.transformation import TogetherAIRerankConfi
from litellm.types.rerank import RerankRequest, RerankResponse
def _rerank_url(api_base: str) -> str:
return f"{api_base.rstrip('/')}/rerank"
class TogetherAIRerank(BaseLLM):
def rerank(
self,
model: str,
api_key: str,
api_base: str,
query: str,
documents: list[str | dict[str, Any]],
top_n: int | None = None,
@ -46,10 +51,10 @@ class TogetherAIRerank(BaseLLM):
raise ValueError("TogetherAI does not support max_chunks_per_doc")
if _is_async:
return self.async_rerank(request_data_dict, api_key) # Call async method
return self.async_rerank(request_data_dict, api_key, api_base)
response: Final = client.post(
"https://api.together.xyz/v1/rerank",
_rerank_url(api_base),
headers={
"accept": "application/json",
"content-type": "application/json",
@ -69,11 +74,12 @@ class TogetherAIRerank(BaseLLM):
self,
request_data_dict: dict[str, Any],
api_key: str,
api_base: str,
) -> RerankResponse:
client: Final = get_async_httpx_client(llm_provider=litellm.LlmProviders.TOGETHER_AI) # Use async client
response: Final = await client.post(
"https://api.together.xyz/v1/rerank",
_rerank_url(api_base),
headers={
"accept": "application/json",
"content-type": "application/json",

0
litellm/llms/vertex_ai/interactions/__init__.py Normal file
View file

149
litellm/llms/vertex_ai/interactions/transformation.py Normal file
View file

@ -0,0 +1,149 @@
from collections.abc import Callable, Mapping
from dataclasses import dataclass
from typing import Final
from litellm.litellm_core_utils.url_utils import encode_url_path_segment
from litellm.llms.gemini.interactions.transformation import GoogleAIStudioInteractionsConfig
from litellm.llms.vertex_ai.common_utils import validate_vertex_location
from litellm.llms.vertex_ai.vertex_llm_base import VertexBase
from litellm.types.llms.vertex_ai import VERTEX_CREDENTIALS_TYPES
from litellm.types.router import GenericLiteLLMParams
from litellm.types.utils import LlmProviders
VERTEX_INTERACTIONS_API_VERSION: Final = "v1beta1"
VERTEX_INTERACTIONS_DEFAULT_LOCATION: Final = "global"
@dataclass(frozen=True, slots=True)
class VertexInteractionsTarget:
base_url: str
project_id: str
location: str
@property
def collection_url(self) -> str:
return (
f"{self.base_url}/{VERTEX_INTERACTIONS_API_VERSION}"
f"/projects/{self.project_id}/locations/{self.location}/interactions"
)
def interaction_url(self, interaction_id: str) -> str:
encoded_interaction_id: Final = encode_url_path_segment(interaction_id, field_name="interaction_id")
return f"{self.collection_url}/{encoded_interaction_id}"
class VertexAIInteractionsConfig(VertexBase, GoogleAIStudioInteractionsConfig):
def __init__(
self,
mint_access_token: Callable[[VERTEX_CREDENTIALS_TYPES | None, str | None], tuple[str, str]] | None = None,
) -> None:
super().__init__()
self._mint_access_token: Final[Callable[[VERTEX_CREDENTIALS_TYPES | None, str | None], tuple[str, str]]] = (
mint_access_token or self._mint_access_token_with_vertex_base
)
def _mint_access_token_with_vertex_base(
self,
credentials: VERTEX_CREDENTIALS_TYPES | None,
project_id: str | None,
) -> tuple[str, str]:
return self._ensure_access_token(
credentials=credentials, project_id=project_id, custom_llm_provider="vertex_ai"
)
@property
def custom_llm_provider(self) -> LlmProviders:
return LlmProviders.VERTEX_AI
@property
def api_version(self) -> str:
return VERTEX_INTERACTIONS_API_VERSION
def get_default_vertex_location(self) -> str:
return VERTEX_INTERACTIONS_DEFAULT_LOCATION
def _mint(self, litellm_params: GenericLiteLLMParams) -> tuple[str, str]:
raw_params: Final = litellm_params.model_dump()
return self._mint_access_token(
self.safe_get_vertex_ai_credentials(raw_params),
self.safe_get_vertex_ai_project(raw_params),
)
def _target(self, api_base: str | None, litellm_params: GenericLiteLLMParams) -> VertexInteractionsTarget:
_, project_id = self._mint(litellm_params)
if not project_id:
raise ValueError(
"Vertex AI project is required. Set vertex_project, litellm.vertex_project, or VERTEXAI_PROJECT"
)
location: Final = validate_vertex_location(
self.explicit_vertex_ai_location(litellm_params.model_dump()) or VERTEX_INTERACTIONS_DEFAULT_LOCATION
)
return VertexInteractionsTarget(
base_url=self.get_api_base(api_base or None, location),
project_id=project_id,
location=location,
)
def validate_environment(
self,
headers: Mapping[str, str],
model: str,
litellm_params: GenericLiteLLMParams | None,
) -> dict: # mutable-ok: BaseInteractionsAPIConfig declares plain-dict headers
access_token, _ = self._mint(litellm_params or GenericLiteLLMParams())
return { # mutable-ok: BaseInteractionsAPIConfig declares plain-dict headers
"Content-Type": "application/json",
"Authorization": f"Bearer {access_token}",
**headers,
}
def get_complete_url(
self,
api_base: str | None,
model: str | None,
agent: str | None = None,
litellm_params: Mapping[str, object] | None = None,
stream: bool | None = None,
) -> str:
params: Final = (
GenericLiteLLMParams.model_validate(litellm_params) if litellm_params else GenericLiteLLMParams()
)
collection_url: Final = self._target(api_base, params).collection_url
return f"{collection_url}?alt=sse" if stream else collection_url
def _interaction_by_id_request(
self,
interaction_id: str,
api_base: str,
litellm_params: GenericLiteLLMParams,
url_suffix: str = "",
) -> tuple[str, dict]: # mutable-ok: BaseInteractionsAPIConfig declares a plain-dict request body
target: Final = self._target(api_base or None, litellm_params)
return f"{target.interaction_url(interaction_id)}{url_suffix}", {} # mutable-ok: same base contract
def transform_get_interaction_request(
self,
interaction_id: str,
api_base: str,
litellm_params: GenericLiteLLMParams,
headers: Mapping[str, str],
) -> tuple[str, dict]: # mutable-ok: BaseInteractionsAPIConfig declares a plain-dict request body
return self._interaction_by_id_request(interaction_id, api_base, litellm_params)
def transform_delete_interaction_request(
self,
interaction_id: str,
api_base: str,
litellm_params: GenericLiteLLMParams,
headers: Mapping[str, str],
) -> tuple[str, dict]: # mutable-ok: BaseInteractionsAPIConfig declares a plain-dict request body
return self._interaction_by_id_request(interaction_id, api_base, litellm_params)
def transform_cancel_interaction_request(
self,
interaction_id: str,
api_base: str,
litellm_params: GenericLiteLLMParams,
headers: Mapping[str, str],
) -> tuple[str, dict]: # mutable-ok: BaseInteractionsAPIConfig declares a plain-dict request body
return self._interaction_by_id_request(interaction_id, api_base, litellm_params, url_suffix=":cancel")

2
litellm/main.py
View file

@ -602,7 +602,7 @@ async def acompletion(
_, custom_llm_provider, _, _ = get_llm_provider(
model=model,
custom_llm_provider=custom_llm_provider,
api_base=base_url,
api_base=kwargs.get("api_base") or base_url,
)
fallbacks = fallbacks or litellm.model_fallbacks

9
litellm/model_prices_and_context_window_backup.json
View file

@ -49016,12 +49016,13 @@
"output_cost_per_token": 3.3e-05,
"output_cost_per_token_above_272k_tokens": 4.95e-05,
"litellm_provider": "bedrock_mantle",
"max_input_tokens": 1000000,
"max_input_tokens": 1050000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "responses",
"use_openai_responses_path": true,
"supported_endpoints": [
"/v1/chat/completions",
"/v1/responses"
],
"supported_modalities": [
@ -49048,12 +49049,13 @@
"output_cost_per_token": 1.32e-05,
"output_cost_per_token_above_272k_tokens": 1.98e-05,
"litellm_provider": "bedrock_mantle",
"max_input_tokens": 1000000,
"max_input_tokens": 1050000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "responses",
"use_openai_responses_path": true,
"supported_endpoints": [
"/v1/chat/completions",
"/v1/responses"
],
"supported_modalities": [
@ -49080,12 +49082,13 @@
"output_cost_per_token": 1.32e-06,
"output_cost_per_token_above_272k_tokens": 1.98e-06,
"litellm_provider": "bedrock_mantle",
"max_input_tokens": 1000000,
"max_input_tokens": 1050000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "responses",
"use_openai_responses_path": true,
"supported_endpoints": [
"/v1/chat/completions",
"/v1/responses"
],
"supported_modalities": [

2
litellm/passthrough/main.py
View file

@ -199,7 +199,7 @@ def llm_passthrough_route(
api_key=api_key,
)
litellm_params_dict: Final = get_litellm_params(**kwargs)
litellm_params_dict: Final = get_litellm_params(api_key=api_key, api_base=api_base, **kwargs)
if client is None:
from litellm.llms.custom_httpx.http_handler import (

11
litellm/proxy/_types.py
View file

@ -815,6 +815,7 @@ class LiteLLMRoutes(enum.Enum):
"/team/member_add",
"/team/member_delete",
"/team/member_update",
"/team/{team_id}/member/{user_id}/reset_spend",
"/team/permissions_list",
"/team/permissions_update",
"/team/daily/activity",
@ -1287,6 +1288,16 @@ class RegenerateKeyRequest(GenerateKeyRequest):
class ResetSpendRequest(LiteLLMPydanticObjectBase):
reset_to: float
@field_validator("reset_to", mode="before")
@classmethod
def reject_bool_reset_to(cls, v):
# bool is a subclass of int, so pydantic silently coerces True/False into
# 1.0/0.0 for a `float` field: a caller who accidentally sends a boolean
# would otherwise get an unintended spend reset instead of a 422.
if isinstance(v, bool):
raise ValueError("reset_to must be a number, not a boolean") # noqa: TRY004 # pydantic needs ValueError
return v
class KeyRequest(LiteLLMPydanticObjectBase):
keys: list[str] | None = None

114
litellm/proxy/auth/auth_checks.py
View file

@ -87,6 +87,8 @@ from litellm.proxy.common_utils.user_api_key_cache import (
object_permission_cache_key,
tag_cache_key,
tag_registry_cache_key,
team_membership_auth_cache_key,
team_membership_reservation_cache_key,
)
from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler
from litellm.proxy.guardrails.tool_name_extraction import (
@ -1967,7 +1969,7 @@ async def get_team_membership(
if user_id is None or team_id is None:
return None
_key: Final = f"team_membership:{user_id}:{team_id}"
_key: Final = team_membership_reservation_cache_key(user_id=user_id, team_id=team_id)
# check if in cache
cached_membership_obj: Final = await user_api_key_cache.async_get_cache(
@ -2402,6 +2404,116 @@ async def _cache_team_object(
)
async def invalidate_team_member_spend_state(
user_id: str,
team_id: str,
user_api_key_cache: UserApiKeyCache,
new_spend: float | None = None,
) -> None:
"""
Clear every cached read path for one team member's budget so a spend
reset or a raised cap takes effect on the next request instead of
waiting on the membership cache's TTL.
Two independently-keyed cache entries hold the same LiteLLM_TeamMembership
row: user_api_key_auth.py's admission check writes ``{team_id}_{user_id}``,
while budget_reservation.py's pre-call reservation and auth_checks.py's own
get_team_membership() (used by _check_team_member_budget) both write
``team_membership:{user_id}:{team_id}``. Both formats must be invalidated
explicitly; writing one does not refresh the other. All keys are also
broadcast (LIT-3803): each worker's own in-memory copy (membership object,
spend counter, or the counter's own short-TTL DB-floor marker) survives
eviction elsewhere until its TTL, so the handling worker alone clearing its
copy leaves every other worker still enforcing the pre-reset budget.
``new_spend`` is only passed by reset_team_member_spend_fn, which knows the
exact post-reset value: it is SET everywhere (matching /key/{key}/reset_spend's
own precedent) rather than deleted, so a worker's next read reflects it
directly instead of re-deriving it through a DB reseed. team_member_update
only changes the budget cap, not the tracked spend, so it passes no
new_spend; the live spend counter is untouched in that case (deleting it
would force a reseed from the DB's own spend column, which lags the live
counter via periodic batch writes, briefly under-enforcing the raised cap
against a spend value lower than what was actually tracked) and only the
membership caches carrying the new cap are invalidated.
The floor marker (``spend_db_floor:``, proxy_server.py's
_authoritative_floor_spend) caches the pre-reset DB spend for
SPEND_DB_FLOOR_CACHE_TTL_SECONDS; left stale after a real reset, a request
landing on the pod that cached it can read that higher floor and raise the
counter right back above the just-reset spend. It is overwritten here with
the post-reset floor (not merely deleted) and _authoritative_floor_spend
re-checks the marker after its DB read, so a floor read already in flight
on this pod when the reset commits cannot clobber it with the pre-reset
value. Both keys are broadcast as SETs carrying new_spend, not deletes:
every subscriber (remote pods AND this pod's own, which receives its own
message) writes the post-reset value, so the self-delivered message cannot
erase the guard just written here.
Raises HTTPException(503) if Redis still holds the stale pre-reset counter
after both the SET and the fallback DELETE fail: budget checks read Redis
first, so returning success would leave the old value authoritative for
every worker despite the DB write having committed.
"""
from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import (
evict_and_broadcast,
publish_auth_cache_invalidation,
)
if new_spend is not None:
from litellm.proxy.proxy_server import SPEND_DB_FLOOR_CACHE_TTL_SECONDS, spend_counter_cache
spend_counter_key: Final = f"spend:team_member:{user_id}:{team_id}"
spend_db_floor_key: Final = f"spend_db_floor:{spend_counter_key}"
spend_counter_cache.in_memory_cache.set_cache(key=spend_counter_key, value=new_spend, ttl=60)
if spend_counter_cache.redis_cache is not None:
try:
await spend_counter_cache.redis_cache.async_set_cache(key=spend_counter_key, value=new_spend, ttl=60)
except Exception as e: # noqa: BLE001 # fall back to deleting the stale entry before giving up
verbose_proxy_logger.warning(
"Failed to set spend counter %s in Redis after reset: %s; deleting it instead so the next "
"read reseeds from the DB rather than keeping the stale pre-reset value authoritative",
spend_counter_key,
e,
)
try:
await spend_counter_cache.redis_cache.async_delete_cache(key=spend_counter_key)
except Exception: # noqa: BLE001 # stale value now authoritative in Redis; surface instead of reporting success
verbose_proxy_logger.warning(
"Failed to delete stale spend counter %s in Redis after a failed reset write",
spend_counter_key,
exc_info=True,
)
raise HTTPException(
status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
detail={ # mutable-ok: HTTPException.detail takes a dict
"error": "Spend was reset in the database, but Redis is unreachable and still "
"holds the pre-reset counter. Retry once Redis is reachable."
},
) from e
spend_counter_cache.in_memory_cache.set_cache(
key=spend_db_floor_key,
value=new_spend,
ttl=SPEND_DB_FLOOR_CACHE_TTL_SECONDS,
)
await publish_auth_cache_invalidation(cache_key=spend_counter_key, new_value=new_spend, ttl=60)
await publish_auth_cache_invalidation(
cache_key=spend_db_floor_key,
new_value=new_spend,
ttl=SPEND_DB_FLOOR_CACHE_TTL_SECONDS,
)
await evict_and_broadcast(
cache_keys=(
team_membership_auth_cache_key(team_id=team_id, user_id=user_id),
team_membership_reservation_cache_key(user_id=user_id, team_id=team_id),
),
user_api_key_cache=user_api_key_cache,
)
async def delete_cache_team_object(
team_id: str,
team_alias: str | None,

50
litellm/proxy/auth/user_api_key_auth.py
View file

@ -87,7 +87,10 @@ from litellm.proxy.common_utils.http_parsing_utils import (
populate_request_with_path_params,
)
from litellm.proxy.common_utils.realtime_utils import _realtime_request_body
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.proxy.common_utils.user_api_key_cache import (
UserApiKeyCache,
team_membership_auth_cache_key,
)
from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler
from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup
from litellm.proxy.utils import (
@ -1970,8 +1973,10 @@ async def _user_api_key_auth_builder(
# Check 3. Check if user is in their team budget
if not skip_budget_checks and valid_token.team_member_spend is not None:
if prisma_client is not None:
_cache_key: Final = f"{valid_token.team_id}_{valid_token.user_id}"
_user_id: Final = valid_token.user_id
_team_id: Final = valid_token.team_id
if prisma_client is not None and _user_id is not None and _team_id is not None:
_cache_key: Final = team_membership_auth_cache_key(team_id=_team_id, user_id=_user_id)
team_member_info = await user_api_key_cache.async_get_cache(
key=_cache_key,
@ -1979,25 +1984,21 @@ async def _user_api_key_auth_builder(
)
if team_member_info is None:
# read from DB
_user_id: Final = valid_token.user_id
_team_id: Final = valid_token.team_id
if _user_id is not None and _team_id is not None:
_db_member: Final = await TeamMembershipRepository(prisma_client).table.find_first(
where={
"user_id": _user_id,
"team_id": _team_id,
},
include={"litellm_budget_table": True},
_db_member: Final = await TeamMembershipRepository(prisma_client).table.find_first(
where={
"user_id": _user_id,
"team_id": _team_id,
},
include={"litellm_budget_table": True},
)
if _db_member is not None:
team_member_info = LiteLLM_TeamMembership(**_db_member.dict())
await user_api_key_cache.async_set_cache(
key=_cache_key,
value=team_member_info,
model_type=LiteLLM_TeamMembership,
ttl=5,
)
if _db_member is not None:
team_member_info = LiteLLM_TeamMembership(**_db_member.dict())
await user_api_key_cache.async_set_cache(
key=_cache_key,
value=team_member_info,
model_type=LiteLLM_TeamMembership,
ttl=5,
)
if team_member_info is not None and team_member_info.litellm_budget_table is not None:
team_member_budget: Final = team_member_info.litellm_budget_table.max_budget
@ -2013,11 +2014,16 @@ async def _user_api_key_auth_builder(
max_budget=team_member_budget,
)
if team_member_spend > team_member_budget:
_entity_id: Final = f"{valid_token.user_id}:{valid_token.team_id}"
raise litellm.BudgetExceededError(
current_cost=team_member_spend,
max_budget=team_member_budget,
message=(
f"Budget has been exceeded! TeamMember={_entity_id} "
f"Current cost: {team_member_spend}, Max budget: {team_member_budget}"
),
entity_type=Litellm_EntityType.TEAM_MEMBER.value,
entity_id=f"{valid_token.user_id}:{valid_token.team_id}",
entity_id=_entity_id,
)
# Check 3. If token is expired

57
litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py
View file

@ -12,6 +12,7 @@ from litellm.proxy.common_utils.config_sync_pubsub import (
)
if TYPE_CHECKING:
from litellm.caching.in_memory_cache import InMemoryCache
from litellm.caching.redis_cache import RedisCache
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
@ -30,15 +31,24 @@ def auth_cache_invalidation_channel(redis_cache: "RedisCache") -> str:
@dataclass(frozen=True, slots=True)
class _CacheInvalidationMessage:
cache_key: str
new_value: float | None = None
ttl: float | None = None
def _cache_invalidation_message_json(cache_key: str) -> str:
return json.dumps(asdict(_CacheInvalidationMessage(cache_key=cache_key)))
def _cache_invalidation_message_json(cache_key: str, new_value: float | None = None, ttl: float | None = None) -> str:
message: Final = asdict(_CacheInvalidationMessage(cache_key=cache_key, new_value=new_value, ttl=ttl))
return json.dumps({field: value for field, value in message.items() if value is not None})
def _cache_key_from_message_data(data: object) -> str | None:
def _finite_number_or_none(value: object) -> float | None:
if isinstance(value, bool) or not isinstance(value, (int, float)):
return None
return float(value)
def _message_from_data(data: object) -> _CacheInvalidationMessage | None:
if isinstance(data, bytes):
data = data.decode("utf-8", errors="replace")
data = data.decode("utf-8", errors="replace") # rebind-ok: normalizing the wire payload to str
if not isinstance(data, str):
return None
try:
@ -48,14 +58,28 @@ def _cache_key_from_message_data(data: object) -> str | None:
if not isinstance(parsed, dict):
return None
cache_key: Final = parsed.get("cache_key")
return cache_key if isinstance(cache_key, str) else None
if not isinstance(cache_key, str):
return None
return _CacheInvalidationMessage(
cache_key=cache_key,
new_value=_finite_number_or_none(parsed.get("new_value")),
ttl=_finite_number_or_none(parsed.get("ttl")),
)
async def publish_auth_cache_invalidation(cache_key: str) -> None:
async def publish_auth_cache_invalidation(
cache_key: str, new_value: float | None = None, ttl: float | None = None
) -> None:
"""
Best-effort broadcast so every worker drops its local in-memory copy of a
mutated management object; without this, only the handling worker and Redis
are evicted and other workers keep serving the stale object until its TTL.
Passing ``new_value`` broadcasts a SET instead of a delete: every subscriber
(including the publishing worker's own, which receives its own message)
writes the value into its additional in-memory caches rather than deleting
the key. A spend reset uses this so the handler's self-delivered message
cannot erase the freshly-written post-reset counter or floor marker.
"""
redis_cache: Final = coordination_redis_cache()
if redis_cache is None:
@ -68,7 +92,10 @@ async def publish_auth_cache_invalidation(cache_key: str) -> None:
cache_key,
)
return
await client.publish(auth_cache_invalidation_channel(redis_cache), _cache_invalidation_message_json(cache_key))
await client.publish(
auth_cache_invalidation_channel(redis_cache),
_cache_invalidation_message_json(cache_key, new_value=new_value, ttl=ttl),
)
except Exception as e: # noqa: BLE001 # best-effort publish; mutations must never fail on redis errors
verbose_proxy_logger.warning("auth cache invalidation publish for %s failed: %s", cache_key, e)
@ -95,15 +122,17 @@ async def evict_and_broadcast(cache_keys: Sequence[str], user_api_key_cache: "Us
class AuthCacheInvalidationSubscriber:
__slots__ = ("_redis_cache", "_task", "_user_api_key_cache")
__slots__ = ("_additional_in_memory_caches", "_redis_cache", "_task", "_user_api_key_cache")
def __init__(
self,
redis_cache: "RedisCache",
user_api_key_cache: "UserApiKeyCache",
additional_in_memory_caches: Sequence["InMemoryCache"] = (),
) -> None:
self._redis_cache = redis_cache
self._user_api_key_cache = user_api_key_cache
self._additional_in_memory_caches = tuple(additional_in_memory_caches)
self._task: asyncio.Task[None] | None = None
def start(self) -> None:
@ -160,12 +189,18 @@ class AuthCacheInvalidationSubscriber:
def _apply_message(self, message: object) -> None:
data: Final = message.get("data") if isinstance(message, dict) else None
cache_key: Final = _cache_key_from_message_data(data)
if cache_key is None:
parsed: Final = _message_from_data(data)
if parsed is None:
return
if parsed.new_value is not None:
for additional_cache in self._additional_in_memory_caches:
additional_cache.set_cache(parsed.cache_key, parsed.new_value, ttl=parsed.ttl)
return
in_memory_cache: Final = self._user_api_key_cache.in_memory_cache
if in_memory_cache is not None:
in_memory_cache.delete_cache(cache_key)
in_memory_cache.delete_cache(parsed.cache_key)
for additional_cache in self._additional_in_memory_caches:
additional_cache.delete_cache(parsed.cache_key)
@staticmethod
async def _close_pubsub(pubsub: _ConfigSyncPubSub) -> None:

15
litellm/proxy/common_utils/user_api_key_cache.py
View file

@ -200,6 +200,21 @@ def end_user_restricted_registry_cache_key() -> str:
return "end_user_restricted_registry"
def team_membership_auth_cache_key(team_id: str, user_id: str) -> str:
"""Cache key one team member's ``LiteLLM_TeamMembership`` row is stored under for the admission check."""
return f"{team_id}_{user_id}"
def team_membership_reservation_cache_key(user_id: str, team_id: str) -> str:
"""Cache key the pre-call budget reservation stores the same ``LiteLLM_TeamMembership`` row under.
Deliberately not unified with ``team_membership_auth_cache_key``: the two readers wrote independent
keys before this file existed, so a fix that invalidates one must invalidate both explicitly rather
than assume a single write is visible to both.
"""
return f"team_membership:{user_id}:{team_id}"
def get_management_object_ttl(cache: DualCache) -> float:
"""
In-memory TTL for management-object cache writes (keys, teams, users, budgets, ...).

131
litellm/proxy/management_endpoints/team_endpoints.py
View file

@ -16,7 +16,7 @@ import traceback
from collections.abc import Mapping, Sequence
from datetime import datetime, timezone
from types import MappingProxyType
from typing import Annotated, Final, NamedTuple, Protocol, TypedDict, TypeVar, cast
from typing import Annotated, Final, NamedTuple, NoReturn, Protocol, TypedDict, TypeVar, cast
import fastapi
from fastapi import APIRouter, Depends, Header, HTTPException, Request, status
@ -56,6 +56,7 @@ from litellm.proxy._types import (
PatchTeamRequest,
ProxyErrorTypes,
ProxyException,
ResetSpendRequest,
SpecialManagementEndpointEnums,
SpecialModelNames,
SpecialProxyStrings,
@ -84,6 +85,7 @@ from litellm.proxy.auth.auth_checks import (
get_team_membership,
get_team_object,
get_user_object,
invalidate_team_member_spend_state,
)
from litellm.proxy.auth.auth_utils import (
enforce_batch_enqueued_token_limit_is_admin_only,
@ -3392,7 +3394,7 @@ async def team_member_update(
Update team member budgets and team member role
"""
from litellm.proxy.proxy_server import premium_user, prisma_client
from litellm.proxy.proxy_server import premium_user, prisma_client, user_api_key_cache
if prisma_client is None:
raise HTTPException(status_code=500, detail={"error": "No db connected"})
@ -3491,6 +3493,12 @@ async def team_member_update(
budget_patch=budget_patch,
team_default_budget_id=team_default_budget_id,
)
if budget_patch:
await invalidate_team_member_spend_state(
user_id=received_user_id,
team_id=data.team_id,
user_api_key_cache=user_api_key_cache,
)
### update team member role
if data.role is not None:
@ -3527,6 +3535,125 @@ async def team_member_update(
)
def _check_not_resetting_own_spend(user_id: str, user_api_key_dict: UserAPIKeyAuth) -> None:
"""
_verify_team_access authorizes a team admin (or org admin) over their own
team, with no check that the target user_id differs from the caller. Left
unchecked, that admin could target their own LiteLLM_TeamMembership row and
repeatedly reset it to 0 right before it crosses their per-member cap,
consuming the shared team budget without the configured limit ever binding.
Only a proxy admin may reset an admin's own spend.
"""
if user_id == user_api_key_dict.user_id and user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN:
_raise_reset_spend_error(status.HTTP_403_FORBIDDEN, "Cannot reset your own spend. Ask a proxy admin.")
def _raise_reset_spend_error(status_code: int, message: str) -> NoReturn:
detail: Final = {"error": message} # mutable-ok: HTTPException.detail takes a dict
raise HTTPException(status_code=status_code, detail=detail)
def _validate_team_member_reset_spend_value(
reset_to: object,
membership: LiteLLM_TeamMembership,
) -> float:
if not isinstance(reset_to, (int, float)):
_raise_reset_spend_error(status.HTTP_400_BAD_REQUEST, "reset_to must be a float")
reset_to_float: Final = float(reset_to)
if not math.isfinite(reset_to_float) or reset_to_float < 0:
_raise_reset_spend_error(status.HTTP_400_BAD_REQUEST, "reset_to must be a finite number >= 0")
current_spend: Final = membership.spend or 0.0
if reset_to_float > current_spend:
_raise_reset_spend_error(
status.HTTP_400_BAD_REQUEST,
f"reset_to ({reset_to_float}) must be <= current spend ({current_spend})",
)
max_budget: Final = membership.litellm_budget_table.max_budget if membership.litellm_budget_table else None
if max_budget is not None and reset_to_float > max_budget:
_raise_reset_spend_error(
status.HTTP_400_BAD_REQUEST,
f"reset_to ({reset_to_float}) must be <= budget ({max_budget})",
)
return reset_to_float
@router.post(
"/team/{team_id}/member/{user_id}/reset_spend",
tags=["team management"], # mutable-ok: FastAPI's `tags` param is typed as list[str], not Sequence
dependencies=(Depends(user_api_key_auth),),
)
@management_endpoint_wrapper
async def reset_team_member_spend_fn(
team_id: str,
user_id: str,
data: ResetSpendRequest,
user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)],
):
"""
Reset a team member's tracked spend against their per-member budget.
A member's spend is tracked separately from both their own personal
budget and the team's own budget (LiteLLM_TeamMembership.spend), so
neither /user/update nor /team/update can clear it: this is the only
endpoint that does. The cross-pod spend counter and cached membership
reads are invalidated so the reset takes effect on the member's next
request rather than waiting on the membership cache's TTL.
"""
from litellm.proxy.proxy_server import prisma_client, proxy_logging_obj, user_api_key_cache
if prisma_client is None:
_raise_reset_spend_error(status.HTTP_500_INTERNAL_SERVER_ERROR, "DB not connected. prisma_client is None")
team_obj: Final = await get_team_object(
team_id=team_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=None,
proxy_logging_obj=proxy_logging_obj,
check_db_only=True,
)
await _verify_team_access(team_obj=team_obj, user_api_key_dict=user_api_key_dict)
_check_not_resetting_own_spend(user_id=user_id, user_api_key_dict=user_api_key_dict)
membership_where: Final = { # mutable-ok: prisma client requires a plain dict where= argument
"user_id_team_id": {"user_id": user_id, "team_id": team_id} # mutable-ok: same prisma where= argument
}
_membership_row: Final = await _team_membership_db(prisma_client).find_unique(
where=membership_where,
include={"litellm_budget_table": True}, # mutable-ok: prisma client requires a plain dict include= argument
)
if _membership_row is None:
_raise_reset_spend_error(status.HTTP_404_NOT_FOUND, f"User {user_id} is not a member of team {team_id}.")
membership: Final = LiteLLM_TeamMembership.model_validate(_membership_row.model_dump())
current_spend: Final = membership.spend or 0.0
reset_to: Final = _validate_team_member_reset_spend_value(data.reset_to, membership)
await _team_membership_db(prisma_client).update(
where=membership_where,
data={"spend": reset_to}, # mutable-ok: prisma client requires a plain dict data= argument
)
await invalidate_team_member_spend_state(
user_id=user_id,
team_id=team_id,
user_api_key_cache=user_api_key_cache,
new_spend=reset_to,
)
return { # mutable-ok: matches this router's established untyped-response-dict convention
"team_id": team_id,
"user_id": user_id,
"spend": reset_to,
"previous_spend": current_spend,
"max_budget": membership.litellm_budget_table.max_budget if membership.litellm_budget_table else None,
}
def _create_results_from_response(
members: list[Member],
response: TeamAddMemberResponse,

7
litellm/proxy/proxy_server.py
View file

@ -2555,6 +2555,12 @@ async def _authoritative_floor_spend(
if db_spend is None:
return None
# a spend reset that committed during the DB read above wrote the post-reset
# floor to the marker; keep it over this read's now-stale pre-commit value
rechecked: Final = spend_counter_cache.in_memory_cache.get_cache(key=marker_key)
if rechecked is not None:
return float(rechecked)
spend_counter_cache.in_memory_cache.set_cache(
key=marker_key,
value=db_spend,
@ -6798,6 +6804,7 @@ class ProxyConfig:
subscriber: Final = AuthCacheInvalidationSubscriber(
redis_cache=redis_cache,
user_api_key_cache=user_api_key_cache,
additional_in_memory_caches=(spend_counter_cache.in_memory_cache,),
)
self.auth_cache_invalidation_subscriber = subscriber
subscriber.start()

10
litellm/proxy/spend_tracking/budget_reservation.py
View file

@ -25,7 +25,11 @@ from litellm.proxy._types import (
from litellm.proxy.auth.auth_utils import get_model_from_request
from litellm.proxy.auth.budget_throttle import should_throttle_budget_exceeded
from litellm.proxy.auth.route_checks import RouteChecks
from litellm.proxy.common_utils.user_api_key_cache import end_user_cache_key, tag_cache_key
from litellm.proxy.common_utils.user_api_key_cache import (
end_user_cache_key,
tag_cache_key,
team_membership_reservation_cache_key,
)
from litellm.proxy.utils import PrismaClient, ProxyLogging
from litellm.router import Router
@ -546,7 +550,9 @@ async def _get_team_member_budget_counter(
if team_object is None or team_object.team_id is None or user_object is None or valid_token.user_id is None:
return None
membership_cache_key: Final = f"team_membership:{valid_token.user_id}:{team_object.team_id}"
membership_cache_key: Final = team_membership_reservation_cache_key(
user_id=valid_token.user_id, team_id=team_object.team_id
)
cached_team_membership: Final = await user_api_key_cache.async_get_cache(key=membership_cache_key)
team_membership: LiteLLM_TeamMembership | None = None
if isinstance(cached_team_membership, LiteLLM_TeamMembership):

3
litellm/rerank_api/main.py
View file

@ -277,6 +277,8 @@ def rerank(
if api_key is None:
raise ValueError("TogetherAI API key is required, please set 'TOGETHERAI_API_KEY' in your environment")
api_base = dynamic_api_base or optional_params.api_base or litellm.api_base or "https://api.together.ai/v1"
response = together_rerank.rerank(
model=model,
query=query,
@ -286,6 +288,7 @@ def rerank(
return_documents=return_documents,
max_chunks_per_doc=max_chunks_per_doc,
api_key=api_key,
api_base=api_base,
_is_async=_is_async,
)
elif _custom_llm_provider == litellm.LlmProviders.JINA_AI:

1
litellm/router.py
View file

@ -8346,6 +8346,7 @@ class Router:
) = litellm.get_llm_provider(
model=deployment.litellm_params.model,
custom_llm_provider=deployment.litellm_params.get("custom_llm_provider", None),
api_base=deployment.litellm_params.api_base,
)
# done reading model["litellm_params"]
# Check if provider is supported: either in enum or JSON-configured

1
litellm/types/llms/anthropic.py
View file

@ -685,6 +685,7 @@ ANTHROPIC_API_ONLY_HEADERS: Final = { # fails if calling anthropic on vertex ai
class AnthropicThinkingParam(TypedDict, total=False):
type: ReadOnly[Literal["enabled", "adaptive", "disabled"]]
budget_tokens: int
display: ReadOnly[Literal["summarized", "omitted"]]
class ANTHROPIC_HOSTED_TOOLS(str, Enum):

3
litellm/types/llms/bedrock.py
View file

@ -1,4 +1,5 @@
import json
from collections.abc import Sequence
from enum import Enum
from typing import TYPE_CHECKING, Any, Final, Literal
@ -396,7 +397,7 @@ class OutputConfigBlock(TypedDict, total=False):
class CommonRequestObject(TypedDict, total=False): # common request object across sync + async flows
additionalModelRequestFields: dict
additionalModelResponseFieldPaths: list[str]
additionalModelResponseFieldPaths: Sequence[str]
inferenceConfig: InferenceConfig
system: list[SystemContentBlock]
toolConfig: ToolConfigBlock

6
litellm/utils.py
View file

@ -8610,6 +8610,12 @@ class ProviderConfigManager:
)
return BedrockPassthroughConfig()
elif LlmProviders.BEDROCK_MANTLE == provider:
from litellm.llms.bedrock_mantle.passthrough.transformation import (
BedrockMantlePassthroughConfig,
)
return BedrockMantlePassthroughConfig()
elif LlmProviders.VLLM == provider or LlmProviders.HOSTED_VLLM == provider:
from litellm.llms.vllm.passthrough.transformation import (
VLLMPassthroughConfig,

9
model_prices_and_context_window.json
View file

@ -49016,12 +49016,13 @@
"output_cost_per_token": 3.3e-05,
"output_cost_per_token_above_272k_tokens": 4.95e-05,
"litellm_provider": "bedrock_mantle",
"max_input_tokens": 1000000,
"max_input_tokens": 1050000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "responses",
"use_openai_responses_path": true,
"supported_endpoints": [
"/v1/chat/completions",
"/v1/responses"
],
"supported_modalities": [
@ -49048,12 +49049,13 @@
"output_cost_per_token": 1.32e-05,
"output_cost_per_token_above_272k_tokens": 1.98e-05,
"litellm_provider": "bedrock_mantle",
"max_input_tokens": 1000000,
"max_input_tokens": 1050000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "responses",
"use_openai_responses_path": true,
"supported_endpoints": [
"/v1/chat/completions",
"/v1/responses"
],
"supported_modalities": [
@ -49080,12 +49082,13 @@
"output_cost_per_token": 1.32e-06,
"output_cost_per_token_above_272k_tokens": 1.98e-06,
"litellm_provider": "bedrock_mantle",
"max_input_tokens": 1000000,
"max_input_tokens": 1050000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "responses",
"use_openai_responses_path": true,
"supported_endpoints": [
"/v1/chat/completions",
"/v1/responses"
],
"supported_modalities": [

6
ruff-strict-budget.json
View file

@ -1,6 +1,6 @@
{
"ANN001": {
"limit": 3018
"limit": 3016
},
"ANN002": {
"limit": 71
@ -9,7 +9,7 @@
"limit": 827
},
"ANN201": {
"limit": 2016
"limit": 2015
},
"ANN202": {
"limit": 852
@ -57,7 +57,7 @@
"limit": 3
},
"BLE001": {
"limit": 2919
"limit": 2918
},
"C401": {
"limit": 8

3
tests/mcp_tests/conftest.py
View file

@ -7,6 +7,7 @@ import pytest
import litellm
import asyncio
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
@pytest.fixture(scope="session")
@ -38,6 +39,8 @@ def setup_and_teardown():
yield
# Teardown code (executes after the yield point)
# LoggingWorker carries still-queued coroutines onto the next test's loop, where they'd log into that test's callbacks
asyncio.run(GLOBAL_LOGGING_WORKER.clear_queue())
loop.close() # Close the loop created earlier
asyncio.set_event_loop(None) # Remove the reference to the loop

45
tests/mcp_tests/test_mcp_logging.py
View file

@ -1,6 +1,9 @@
import os
import pytest
import asyncio
import subprocess
import sys
from pathlib import Path
from typing import Optional
from unittest.mock import AsyncMock, patch
@ -458,3 +461,45 @@ async def test_mcp_tool_call_hook():
logged_standard_logging_payload is not None
), "Standard logging payload should not be None"
assert logged_standard_logging_payload["response_cost"] == 1.42
_QUEUED_LOGGING_OUTLIVES_TEST = '''
import time
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
ran_at = []
async def _record_run():
ran_at.append(time.monotonic())
async def test_1_leaves_logging_queued_behind_a_stopped_worker():
GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(_record_run())
await GLOBAL_LOGGING_WORKER.stop()
assert ran_at == []
async def test_2_starts_after_the_previous_tests_logging_ran():
started_at = time.monotonic()
GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(_record_run())
await GLOBAL_LOGGING_WORKER.flush()
assert [t < started_at for t in ran_at] == [True, False]
'''
def test_logging_queued_by_one_test_is_drained_before_the_next(tmp_path: Path):
"""Regression: a logging coroutine queued by one test must not run inside a later test (it would log into that
test's callbacks, which is how test_mcp_tool_call_hook captured a gpt-4o-mini payload under xdist)."""
(tmp_path / "conftest.py").write_text((Path(__file__).parent / "conftest.py").read_text())
(tmp_path / "pyproject.toml").write_text('[tool.pytest.ini_options]\nasyncio_mode = "auto"\n')
(tmp_path / "test_queued_logging.py").write_text(_QUEUED_LOGGING_OUTLIVES_TEST)
result = subprocess.run(
[sys.executable, "-m", "pytest", "-q", "-p", "no:cacheprovider", "test_queued_logging.py"],
cwd=tmp_path,
capture_output=True,
text=True,
timeout=120,
)
assert result.returncode == 0, result.stdout + result.stderr

152
tests/proxy_behavior/management/test_team_member_reset_spend.py Normal file
View file

@ -0,0 +1,152 @@
import uuid
import pytest
from .actors import Actor
from .conftest import create_scratch_team
pytestmark = pytest.mark.asyncio(loop_scope="session")
_SEED_SPEND = 5.0
_RESET_TO = 2.0
# POST /team/{team_id}/member/{user_id}/reset_spend. The handler gate is
# _verify_team_access (proxy admin / team admin of this team / org admin of
# the team's org) — the same gate /team/member_update uses, so this mirrors
# that file's matrix exactly.
_MATRIX = [
("alpha/proxy_admin", Actor.PROXY_ADMIN, "alpha", 200),
("alpha/org_admin", Actor.ORG_ADMIN, "alpha", 200),
("alpha/team_admin", Actor.TEAM_ADMIN, "alpha", 200),
("alpha/internal_user", Actor.INTERNAL_USER, "alpha", 403),
("alpha/owner", Actor.OWNER, "alpha", 403),
("alpha/unrelated_same_org", Actor.UNRELATED_SAME_ORG, "alpha", 403),
("alpha/cross_org_user", Actor.CROSS_ORG_USER, "alpha", 403),
("alpha/service_account", Actor.SERVICE_ACCOUNT, "alpha", 403),
("alpha/org_b_admin", Actor.ORG_B_ADMIN, "alpha", 403),
("beta/proxy_admin", Actor.PROXY_ADMIN, "beta", 200),
("beta/org_admin", Actor.ORG_ADMIN, "beta", 403),
("beta/team_admin", Actor.TEAM_ADMIN, "beta", 403),
("beta/org_b_admin", Actor.ORG_B_ADMIN, "beta", 200),
]
async def _seed_target(prisma, world, shape: str, team_id: str, member_id: str) -> None:
if shape == "alpha":
await create_scratch_team(
prisma,
team_id,
organization_id=world.org_a_id,
admin_user_ids=[world.keys[Actor.TEAM_ADMIN].user_id],
)
elif shape == "beta":
await create_scratch_team(prisma, team_id, organization_id=world.org_b_id)
else: # pragma: no cover - guard
pytest.fail(f"unknown shape={shape}")
await prisma.db.litellm_teammembership.create(
data={"user_id": member_id, "team_id": team_id, "spend": _SEED_SPEND}
)
@pytest.mark.parametrize(
"actor,shape,expected_status",
[(a, sh, s) for (_id, a, sh, s) in _MATRIX],
ids=[s[0] for s in _MATRIX],
)
async def test_team_member_reset_spend_authz_matrix(
actor: Actor,
shape: str,
expected_status: int,
proxy_client,
prisma,
scratch,
world,
):
member_id = scratch.tag("member")
await _seed_target(prisma, world, shape, scratch.prefix, member_id)
caller = world.keys[actor]
resp = await proxy_client.post(
f"/team/{scratch.prefix}/member/{member_id}/reset_spend",
headers={"Authorization": f"Bearer {caller.cleartext}"},
json={"reset_to": _RESET_TO},
)
assert (
resp.status_code == expected_status
), f"{actor.value} {shape}: {resp.status_code} {resp.text}"
row = await prisma.db.litellm_teammembership.find_unique(
where={"user_id_team_id": {"user_id": member_id, "team_id": scratch.prefix}}
)
assert row is not None
if expected_status == 200:
assert row.spend == _RESET_TO
else:
assert row.spend == _SEED_SPEND, "denied but spend reset"
async def test_team_member_reset_spend_missing_team_is_404(proxy_client, world):
resp = await proxy_client.post(
f"/team/behavior-pin-no-such-team/member/{uuid.uuid4().hex}/reset_spend",
headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"},
json={"reset_to": 0.0},
)
assert resp.status_code == 404, resp.text
async def test_team_member_reset_spend_missing_membership_is_404(
proxy_client, prisma, scratch, world
):
"""A well-formed team but a user_id with no LiteLLM_TeamMembership row is 404."""
await create_scratch_team(prisma, scratch.prefix, organization_id=world.org_a_id)
resp = await proxy_client.post(
f"/team/{scratch.prefix}/member/{uuid.uuid4().hex}/reset_spend",
headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"},
json={"reset_to": 0.0},
)
assert resp.status_code == 404, resp.text
async def test_team_member_reset_spend_above_current_spend_is_400(
proxy_client, prisma, scratch, world
):
member_id = scratch.tag("member")
await create_scratch_team(prisma, scratch.prefix, organization_id=world.org_a_id)
await prisma.db.litellm_teammembership.create(
data={"user_id": member_id, "team_id": scratch.prefix, "spend": 1.0}
)
resp = await proxy_client.post(
f"/team/{scratch.prefix}/member/{member_id}/reset_spend",
headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"},
json={"reset_to": 5.0},
)
assert resp.status_code == 400, resp.text
async def test_team_member_reset_spend_team_admin_cannot_reset_own_spend(
proxy_client, prisma, scratch, world
):
"""A team admin targeting their own LiteLLM_TeamMembership row is 403: unchecked, an
admin could repeatedly zero their own spend right before it crosses their per-member
cap, consuming the shared team budget without the configured limit ever binding."""
team_admin = world.keys[Actor.TEAM_ADMIN]
await create_scratch_team(
prisma,
scratch.prefix,
organization_id=world.org_a_id,
admin_user_ids=[team_admin.user_id],
)
await prisma.db.litellm_teammembership.create(
data={"user_id": team_admin.user_id, "team_id": scratch.prefix, "spend": _SEED_SPEND}
)
resp = await proxy_client.post(
f"/team/{scratch.prefix}/member/{team_admin.user_id}/reset_spend",
headers={"Authorization": f"Bearer {team_admin.cleartext}"},
json={"reset_to": 0.0},
)
assert resp.status_code == 403, resp.text
row = await prisma.db.litellm_teammembership.find_unique(
where={"user_id_team_id": {"user_id": team_admin.user_id, "team_id": scratch.prefix}}
)
assert row is not None and row.spend == _SEED_SPEND, "denied but spend reset"

28
tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py
View file

@ -2,7 +2,7 @@ import datetime
import json
import os
import unittest
from typing import TYPE_CHECKING, List, Literal, Optional, Tuple
from typing import TYPE_CHECKING, Final, List, Literal, Optional, Tuple
from unittest.mock import ANY, MagicMock, Mock, patch
import httpx
@ -2444,6 +2444,32 @@ def test_map_optional_params_preserves_reasoning_summary():
assert responses_api_request["reasoning"]["summary"] == "detailed"
@pytest.mark.parametrize("reasoning_effort", ["max", "high"])
def test_transform_request_bedrock_mantle_tools_keeps_reasoning_effort(monkeypatch, reasoning_effort):
"""Regression for reasoning_effort=max being dropped on the chat -> Responses bridge (issue #38084)."""
from litellm.completion_extras.litellm_responses_transformation.transformation import (
LiteLLMResponsesTransformationHandler,
)
monkeypatch.setattr(litellm, "reasoning_auto_summary", False)
monkeypatch.delenv("LITELLM_REASONING_AUTO_SUMMARY", raising=False)
handler: Final = LiteLLMResponsesTransformationHandler()
result: Final = handler.transform_request(
model="openai.gpt-5.6-sol",
messages=[{"role": "user", "content": "Say pong"}],
optional_params={
"reasoning_effort": reasoning_effort,
"tools": [{"type": "function", "function": {"name": "get_weather", "parameters": {"type": "object"}}}],
},
litellm_params={"custom_llm_provider": "bedrock_mantle"},
headers={},
litellm_logging_obj=Mock(),
)
assert result["reasoning"] == {"effort": reasoning_effort}
def test_map_optional_params_tool_choice_chat_nested_to_responses_api():
"""Chat tool_choice must become Responses ToolChoiceFunction (top-level name)."""
from litellm.completion_extras.litellm_responses_transformation.transformation import (

4
tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py
View file

@ -478,10 +478,10 @@ def test_generic_cost_per_token_minimax_m3_above_512k_tokens(_local_model_cost_m
],
)
def test_generic_cost_per_token_bedrock_mantle_gpt56_long_context(_local_model_cost_map, model):
"""Bedrock GPT-5.6 supports a 1M context window, billed at the long-context rates above 272K."""
"""Bedrock GPT-5.6 enforces a 1,050,000-token context window, billed at the long-context rates above 272K."""
model_cost_map = litellm.model_cost[model]
assert model_cost_map["max_input_tokens"] == 1000000
assert model_cost_map["max_input_tokens"] == 1050000
cached_tokens = 100000
completion_tokens = 1000

51
tests/test_litellm/litellm_core_utils/test_get_llm_provider_endpoint_match.py
View file

@ -133,3 +133,54 @@ class TestGetLlmProviderRejectsAttackerSmuggledApiBase:
assert provider == "groq"
assert dynamic_api_key == "server-real-groq-key"
class TestTogetherApiBaseResolvesProvider:
"""
Regression for the Together host migration: both the current
``api.together.ai`` host and the legacy ``api.together.xyz`` host must
resolve to ``together_ai`` when passed as ``api_base``. Before the fix
the endpoint list carried the legacy host but the provider-mapping
chain had no branch for it, so the match fell through with a None
provider and the deployment failed with "LLM Provider NOT provided".
"""
@pytest.mark.parametrize(
"api_base",
[
"https://api.together.ai/v1",
"https://api.together.xyz/v1",
],
)
def test_together_api_base_resolves_to_together_ai(self, api_base, monkeypatch):
monkeypatch.setenv("TOGETHER_API_KEY", "together-key-from-env")
model, provider, dynamic_api_key, returned_api_base = get_llm_provider(
model="some-model",
api_base=api_base,
)
assert provider == "together_ai"
assert dynamic_api_key == "together-key-from-env"
assert returned_api_base == api_base
assert model == "some-model"
def test_explicit_api_key_beats_together_env_key(self, monkeypatch):
monkeypatch.setenv("TOGETHER_API_KEY", "together-key-from-env")
_, provider, dynamic_api_key, _ = get_llm_provider(
model="some-model",
api_base="https://api.together.ai/v1",
api_key="explicit-caller-key",
)
assert provider == "together_ai"
assert dynamic_api_key == "explicit-caller-key"
def test_together_default_api_base_is_together_ai(self, monkeypatch):
monkeypatch.delenv("TOGETHER_AI_API_BASE", raising=False)
_, provider, _, api_base = get_llm_provider(model="together_ai/some-model")
assert provider == "together_ai"
assert api_base == "https://api.together.ai/v1"

2
tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_reasoning_effort_translation.py
View file

@ -42,7 +42,7 @@ def test_reasoning_effort_maps_to_output_config_for_adaptive_model(
)
assert "reasoning_effort" not in result
assert result.get("thinking") == {"type": "adaptive"}
assert result.get("thinking") == {"type": "adaptive", "display": "summarized"}
assert result.get("output_config") == {"effort": expected_effort}

12
tests/test_litellm/llms/anthropic/test_anthropic_reasoning_effort.py
View file

@ -5,6 +5,8 @@ Verifies that reasoning_effort=None returns None for all models,
including Claude Opus 4.6.
"""
import pytest
from litellm.llms.anthropic.chat.transformation import AnthropicConfig
@ -35,6 +37,16 @@ class TestMapReasoningEffort:
)
assert result["type"] == "adaptive"
@pytest.mark.parametrize("effort", ["low", "medium", "high"])
def test_adaptive_mapping_requests_summarized_display(self, effort):
"""Regression LIT-5714: adaptive thinking without ``display`` makes Anthropic
return a blank thinking block, so reasoning_effort callers always got
``reasoning_content: ""``."""
result = AnthropicConfig._map_reasoning_effort(
reasoning_effort=effort, model="claude-opus-4-6", custom_llm_provider="anthropic"
)
assert result["display"] == "summarized"
def test_other_model_low_returns_enabled_with_budget(self):
result = AnthropicConfig._map_reasoning_effort(
reasoning_effort="low", model="claude-4-sonnet-20250514", custom_llm_provider="anthropic"

2
tests/test_litellm/llms/azure_ai/claude/test_azure_anthropic_messages_transformation.py
View file

@ -341,7 +341,7 @@ def test_messages_thinking_shape_follows_exact_azure_entry_flag(local_model_cost
)
result = transform()
assert result.get("thinking") == {"type": "adaptive"}
assert result.get("thinking") == {"type": "adaptive", "display": "summarized"}
assert result.get("output_config") == {"effort": "medium"}
monkeypatch.setitem(

90
tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py
View file

@ -366,6 +366,96 @@ def test_output_config_effort_forwarded_into_additional_request_fields(model):
assert additional.get("output_config") == {"effort": "high"}
def test_reasoning_effort_requests_summarized_display_converse():
"""Regression LIT-5714: adaptive thinking synthesized from reasoning_effort must
request the summarized display, otherwise the provider returns a blank thinking
block and reasoning_content is always empty."""
config = AmazonConverseConfig()
optional_params = config.map_openai_params(
non_default_params={"reasoning_effort": "high"},
optional_params={},
model="bedrock/converse/us.anthropic.claude-opus-4-7",
drop_params=False,
)
assert optional_params["thinking"]["type"] == "adaptive"
assert optional_params["thinking"]["display"] == "summarized"
def test_thinking_request_adds_output_tokens_details_response_path():
"""Regression LIT-5714: the Converse usage block has no thinking-token field, so
thinking requests must ask for ``/usage/output_tokens_details`` via
``additionalModelResponseFieldPaths``."""
config = AmazonConverseConfig()
result = config._transform_request(
model="bedrock/converse/us.anthropic.claude-opus-4-7",
messages=[{"role": "user", "content": "hi"}],
optional_params={
"maxTokens": 256,
"thinking": {"type": "adaptive", "display": "summarized"},
"output_config": {"effort": "high"},
},
litellm_params={},
headers={},
)
assert result["additionalModelResponseFieldPaths"] == ("/usage/output_tokens_details",)
def test_request_without_thinking_omits_response_field_paths():
config = AmazonConverseConfig()
result = config._transform_request(
model="bedrock/converse/us.anthropic.claude-opus-4-7",
messages=[{"role": "user", "content": "hi"}],
optional_params={"maxTokens": 256},
litellm_params={},
headers={},
)
assert "additionalModelResponseFieldPaths" not in result
def test_transform_usage_prefers_provider_reasoning_tokens():
"""Regression LIT-5714: provider-reported thinking tokens must win over the
token_counter estimate derived from visible reasoning text."""
config = AmazonConverseConfig()
usage = config.transform_usage(
{"inputTokens": 40, "outputTokens": 3002, "totalTokens": 3042},
reasoning_content="a short reasoning summary",
thinking_ran=True,
provider_reasoning_tokens=1033,
)
assert usage.completion_tokens_details.reasoning_tokens == 1033
assert usage.completion_tokens_details.text_tokens == 3002 - 1033
def test_transform_usage_falls_back_to_estimate_without_provider_tokens():
config = AmazonConverseConfig()
usage = config.transform_usage(
{"inputTokens": 40, "outputTokens": 300, "totalTokens": 340},
reasoning_content="a short reasoning summary",
thinking_ran=True,
)
assert usage.completion_tokens_details.reasoning_tokens > 0
assert usage.completion_tokens_details.reasoning_tokens < 300
def test_thinking_tokens_parsed_from_additional_model_response_fields():
parsed = AmazonConverseConfig.thinking_tokens_from_additional_fields(
{"usage": {"output_tokens_details": {"thinking_tokens": 92}}}
)
assert parsed == 92
assert AmazonConverseConfig.thinking_tokens_from_additional_fields(None) is None
assert AmazonConverseConfig.thinking_tokens_from_additional_fields({"usage": {}}) is None
@pytest.mark.parametrize(
"model,effort,expected_effort",
[

23
tests/test_litellm/llms/bedrock/chat/test_invoke_handler.py
View file

@ -208,6 +208,29 @@ def test_bedrock_converse_streaming_consistent_id():
), "All chunk IDs must match the one captured from the messageStart event"
def test_converse_streaming_usage_uses_provider_thinking_tokens():
"""Regression LIT-5714: the messageStop event carries provider thinking tokens
under ``additionalModelResponseFields``; the usage chunk must report them instead
of a token_counter estimate."""
chunks = [
{
"contentBlockIndex": 0,
"delta": {"reasoningContent": {"text": "thinking about it"}},
},
{
"stopReason": "end_turn",
"additionalModelResponseFields": {"usage": {"output_tokens_details": {"thinking_tokens": 1033}}},
},
{"usage": {"inputTokens": 40, "outputTokens": 3002, "totalTokens": 3042}},
]
decoder = AWSEventStreamDecoder(model="bedrock/anthropic.claude-opus-4-7")
parsed = [decoder.converse_chunk_parser(chunk) for chunk in chunks]
usage = parsed[-1].usage
assert usage.completion_tokens_details.reasoning_tokens == 1033
@pytest.mark.asyncio
async def test_make_call_does_not_rechunk_stream_by_default():
"""Re-chunking the event stream into fixed 1024-byte blocks holds small

4
tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py
View file

@ -1387,7 +1387,7 @@ def test_bedrock_messages_maps_reasoning_effort_for_adaptive_model(
)
assert "reasoning_effort" not in result
assert result.get("thinking") == {"type": "adaptive"}
assert result.get("thinking") == {"type": "adaptive", "display": "summarized"}
assert result.get("output_config") == {"effort": expected_effort}
@ -2935,7 +2935,7 @@ def test_bedrock_messages_thinking_shape_follows_exact_bedrock_entry_flag(
)
result = transform()
assert result.get("thinking") == {"type": "adaptive"}
assert result.get("thinking") == {"type": "adaptive", "display": "summarized"}
assert result.get("output_config") == {"effort": "medium"}
monkeypatch.setitem(litellm.model_cost[model], "supports_adaptive_thinking", False)

197
tests/test_litellm/llms/bedrock_mantle/passthrough/test_bedrock_mantle_passthrough_transformation.py Normal file
View file

@ -0,0 +1,197 @@
import json
from unittest.mock import MagicMock, patch
import httpx
import pytest
from botocore.credentials import Credentials
from litellm.llms.bedrock.passthrough.transformation import BedrockPassthroughConfig
from litellm.llms.bedrock_mantle.passthrough.transformation import BedrockMantlePassthroughConfig
from litellm.llms.custom_httpx.http_handler import HTTPHandler
from litellm.passthrough.main import llm_passthrough_route
from litellm.types.utils import LlmProviders
from litellm.utils import ProviderConfigManager
MANTLE_API_BASE = "https://bedrock-mantle.us-east-2.api.aws"
INVOKE_ENDPOINT = "model/us.openai.gpt-5.6-sol/invoke"
CONVERSE_ENDPOINT = "model/us.openai.gpt-5.6-sol/converse"
REQUEST_BODY = {"messages": [{"role": "user", "content": "say pong"}], "max_completion_tokens": 64}
@pytest.fixture
def no_ambient_aws(monkeypatch):
for name in (
"AWS_BEARER_TOKEN_BEDROCK",
"BEDROCK_MANTLE_API_KEY",
"BEDROCK_MANTLE_API_BASE",
"BEDROCK_MANTLE_REGION",
"AWS_BEDROCK_RUNTIME_ENDPOINT",
"AWS_REGION_NAME",
"AWS_REGION",
"AWS_DEFAULT_REGION",
):
monkeypatch.delenv(name, raising=False)
def test_bedrock_mantle_registers_its_own_bedrock_passthrough_config():
config = ProviderConfigManager.get_provider_passthrough_config(
model="us.openai.gpt-5.6-sol", provider=LlmProviders.BEDROCK_MANTLE
)
assert isinstance(config, BedrockMantlePassthroughConfig)
assert isinstance(config, BedrockPassthroughConfig)
def test_mantle_api_base_only_lends_its_region_to_the_runtime_url(no_ambient_aws):
url, base_url = BedrockMantlePassthroughConfig().get_complete_url(
api_base=MANTLE_API_BASE,
api_key=None,
model="us.openai.gpt-5.6-sol",
endpoint=INVOKE_ENDPOINT,
request_query_params=None,
litellm_params={"api_base": MANTLE_API_BASE},
)
assert str(url) == f"https://bedrock-runtime.us-east-2.amazonaws.com/{INVOKE_ENDPOINT}"
assert base_url == "https://bedrock-runtime.us-east-2.amazonaws.com"
def test_explicit_region_and_non_mantle_api_base_are_kept(no_ambient_aws):
vpc_endpoint = "https://vpce-0123.bedrock-runtime.us-east-1.vpce.amazonaws.com"
url, base_url = BedrockMantlePassthroughConfig().get_complete_url(
api_base=vpc_endpoint,
api_key=None,
model="us.openai.gpt-5.6-sol",
endpoint=INVOKE_ENDPOINT,
request_query_params=None,
litellm_params={"api_base": vpc_endpoint, "aws_region_name": "us-east-1"},
)
assert str(url) == f"{vpc_endpoint}/{INVOKE_ENDPOINT}"
assert base_url == vpc_endpoint
def test_region_falls_back_to_the_mantle_default_without_any_hint(no_ambient_aws):
url, _ = BedrockMantlePassthroughConfig().get_complete_url(
api_base=None,
api_key=None,
model="us.openai.gpt-5.6-sol",
endpoint=INVOKE_ENDPOINT,
request_query_params=None,
litellm_params={},
)
assert str(url) == f"https://bedrock-runtime.us-east-1.amazonaws.com/{INVOKE_ENDPOINT}"
@pytest.mark.parametrize(
("litellm_params", "env", "expected_bearer"),
[
({"api_key": "deployment-bedrock-api-key"}, {}, "deployment-bedrock-api-key"),
({}, {"BEDROCK_MANTLE_API_KEY": "mantle-env-key"}, "mantle-env-key"),
({}, {"AWS_BEARER_TOKEN_BEDROCK": "aws-env-key"}, "aws-env-key"),
],
)
def test_sign_request_uses_the_deployment_bearer_token(no_ambient_aws, monkeypatch, litellm_params, env, expected_bearer):
for name, value in env.items():
monkeypatch.setenv(name, value)
headers, body = BedrockMantlePassthroughConfig().sign_request(
headers={},
litellm_params=litellm_params,
request_data=REQUEST_BODY,
api_base=f"https://bedrock-runtime.us-east-1.amazonaws.com/{INVOKE_ENDPOINT}",
model="us.openai.gpt-5.6-sol",
)
assert headers["Authorization"] == f"Bearer {expected_bearer}"
assert body is not None
assert json.loads(body) == REQUEST_BODY
def test_sign_request_falls_back_to_sigv4_scoped_to_the_mantle_region(no_ambient_aws):
config = BedrockMantlePassthroughConfig()
with patch.object(config, "get_credentials", return_value=Credentials("AKIA", "secret")):
headers, body = config.sign_request(
headers={},
litellm_params={"api_base": MANTLE_API_BASE},
request_data=REQUEST_BODY,
api_base=f"https://bedrock-runtime.us-east-2.amazonaws.com/{INVOKE_ENDPOINT}",
model="us.openai.gpt-5.6-sol",
)
assert headers["Authorization"].startswith("AWS4-HMAC-SHA256 Credential=AKIA/")
assert "/us-east-2/bedrock/aws4_request" in headers["Authorization"]
assert body is not None
assert json.loads(body) == REQUEST_BODY
@pytest.mark.parametrize(
("route_kwargs", "env", "expected_bearer"),
[
({"api_key": "deployment-bedrock-api-key"}, {}, "deployment-bedrock-api-key"),
({}, {"BEDROCK_MANTLE_API_KEY": "mantle-env-key"}, "mantle-env-key"),
],
)
def test_invoke_passthrough_route_reaches_bedrock_runtime_for_a_mantle_deployment(
no_ambient_aws, monkeypatch, route_kwargs, env, expected_bearer
):
for name, value in env.items():
monkeypatch.setenv(name, value)
client = HTTPHandler()
with (
patch.object(client.client, "send", return_value=MagicMock(status_code=200)),
patch.object(client.client, "build_request", wraps=client.client.build_request) as build_request,
):
response = llm_passthrough_route(
model="bedrock_mantle/us.openai.gpt-5.6-sol",
endpoint=INVOKE_ENDPOINT,
method="POST",
api_base=MANTLE_API_BASE,
json=dict(REQUEST_BODY),
client=client,
litellm_logging_obj=MagicMock(),
**route_kwargs,
)
assert response.status_code == 200
sent = build_request.call_args.kwargs
assert str(sent["url"]) == f"https://bedrock-runtime.us-east-2.amazonaws.com/{INVOKE_ENDPOINT}"
assert sent["headers"]["Authorization"] == f"Bearer {expected_bearer}"
assert json.loads(sent["content"]) == REQUEST_BODY
def _logged_model_response(endpoint, body):
request = httpx.Request("POST", f"https://bedrock-runtime.us-east-1.amazonaws.com/{endpoint}")
return BedrockMantlePassthroughConfig().logging_non_streaming_response(
model="us.openai.gpt-5.6-sol",
custom_llm_provider="bedrock_mantle",
httpx_response=httpx.Response(200, json=body, request=request),
request_data={"messages": [{"role": "user", "content": [{"text": "say pong"}]}]},
logging_obj=MagicMock(),
endpoint=endpoint,
)
def test_converse_logging_parses_the_converse_response_shape():
result = _logged_model_response(
CONVERSE_ENDPOINT,
{
"metrics": {"latencyMs": 800.0},
"output": {"message": {"content": [{"text": "pong"}], "role": "assistant"}},
"stopReason": "end_turn",
"usage": {"inputTokens": 8, "outputTokens": 5, "totalTokens": 13},
},
)
assert result.choices[0].message.content == "pong"
assert result.usage.prompt_tokens == 8
assert result.usage.completion_tokens == 5
def test_invoke_logging_parses_the_openai_chat_response_shape():
result = _logged_model_response(
INVOKE_ENDPOINT,
{
"choices": [{"finish_reason": "stop", "index": 0, "message": {"content": "pong", "role": "assistant"}}],
"created": 1787677792,
"id": "chatcmpl-regression",
"model": "us.openai.gpt-5.6-sol",
"object": "chat.completion",
"usage": {"completion_tokens": 5, "prompt_tokens": 8, "total_tokens": 13},
},
)
assert result.choices[0].message.content == "pong"
assert result.usage.prompt_tokens == 8
assert result.usage.completion_tokens == 5

44
tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py
View file

@ -8,7 +8,8 @@ gate, the URL construction for both paths, and the shared Bearer auth.
"""
import copy
import json
from pathlib import Path
import pytest
from botocore.exceptions import (
@ -1523,7 +1524,7 @@ class TestBedrockMantleResponsesPricing:
assert info["cache_creation_input_token_cost"] == pytest.approx(cache_creation_cost)
assert info["cache_read_input_token_cost"] == pytest.approx(cache_read_cost)
assert info["output_cost_per_token"] == pytest.approx(output_cost)
assert info["max_input_tokens"] == 1000000
assert info["max_input_tokens"] == 1050000
assert info["input_cost_per_token_above_272k_tokens"] == pytest.approx(input_cost * 2)
assert info["cache_creation_input_token_cost_above_272k_tokens"] == pytest.approx(cache_creation_cost * 2)
assert info["cache_read_input_token_cost_above_272k_tokens"] == pytest.approx(cache_read_cost * 2)
@ -1565,3 +1566,42 @@ class TestBedrockMantleResponsesPricing:
def test_models_registered(self, local_cost_map):
assert "bedrock_mantle/openai.gpt-5.5" in litellm.bedrock_mantle_models
assert "bedrock_mantle/openai.gpt-5.4" in litellm.bedrock_mantle_models
def _repo_cost_map(map_name: str) -> dict[str, dict[str, object]]:
repo_root = Path(__file__).resolve().parents[4]
paths = {
"root": repo_root / "model_prices_and_context_window.json",
"bundled_backup": repo_root / "litellm" / "model_prices_and_context_window_backup.json",
}
return json.loads(paths[map_name].read_text())
class TestGpt56MantleRegistryEntries:
"""Locks the gpt-5.6 frontier entries to Bedrock Mantle's live behavior.
Mantle enforces a 1,050,000-token prompt maximum for gpt-5.6 sol/terra/luna
(oversize requests 400 with "prompt tokens (N) exceed model maximum
(1050000)", and a 1,030,590-token request completes), matching the OpenAI
Bedrock guide. mode must stay "responses": Mantle's native
/v1/chat/completions rejects function tools unless reasoning_effort is
"none", so chat traffic has to keep bridging to the Responses API
(see the responses_api_bridge tests above).
"""
@pytest.mark.parametrize("map_name", ("root", "bundled_backup"))
@pytest.mark.parametrize(
"key",
(
"bedrock_mantle/openai.gpt-5.6-sol",
"bedrock_mantle/openai.gpt-5.6-terra",
"bedrock_mantle/openai.gpt-5.6-luna",
),
)
def test_entry_matches_mantle_enforced_limits(self, map_name, key):
entry = _repo_cost_map(map_name)[key]
assert entry["max_input_tokens"] == 1050000
assert entry["max_output_tokens"] == 128000
assert entry["mode"] == "responses"
assert entry["use_openai_responses_path"] is True
assert entry["supported_endpoints"] == ["/v1/chat/completions", "/v1/responses"]

231
tests/test_litellm/llms/vertex_ai/interactions/test_vertex_ai_interactions_transformation.py Normal file
View file

@ -0,0 +1,231 @@
import pytest
import litellm
from litellm.interactions.utils import get_provider_interactions_api_config
from litellm.llms.gemini.interactions.transformation import (
GoogleAIStudioInteractionsConfig,
)
from litellm.llms.vertex_ai.interactions.transformation import (
VertexAIInteractionsConfig,
)
from litellm.types.llms.vertex_ai import VERTEX_CREDENTIALS_TYPES
from litellm.types.router import GenericLiteLLMParams
from litellm.types.utils import LlmProviders
GLOBAL_BASE = "https://aiplatform.googleapis.com/v1beta1/projects/test-proj/locations/global/interactions"
class MinterRecorder:
def __init__(self, resolved_project: str = "creds-proj") -> None:
self.calls: list[tuple[VERTEX_CREDENTIALS_TYPES | None, str | None]] = []
self.resolved_project = resolved_project
def __call__(
self,
credentials: VERTEX_CREDENTIALS_TYPES | None,
project_id: str | None,
) -> tuple[str, str]:
self.calls.append((credentials, project_id))
return "test-token", project_id or self.resolved_project
@pytest.fixture
def minter():
return MinterRecorder()
@pytest.fixture
def config(minter):
return VertexAIInteractionsConfig(mint_access_token=minter)
@pytest.fixture
def litellm_params():
return GenericLiteLLMParams(vertex_project="test-proj", vertex_credentials="creds.json")
class TestRegistration:
def test_vertex_ai_returns_vertex_config(self):
assert isinstance(get_provider_interactions_api_config("vertex_ai"), VertexAIInteractionsConfig)
def test_vertex_ai_beta_returns_vertex_config(self):
assert isinstance(get_provider_interactions_api_config("vertex_ai_beta"), VertexAIInteractionsConfig)
def test_gemini_still_returns_google_ai_studio_config(self):
gemini_config = get_provider_interactions_api_config("gemini")
assert isinstance(gemini_config, GoogleAIStudioInteractionsConfig)
assert not isinstance(gemini_config, VertexAIInteractionsConfig)
def test_lazy_import_resolves(self):
assert litellm.VertexAIInteractionsConfig is VertexAIInteractionsConfig
def test_custom_llm_provider_is_vertex_ai(self, config):
assert config.custom_llm_provider == LlmProviders.VERTEX_AI
class TestValidateEnvironment:
def test_sets_bearer_auth_without_gemini_headers(self, config, minter, litellm_params):
headers = config.validate_environment(
headers={},
model="gemini-omni-flash-preview",
litellm_params=litellm_params,
)
assert headers["Authorization"] == "Bearer test-token"
assert headers["Content-Type"] == "application/json"
assert "x-goog-api-key" not in headers
assert "Api-Revision" not in headers
assert minter.calls == [("creds.json", "test-proj")]
def test_caller_authorization_wins(self, config, litellm_params):
headers = config.validate_environment(
headers={"Authorization": "Bearer caller-token"},
model="gemini-omni-flash-preview",
litellm_params=litellm_params,
)
assert headers["Authorization"] == "Bearer caller-token"
class TestGetCompleteUrl:
def test_defaults_to_global_v1beta1(self, config, litellm_params):
url = config.get_complete_url(
api_base=None,
model="gemini-omni-flash-preview",
litellm_params=dict(litellm_params),
)
assert url == GLOBAL_BASE
def test_stream_appends_alt_sse(self, config, litellm_params):
url = config.get_complete_url(
api_base=None,
model="gemini-omni-flash-preview",
litellm_params=dict(litellm_params),
stream=True,
)
assert url == f"{GLOBAL_BASE}?alt=sse"
def test_multi_region_location_uses_rep_host(self, config):
url = config.get_complete_url(
api_base=None,
model="gemini-omni-flash-preview",
litellm_params={"vertex_project": "test-proj", "vertex_location": "us"},
)
assert url == "https://aiplatform.us.rep.googleapis.com/v1beta1/projects/test-proj/locations/us/interactions"
def test_regional_location_uses_regional_host(self, config):
url = config.get_complete_url(
api_base=None,
model="gemini-omni-flash-preview",
litellm_params={"vertex_project": "test-proj", "vertex_location": "us-central1"},
)
assert url == (
"https://us-central1-aiplatform.googleapis.com"
"/v1beta1/projects/test-proj/locations/us-central1/interactions"
)
def test_location_env_fallback_is_ignored(self, config, monkeypatch):
monkeypatch.setenv("VERTEXAI_LOCATION", "us-east5")
url = config.get_complete_url(
api_base=None,
model="gemini-omni-flash-preview",
litellm_params={"vertex_project": "test-proj"},
)
assert url == GLOBAL_BASE
def test_api_base_override(self, config, litellm_params):
url = config.get_complete_url(
api_base="https://proxy.example.test",
model="gemini-omni-flash-preview",
litellm_params=dict(litellm_params),
)
assert url == "https://proxy.example.test/v1beta1/projects/test-proj/locations/global/interactions"
def test_project_resolved_from_credentials_when_not_passed(self, config, monkeypatch):
monkeypatch.delenv("VERTEXAI_PROJECT", raising=False)
monkeypatch.setattr(litellm, "vertex_project", None)
url = config.get_complete_url(
api_base=None,
model="gemini-omni-flash-preview",
litellm_params={"vertex_credentials": "creds.json"},
)
assert url == "https://aiplatform.googleapis.com/v1beta1/projects/creds-proj/locations/global/interactions"
def test_invalid_location_rejected(self, config):
with pytest.raises(ValueError, match="Invalid vertex_location"):
config.get_complete_url(
api_base=None,
model="gemini-omni-flash-preview",
litellm_params={"vertex_project": "test-proj", "vertex_location": "evil.com#"},
)
def test_missing_project_rejected(self, monkeypatch):
monkeypatch.delenv("VERTEXAI_PROJECT", raising=False)
monkeypatch.setattr(litellm, "vertex_project", None)
def unresolved_minter(
credentials: VERTEX_CREDENTIALS_TYPES | None,
project_id: str | None,
) -> tuple[str, str]:
return "test-token", ""
with pytest.raises(ValueError, match="Vertex AI project is required"):
VertexAIInteractionsConfig(mint_access_token=unresolved_minter).get_complete_url(
api_base=None,
model="gemini-omni-flash-preview",
litellm_params={},
)
class TestInteractionByIdRequests:
def test_get_url(self, config, litellm_params):
url, request_body = config.transform_get_interaction_request(
interaction_id="abc123",
api_base="",
litellm_params=litellm_params,
headers={},
)
assert url == f"{GLOBAL_BASE}/abc123"
assert request_body == {}
def test_get_url_encodes_interaction_id(self, config, litellm_params):
url, _ = config.transform_get_interaction_request(
interaction_id="id/with space",
api_base="",
litellm_params=litellm_params,
headers={},
)
assert url == f"{GLOBAL_BASE}/id%2Fwith%20space"
def test_delete_url(self, config, litellm_params):
url, request_body = config.transform_delete_interaction_request(
interaction_id="abc123",
api_base="",
litellm_params=litellm_params,
headers={},
)
assert url == f"{GLOBAL_BASE}/abc123"
assert request_body == {}
def test_cancel_url(self, config, litellm_params):
url, request_body = config.transform_cancel_interaction_request(
interaction_id="abc123",
api_base="",
litellm_params=litellm_params,
headers={},
)
assert url == f"{GLOBAL_BASE}/abc123:cancel"
assert request_body == {}

2
tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_messages_config.py
View file

@ -538,7 +538,7 @@ def test_messages_thinking_shape_follows_exact_vertex_entry_flag(local_model_cos
)
result = transform()
assert result.get("thinking") == {"type": "adaptive"}
assert result.get("thinking") == {"type": "adaptive", "display": "summarized"}
assert result.get("output_config") == {"effort": "medium"}
monkeypatch.setitem(

305
tests/test_litellm/proxy/auth/test_auth_checks.py
View file

@ -47,6 +47,7 @@ from litellm.proxy.auth.auth_checks import (
_virtual_key_soft_budget_check,
get_key_object,
get_user_object,
invalidate_team_member_spend_state,
vector_store_access_check,
)
from litellm.caching.in_memory_cache import InMemoryCache
@ -6939,3 +6940,307 @@ def test_model_has_no_cost_mapping_alias_to_a_group_priced_through_model_info_is
router = _router_with_a_group_priced_through_model_info()
assert model_has_no_cost_mapping(model="model-info-priced-alias", llm_router=router) is False
@pytest.mark.asyncio
async def test_invalidate_team_member_spend_state_sets_the_spend_counter_and_clears_both_membership_cache_keys():
"""A team-member budget reset (new_spend passed) must SET the spend counter to the reset
value, clear its DB-floor marker, AND invalidate both independently-keyed membership caches
(user_api_key_auth.py's admission check writes one key format, budget_reservation.py and
auth_checks.py's own get_team_membership() write the other) or a stale read keeps 429ing
after the reset. Asserted against real cache reads, not mock call args, so a change that
keeps the call but drops its effect still fails."""
from litellm.caching.dual_cache import DualCache
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
real_cache = UserApiKeyCache()
await real_cache.async_set_cache(key="team-1_user-1", value="stale-membership")
await real_cache.async_set_cache(key="team_membership:user-1:team-1", value="stale-membership")
real_spend_counter_cache = DualCache()
real_spend_counter_cache.in_memory_cache.set_cache(key="spend:team_member:user-1:team-1", value=999.0)
real_spend_counter_cache.in_memory_cache.set_cache(
key="spend_db_floor:spend:team_member:user-1:team-1", value=999.0
)
with patch( # test-quality-ok: injects a real DualCache for the module global, not a behavior mock
"litellm.proxy.proxy_server.spend_counter_cache", real_spend_counter_cache
):
await invalidate_team_member_spend_state(
user_id="user-1",
team_id="team-1",
user_api_key_cache=real_cache,
new_spend=0.0,
)
assert await real_cache.async_get_cache(key="team-1_user-1") is None
assert await real_cache.async_get_cache(key="team_membership:user-1:team-1") is None
assert real_spend_counter_cache.in_memory_cache.get_cache(key="spend:team_member:user-1:team-1") == 0.0
assert (
real_spend_counter_cache.in_memory_cache.get_cache(key="spend_db_floor:spend:team_member:user-1:team-1")
== 0.0
), "the DB-floor marker kept the pre-reset value; a stale-floor read can raise the counter right back up"
@pytest.mark.asyncio
async def test_invalidate_team_member_spend_state_leaves_the_live_spend_counter_alone_without_new_spend():
"""team_member_update only changes the budget cap, not the tracked spend, so it calls
invalidate_team_member_spend_state with no new_spend. Deleting the live spend counter in that
case would force the next read to reseed from the DB's own spend column, which lags the live
counter via periodic batch writes, briefly UNDER-enforcing the raised cap against a spend
value lower than what was actually tracked (regression: PR #37971 Bugbot finding)."""
from litellm.caching.dual_cache import DualCache
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
real_cache = UserApiKeyCache()
await real_cache.async_set_cache(key="team-1_user-1", value="stale-membership")
real_spend_counter_cache = DualCache()
real_spend_counter_cache.in_memory_cache.set_cache(key="spend:team_member:user-1:team-1", value=999.0)
with patch( # test-quality-ok: injects a real DualCache for the module global, not a behavior mock
"litellm.proxy.proxy_server.spend_counter_cache", real_spend_counter_cache
):
await invalidate_team_member_spend_state(
user_id="user-1",
team_id="team-1",
user_api_key_cache=real_cache,
)
assert await real_cache.async_get_cache(key="team-1_user-1") is None
assert real_spend_counter_cache.in_memory_cache.get_cache(key="spend:team_member:user-1:team-1") == 999.0
@pytest.mark.asyncio
async def test_invalidate_team_member_spend_state_sets_new_spend_instead_of_deleting():
"""/key/{key}/reset_spend SETs its counter to the reset value rather than deleting it, so a
worker's next read reflects it directly instead of falling back through a DB reseed. A reset
caller passing new_spend must match that precedent, not merely delete the counter."""
from litellm.caching.dual_cache import DualCache
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
real_cache = UserApiKeyCache()
real_spend_counter_cache = DualCache()
fake_redis_cache = MagicMock()
fake_redis_cache.async_set_cache = AsyncMock()
real_spend_counter_cache.redis_cache = fake_redis_cache
with patch( # test-quality-ok: injects a real DualCache for the module global, not a behavior mock
"litellm.proxy.proxy_server.spend_counter_cache", real_spend_counter_cache
):
await invalidate_team_member_spend_state(
user_id="user-1",
team_id="team-1",
user_api_key_cache=real_cache,
new_spend=2.5,
)
assert real_spend_counter_cache.in_memory_cache.get_cache(key="spend:team_member:user-1:team-1") == 2.5
fake_redis_cache.async_set_cache.assert_awaited_once_with(key="spend:team_member:user-1:team-1", value=2.5, ttl=60)
@pytest.mark.asyncio
async def test_invalidate_team_member_spend_state_deletes_redis_counter_when_set_fails(): # test-quality-ok: only observable effect is the fallback call on the same fake client
"""Redis reads take priority over the local in-memory copy (get_current_spend reads Redis
first), so a failed Redis SET would otherwise leave the OLD pre-reset value authoritative
for every worker even though the reset reported success. On a failed SET, the stale Redis
entry must be deleted instead, so the next read clean-misses and reseeds from the DB."""
from litellm.caching.dual_cache import DualCache
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
real_cache = UserApiKeyCache()
real_spend_counter_cache = DualCache()
fake_redis_cache = MagicMock()
fake_redis_cache.async_set_cache = AsyncMock(side_effect=ConnectionError("redis down"))
fake_redis_cache.async_delete_cache = AsyncMock()
real_spend_counter_cache.redis_cache = fake_redis_cache
with patch( # test-quality-ok: injects a real DualCache for the module global, not a behavior mock
"litellm.proxy.proxy_server.spend_counter_cache", real_spend_counter_cache
):
await invalidate_team_member_spend_state(
user_id="user-1",
team_id="team-1",
user_api_key_cache=real_cache,
new_spend=2.5,
)
fake_redis_cache.async_delete_cache.assert_awaited_once_with(key="spend:team_member:user-1:team-1")
@pytest.mark.asyncio
async def test_invalidate_team_member_spend_state_raises_503_when_both_redis_writes_fail():
"""If the Redis SET fails AND the fallback DELETE fails, the stale pre-reset counter is still
authoritative in Redis for every worker. Reporting success would silently keep 429ing the
member, so the reset must surface a 503 instead (regression: PR #37971 Greptile finding)."""
from fastapi import HTTPException
from litellm.caching.dual_cache import DualCache
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
real_cache = UserApiKeyCache()
real_spend_counter_cache = DualCache()
fake_redis_cache = MagicMock()
fake_redis_cache.async_set_cache = AsyncMock(side_effect=ConnectionError("redis down"))
fake_redis_cache.async_delete_cache = AsyncMock(side_effect=ConnectionError("redis still down"))
real_spend_counter_cache.redis_cache = fake_redis_cache
with (
patch( # test-quality-ok: injects a real DualCache for the module global, not a behavior mock
"litellm.proxy.proxy_server.spend_counter_cache", real_spend_counter_cache
),
pytest.raises(HTTPException) as exc_info,
):
await invalidate_team_member_spend_state(
user_id="user-1",
team_id="team-1",
user_api_key_cache=real_cache,
new_spend=2.5,
)
assert exc_info.value.status_code == 503
@pytest.mark.asyncio
async def test_invalidate_team_member_spend_state_broadcasts_the_spend_counter_to_remote_workers():
"""The test above only proves the handling worker's own spend counter is
cleared. A remote worker's spend counter is a separate DualCache instance;
if the reset never reaches it, that worker keeps enforcing the pre-reset
spend the moment its own Redis read for the counter fails and it falls
back to its own (now-stale) in-memory copy. Drives the actual message
published onto the invalidation channel through a second, independent
AuthCacheInvalidationSubscriber standing in for that remote worker, rather
than asserting on the publish call args."""
from redis.asyncio import Redis
from litellm.caching.dual_cache import DualCache
from litellm.caching.in_memory_cache import InMemoryCache
from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import AuthCacheInvalidationSubscriber
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
published: list[tuple[str, str]] = []
class _RecordingRedisClient(Redis):
def __init__(self) -> None:
pass
async def publish(self, channel: str, message: str) -> int:
published.append((channel, message))
return 1
class _FakeRedisCache:
def __init__(self) -> None:
self.namespace = None
def init_async_client(self) -> object:
return _RecordingRedisClient()
local_spend_counter_cache = DualCache()
remote_user_api_key_cache = UserApiKeyCache()
remote_spend_counter_in_memory_cache = InMemoryCache()
remote_spend_counter_in_memory_cache.set_cache("spend:team_member:user-1:team-1", 999.0)
remote_spend_counter_in_memory_cache.set_cache("spend_db_floor:spend:team_member:user-1:team-1", 999.0)
with (
patch( # test-quality-ok: injects a real DualCache for the module global, not a behavior mock
"litellm.proxy.proxy_server.spend_counter_cache", local_spend_counter_cache
),
patch( # test-quality-ok: injects a fake pub/sub-capable redis cache; no live redis in this unit test
"litellm.proxy.common_utils.auth_cache_invalidation_pubsub.coordination_redis_cache",
return_value=_FakeRedisCache(),
),
):
await invalidate_team_member_spend_state(
user_id="user-1",
team_id="team-1",
user_api_key_cache=UserApiKeyCache(),
new_spend=0.0,
)
def _published_message_for(cache_key: str) -> str:
matches = [message for _, message in published if json.loads(message)["cache_key"] == cache_key]
assert matches, f"{cache_key} never reached the cross-worker invalidation channel"
return matches[-1]
remote_subscriber = AuthCacheInvalidationSubscriber(
redis_cache=_FakeRedisCache(),
user_api_key_cache=remote_user_api_key_cache,
additional_in_memory_caches=(remote_spend_counter_in_memory_cache,),
)
for cache_key in ("spend:team_member:user-1:team-1", "spend_db_floor:spend:team_member:user-1:team-1"):
remote_subscriber._apply_message( # pyright: ignore[reportPrivateUsage] # exercising the real cross-worker message handler, not a public API
{"type": "message", "data": _published_message_for(cache_key)}
)
assert remote_spend_counter_in_memory_cache.get_cache("spend:team_member:user-1:team-1") == 0.0
assert (
remote_spend_counter_in_memory_cache.get_cache("spend_db_floor:spend:team_member:user-1:team-1") == 0.0
), "the DB-floor marker was not broadcast; a remote worker can re-raise the counter off its stale floor"
@pytest.mark.asyncio
async def test_invalidate_team_member_spend_state_self_delivered_broadcast_does_not_erase_the_reset():
"""The handling worker subscribes to the same invalidation channel it publishes on, so it
receives its own reset message. A delete-style broadcast would erase the post-reset counter
and floor marker the handler just wrote, reopening the stale-floor race the reset closed
(regression: PR #37971 Greptile finding). The broadcast carries the reset value as a SET, so
applying the self-delivered message must leave both keys at the post-reset value."""
from redis.asyncio import Redis
from litellm.caching.dual_cache import DualCache
from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import AuthCacheInvalidationSubscriber
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
published: list[tuple[str, str]] = []
class _RecordingRedisClient(Redis):
def __init__(self) -> None:
pass
async def publish(self, channel: str, message: str) -> int:
published.append((channel, message))
return 1
class _FakeRedisCache:
def __init__(self) -> None:
self.namespace = None
def init_async_client(self) -> object:
return _RecordingRedisClient()
local_spend_counter_cache = DualCache()
local_user_api_key_cache = UserApiKeyCache()
with (
patch( # test-quality-ok: injects a real DualCache for the module global, not a behavior mock
"litellm.proxy.proxy_server.spend_counter_cache", local_spend_counter_cache
),
patch( # test-quality-ok: injects a fake pub/sub-capable redis cache; no live redis in this unit test
"litellm.proxy.common_utils.auth_cache_invalidation_pubsub.coordination_redis_cache",
return_value=_FakeRedisCache(),
),
):
await invalidate_team_member_spend_state(
user_id="user-1",
team_id="team-1",
user_api_key_cache=local_user_api_key_cache,
new_spend=0.0,
)
own_subscriber = AuthCacheInvalidationSubscriber(
redis_cache=_FakeRedisCache(),
user_api_key_cache=local_user_api_key_cache,
additional_in_memory_caches=(local_spend_counter_cache.in_memory_cache,),
)
for _, message in published:
own_subscriber._apply_message( # pyright: ignore[reportPrivateUsage] # exercising the real cross-worker message handler, not a public API
{"type": "message", "data": message}
)
assert local_spend_counter_cache.in_memory_cache.get_cache("spend:team_member:user-1:team-1") == 0.0, (
"the handler's self-delivered broadcast erased the post-reset spend counter"
)
assert (
local_spend_counter_cache.in_memory_cache.get_cache("spend_db_floor:spend:team_member:user-1:team-1") == 0.0
), "the handler's self-delivered broadcast erased the post-reset floor marker, reopening the stale-floor race"

32
tests/test_litellm/proxy/common_utils/test_auth_cache_invalidation_pubsub.py
View file

@ -6,6 +6,7 @@ from unittest.mock import patch
import pytest
from redis.asyncio import Redis
from litellm.caching.in_memory_cache import InMemoryCache
from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import (
AUTH_CACHE_INVALIDATION_CHANNEL,
AuthCacheInvalidationSubscriber,
@ -144,6 +145,37 @@ async def test_subscriber_deletes_local_cache_entry_on_message() -> None:
assert pubsub.subscribed_channels == [AUTH_CACHE_INVALIDATION_CHANNEL]
@pytest.mark.asyncio
async def test_subscriber_deletes_additional_in_memory_cache_entry_on_message() -> None:
"""
The spend-counter half of the same cross-worker gap: a remote worker's own
spend counter can hold a stale value (its fallback path when that worker's
own Redis read for the counter fails), and only clearing user_api_key_cache
on message would leave that separate DualCache's in-memory copy untouched.
"""
cache = UserApiKeyCache()
spend_counter_in_memory_cache = InMemoryCache()
spend_counter_in_memory_cache.set_cache("spend:team_member:u-1:t-1", 999.0)
assert spend_counter_in_memory_cache.get_cache("spend:team_member:u-1:t-1") is not None
pubsub = _QueuePubSub(initial_messages=[_invalidation_message("spend:team_member:u-1:t-1")])
subscriber = AuthCacheInvalidationSubscriber(
redis_cache=_FakeRedisCache(client=_ScriptedPubSubRedisClient(pubsubs=[pubsub])),
user_api_key_cache=cache,
additional_in_memory_caches=(spend_counter_in_memory_cache,),
)
subscriber.start()
try:
for _ in range(200):
if spend_counter_in_memory_cache.get_cache("spend:team_member:u-1:t-1") is None:
break
await asyncio.sleep(0.01)
finally:
await subscriber.stop()
assert spend_counter_in_memory_cache.get_cache("spend:team_member:u-1:t-1") is None
@pytest.mark.asyncio
async def test_subscriber_ignores_malformed_messages() -> None:
cache = UserApiKeyCache()

380
tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py
View file

@ -9,11 +9,13 @@ from unittest.mock import AsyncMock, MagicMock, call, patch
import pytest
from fastapi import HTTPException
from fastapi.testclient import TestClient
from pydantic import ValidationError
from litellm._uuid import uuid
from litellm.proxy._types import UserAPIKeyAuth # Import UserAPIKeyAuth
from litellm.proxy._types import (
LiteLLM_BudgetTable,
LiteLLM_BudgetTableFull,
LiteLLM_ModelTable,
LiteLLM_OrganizationMembershipTable,
@ -27,7 +29,9 @@ from litellm.proxy._types import (
Member,
ProxyErrorTypes,
ProxyException,
ResetSpendRequest,
TeamMemberAddRequest,
TeamMemberUpdateRequest,
UpdateTeamRequest,
)
from litellm.proxy.management_endpoints.team_endpoints import (
@ -42,12 +46,15 @@ from litellm.proxy.management_endpoints.team_endpoints import (
_transform_teams_to_deleted_records,
_update_model_table,
_validate_and_populate_member_user_info,
_validate_team_member_reset_spend_value,
_verify_team_access,
delete_team,
list_available_teams,
reset_team_member_spend_fn,
router,
team_member_add_duplication_check,
team_member_delete,
team_member_update,
update_team,
validate_team_org_change,
)
@ -12603,3 +12610,376 @@ async def test_invalidate_access_group_cache_deletes_the_cached_object():
"user_api_key_cache": cache,
"proxy_logging_obj": logging_obj,
}
def test_validate_team_member_reset_spend_value_rejects_non_numeric():
with pytest.raises(HTTPException) as exc:
_validate_team_member_reset_spend_value(
reset_to="not-a-number",
membership=LiteLLM_TeamMembership(user_id="u1", team_id="t1", spend=10.0),
)
assert exc.value.status_code == 400
def test_validate_team_member_reset_spend_value_rejects_negative():
with pytest.raises(HTTPException) as exc:
_validate_team_member_reset_spend_value(
reset_to=-1.0,
membership=LiteLLM_TeamMembership(user_id="u1", team_id="t1", spend=10.0),
)
assert exc.value.status_code == 400
@pytest.mark.parametrize("reset_to", [float("nan"), float("inf"), float("-inf")])
def test_validate_team_member_reset_spend_value_rejects_non_finite(reset_to):
"""NaN and +/-inf are instances of float and compare False against every bound
below (`nan < 0`, `nan > current_spend` are both False), so an isinstance-and-range
check alone lets them through to persist as the member's spend and silently
disable every later budget comparison against it."""
with pytest.raises(HTTPException) as exc:
_validate_team_member_reset_spend_value(
reset_to=reset_to,
membership=LiteLLM_TeamMembership(user_id="u1", team_id="t1", spend=10.0),
)
assert exc.value.status_code == 400
@pytest.mark.parametrize("reset_to", [True, False])
def test_reset_spend_request_rejects_bool_reset_to(reset_to):
"""bool is a subclass of int, so pydantic silently coerces True/False into 1.0/0.0 for a
``float`` field: {"reset_to": true} would otherwise reach _validate_team_member_reset_spend_value
as an indistinguishable 1.0 and reset the member's spend instead of failing the request."""
with pytest.raises(ValidationError):
ResetSpendRequest(reset_to=reset_to)
def test_validate_team_member_reset_spend_value_rejects_above_current_spend():
with pytest.raises(HTTPException) as exc:
_validate_team_member_reset_spend_value(
reset_to=20.0,
membership=LiteLLM_TeamMembership(user_id="u1", team_id="t1", spend=10.0),
)
assert exc.value.status_code == 400
def test_validate_team_member_reset_spend_value_rejects_above_max_budget():
with pytest.raises(HTTPException) as exc:
_validate_team_member_reset_spend_value(
reset_to=10.0,
membership=LiteLLM_TeamMembership(
user_id="u1",
team_id="t1",
spend=10.0,
litellm_budget_table=LiteLLM_BudgetTable(budget_id="b1", max_budget=5.0),
),
)
assert exc.value.status_code == 400
def test_validate_team_member_reset_spend_value_accepts_valid_reset():
result = _validate_team_member_reset_spend_value(
reset_to=0.0,
membership=LiteLLM_TeamMembership(user_id="u1", team_id="t1", spend=10.0),
)
assert result == 0.0
@pytest.mark.asyncio
async def test_reset_team_member_spend_fn_success(monkeypatch):
"""A proxy admin resetting a stuck team member's spend must write the DB
row to reset_to AND invalidate the cached spend/membership state, or the
429 the endpoint exists to clear keeps firing off the stale cache.
Asserted against real cache reads, not mock call args, so a change that
keeps the call but drops its effect still fails."""
from litellm.caching.dual_cache import DualCache
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
mock_prisma_client = MagicMock()
mock_proxy_logging_obj = MagicMock()
real_cache = UserApiKeyCache()
await real_cache.async_set_cache(key="team-1_member-1", value="stale-membership")
await real_cache.async_set_cache(key="team_membership:member-1:team-1", value="stale-membership")
real_spend_counter_cache = DualCache()
real_spend_counter_cache.in_memory_cache.set_cache(key="spend:team_member:member-1:team-1", value=999.0)
membership_row = LiteLLM_TeamMembership(
user_id="member-1",
team_id="team-1",
spend=10.0,
litellm_budget_table=LiteLLM_BudgetTable(budget_id="b1", max_budget=50.0),
)
mock_prisma_client.db.litellm_teammembership.find_unique = AsyncMock(return_value=membership_row)
mock_prisma_client.db.litellm_teammembership.update = AsyncMock(return_value=membership_row)
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", real_cache)
monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging_obj)
monkeypatch.setattr("litellm.proxy.proxy_server.spend_counter_cache", real_spend_counter_cache)
with patch( # test-quality-ok: no live DB here; matches this file's established convention for endpoint-logic unit tests
"litellm.proxy.management_endpoints.team_endpoints.get_team_object",
AsyncMock(return_value=LiteLLM_TeamTable(team_id="team-1")),
):
response = await reset_team_member_spend_fn(
team_id="team-1",
user_id="member-1",
data=ResetSpendRequest(reset_to=0.0),
user_api_key_dict=UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin", user_id="admin-user"
),
)
assert response["spend"] == 0.0
assert response["previous_spend"] == 10.0
assert response["max_budget"] == 50.0
mock_prisma_client.db.litellm_teammembership.update.assert_awaited_once_with(
where={"user_id_team_id": {"user_id": "member-1", "team_id": "team-1"}},
data={"spend": 0.0},
)
assert await real_cache.async_get_cache(key="team-1_member-1") is None
assert await real_cache.async_get_cache(key="team_membership:member-1:team-1") is None
assert real_spend_counter_cache.in_memory_cache.get_cache(key="spend:team_member:member-1:team-1") == 0.0
@pytest.mark.asyncio
async def test_reset_team_member_spend_fn_membership_not_found(monkeypatch):
mock_prisma_client = MagicMock()
mock_prisma_client.db.litellm_teammembership.find_unique = AsyncMock(return_value=None)
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", MagicMock())
monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock())
with patch( # test-quality-ok: no live DB here; matches this file's established convention for endpoint-logic unit tests
"litellm.proxy.management_endpoints.team_endpoints.get_team_object",
AsyncMock(return_value=LiteLLM_TeamTable(team_id="team-1")),
):
with pytest.raises(HTTPException) as exc:
await reset_team_member_spend_fn(
team_id="team-1",
user_id="ghost-user",
data=ResetSpendRequest(reset_to=0.0),
user_api_key_dict=UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin", user_id="admin-user"
),
)
assert exc.value.status_code == 404
@pytest.mark.asyncio
async def test_reset_team_member_spend_fn_team_not_found(monkeypatch):
mock_prisma_client = MagicMock()
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", MagicMock())
monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock())
with patch( # test-quality-ok: no live DB here; matches this file's established convention for endpoint-logic unit tests
"litellm.proxy.management_endpoints.team_endpoints.get_team_object",
AsyncMock(side_effect=HTTPException(status_code=404, detail={"error": "Team doesn't exist in db."})),
):
with pytest.raises(HTTPException) as exc:
await reset_team_member_spend_fn(
team_id="ghost-team",
user_id="member-1",
data=ResetSpendRequest(reset_to=0.0),
user_api_key_dict=UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin", user_id="admin-user"
),
)
assert exc.value.status_code == 404
@pytest.mark.asyncio
async def test_reset_team_member_spend_fn_forbidden_for_non_admin(monkeypatch):
"""A caller who is neither proxy admin, org admin, nor this team's admin must be refused,
matching every other team-mutating endpoint's authorization."""
mock_prisma_client = MagicMock()
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", MagicMock())
monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock())
with patch( # test-quality-ok: no live DB here; matches this file's established convention for endpoint-logic unit tests
"litellm.proxy.management_endpoints.team_endpoints.get_team_object",
AsyncMock(return_value=LiteLLM_TeamTable(team_id="team-1", members_with_roles=[])),
):
with pytest.raises(HTTPException) as exc:
await reset_team_member_spend_fn(
team_id="team-1",
user_id="member-1",
data=ResetSpendRequest(reset_to=0.0),
user_api_key_dict=UserAPIKeyAuth(
user_role=LitellmUserRoles.INTERNAL_USER, api_key="sk-user", user_id="plain-user"
),
)
assert exc.value.status_code == 403
@pytest.mark.asyncio
async def test_reset_team_member_spend_fn_team_admin_cannot_reset_own_spend(monkeypatch):
"""_verify_team_access authorizes a team admin over their own team with no check that the
target differs from the caller. Unchecked, that admin could target their own membership row
and repeatedly zero it right before it crosses their per-member cap, consuming the shared
team budget without the configured limit ever binding (Veria finding on PR #37971)."""
mock_prisma_client = MagicMock()
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", MagicMock())
monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock())
team_admin = UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, api_key="sk-admin", user_id="team-admin-1")
with patch( # test-quality-ok: no live DB here; matches this file's established convention for endpoint-logic unit tests
"litellm.proxy.management_endpoints.team_endpoints.get_team_object",
AsyncMock(
return_value=LiteLLM_TeamTable(
team_id="team-1",
members_with_roles=[Member(user_id="team-admin-1", role="admin")],
)
),
):
with pytest.raises(HTTPException) as exc:
await reset_team_member_spend_fn(
team_id="team-1",
user_id="team-admin-1",
data=ResetSpendRequest(reset_to=0.0),
user_api_key_dict=team_admin,
)
assert exc.value.status_code == 403
mock_prisma_client.db.litellm_teammembership.update.assert_not_called()
@pytest.mark.asyncio
async def test_reset_team_member_spend_fn_proxy_admin_can_reset_own_spend(monkeypatch):
"""The self-reset guard is scoped to non-proxy-admin roles: a proxy admin resetting their
own membership spend is the platform-wide trust boundary, not a team-scoped one."""
mock_prisma_client = MagicMock()
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", MagicMock())
monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock())
membership_row = LiteLLM_TeamMembership(user_id="admin-user", team_id="team-1", spend=10.0)
mock_prisma_client.db.litellm_teammembership.find_unique = AsyncMock(return_value=membership_row)
mock_prisma_client.db.litellm_teammembership.update = AsyncMock(return_value=membership_row)
with patch( # test-quality-ok: no live DB here; matches this file's established convention for endpoint-logic unit tests
"litellm.proxy.management_endpoints.team_endpoints.get_team_object",
AsyncMock(return_value=LiteLLM_TeamTable(team_id="team-1")),
):
response = await reset_team_member_spend_fn(
team_id="team-1",
user_id="admin-user",
data=ResetSpendRequest(reset_to=0.0),
user_api_key_dict=UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin", user_id="admin-user"
),
)
assert response["spend"] == 0.0
@pytest.mark.asyncio
async def test_team_member_update_invalidates_team_member_spend_state_when_budget_patch_applied(monkeypatch):
"""Raising a stuck member's max_budget_in_team via the documented /team/member_update
endpoint must invalidate the cached membership state, or the raised cap never reaches the
admission check and the member stays 429ing. The live spend counter itself must be left
untouched: only the cap changed, and deleting the counter would force a reseed from the
DB's own spend column, which lags the live counter via periodic batch writes, briefly
UNDER-enforcing the raised cap against a spend value lower than what was actually tracked.
Asserted against real cache reads, not mock call args, so a change that keeps the call but
drops its effect still fails."""
from litellm.caching.dual_cache import DualCache
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
mock_prisma_client = MagicMock()
real_cache = UserApiKeyCache()
await real_cache.async_set_cache(key="team-1_member-1", value="stale-membership")
await real_cache.async_set_cache(key="team_membership:member-1:team-1", value="stale-membership")
real_spend_counter_cache = DualCache()
real_spend_counter_cache.in_memory_cache.set_cache(key="spend:team_member:member-1:team-1", value=999.0)
team_row = LiteLLM_TeamTable(team_id="team-1", metadata={}, members_with_roles=[])
team_info_response = {
"team_info": team_row,
"team_memberships": [LiteLLM_TeamMembership(user_id="member-1", team_id="team-1", budget_id=None)],
}
mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=team_row)
monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True)
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", real_cache)
monkeypatch.setattr("litellm.proxy.proxy_server.spend_counter_cache", real_spend_counter_cache)
mock_tx = AsyncMock()
mock_prisma_client.tx.return_value.__aenter__ = AsyncMock(return_value=mock_tx)
mock_prisma_client.tx.return_value.__aexit__ = AsyncMock(return_value=None)
with (
patch( # test-quality-ok: no live DB here; matches this file's established convention for endpoint-logic unit tests
"litellm.proxy.management_endpoints.team_endpoints.team_info",
AsyncMock(return_value=team_info_response),
),
patch( # test-quality-ok: no live DB here; matches this file's established convention for endpoint-logic unit tests
"litellm.proxy.management_endpoints.team_endpoints._upsert_budget_and_membership",
AsyncMock(),
),
):
await team_member_update(
data=TeamMemberUpdateRequest(team_id="team-1", user_id="member-1", max_budget_in_team=999999.0),
http_request=MagicMock(),
user_api_key_dict=UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin", user_id="admin-user"
),
)
assert await real_cache.async_get_cache(key="team-1_member-1") is None
assert await real_cache.async_get_cache(key="team_membership:member-1:team-1") is None
assert real_spend_counter_cache.in_memory_cache.get_cache(key="spend:team_member:member-1:team-1") == 999.0
@pytest.mark.asyncio
async def test_team_member_update_skips_invalidation_when_no_budget_fields_sent(monkeypatch):
"""A role-only update carries an empty budget_patch and touches no budget state,
so the member's cached spend/membership state must be left untouched."""
from litellm.caching.dual_cache import DualCache
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
mock_prisma_client = MagicMock()
real_cache = UserApiKeyCache()
await real_cache.async_set_cache(key="team-1_member-1", value="still-fresh-membership")
real_spend_counter_cache = DualCache()
real_spend_counter_cache.in_memory_cache.set_cache(key="spend:team_member:member-1:team-1", value=1.5)
team_row = LiteLLM_TeamTable(team_id="team-1", metadata={}, members_with_roles=[])
team_info_response = {
"team_info": team_row,
"team_memberships": [LiteLLM_TeamMembership(user_id="member-1", team_id="team-1", budget_id=None)],
}
mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=team_row)
monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True)
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", real_cache)
monkeypatch.setattr("litellm.proxy.proxy_server.spend_counter_cache", real_spend_counter_cache)
mock_tx = AsyncMock()
mock_prisma_client.tx.return_value.__aenter__ = AsyncMock(return_value=mock_tx)
mock_prisma_client.tx.return_value.__aexit__ = AsyncMock(return_value=None)
with (
patch( # test-quality-ok: no live DB here; matches this file's established convention for endpoint-logic unit tests
"litellm.proxy.management_endpoints.team_endpoints.team_info",
AsyncMock(return_value=team_info_response),
),
patch( # test-quality-ok: no live DB here; matches this file's established convention for endpoint-logic unit tests
"litellm.proxy.management_endpoints.team_endpoints._upsert_budget_and_membership",
AsyncMock(),
),
):
await team_member_update(
data=TeamMemberUpdateRequest(team_id="team-1", user_id="member-1"),
http_request=MagicMock(),
user_api_key_dict=UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin", user_id="admin-user"
),
)
assert await real_cache.async_get_cache(key="team-1_member-1") == "still-fresh-membership"
assert real_spend_counter_cache.in_memory_cache.get_cache(key="spend:team_member:member-1:team-1") == 1.5

35
tests/test_litellm/proxy/test_proxy_server.py
View file

@ -11347,3 +11347,38 @@ class TestRouterModelNameOnStreamingChunks:
assert len(frames) >= 3
assert '"router_model_name":"deep-model"' in frames[0]
assert all('"router_model_name":"backup-tier"' in frame for frame in frames[1:])
@pytest.mark.asyncio
async def test_authoritative_floor_spend_keeps_a_reset_marker_written_during_the_db_read():
"""A team-member spend reset writes the post-reset floor to the spend_db_floor marker
(auth_checks.invalidate_team_member_spend_state). A floor read already in flight when the
reset commits would otherwise cache its stale pre-reset DB value over the fresh marker,
letting a budget check raise the counter right back above the just-reset spend
(regression: PR #37971 Greptile finding)."""
from litellm.proxy.proxy_server import _authoritative_floor_spend
real_spend_counter_cache = DualCache()
counter_key = "spend:team_member:user-1:team-1"
marker_key = f"spend_db_floor:{counter_key}"
async def db_read_racing_with_a_reset(prisma_client, counter_key):
real_spend_counter_cache.in_memory_cache.set_cache(key=marker_key, value=0.0)
return 999.0
with (
patch.object( # test-quality-ok: injects a real DualCache for the module global, not a behavior mock
proxy_server_module, "spend_counter_cache", real_spend_counter_cache
),
patch.object( # test-quality-ok: the DB read must race the reset; no injectable seam for module-global prisma reads
proxy_server_module.SpendCounterReseed,
"from_db",
AsyncMock(side_effect=db_read_racing_with_a_reset),
),
):
result = await _authoritative_floor_spend(counter_key=counter_key)
assert result == 0.0
assert real_spend_counter_cache.in_memory_cache.get_cache(key=marker_key) == 0.0, (
"the in-flight DB read clobbered the post-reset floor marker with the stale pre-reset value"
)

67
tests/test_litellm/rerank_api/test_main.py
View file

@ -1,6 +1,10 @@
import logging
from unittest.mock import MagicMock, patch
import httpx
import pytest
import respx
import litellm
@ -62,3 +66,66 @@ def test_rerank_does_not_log_request_content_at_info(caplog):
assert all(
r.levelno == logging.DEBUG for r in optional_params_logs
), "optional_rerank_params must be logged at DEBUG, not INFO"
TOGETHER_RERANK_BODY = {
"id": "rerank-mock-id",
"results": [{"index": 0, "relevance_score": 0.95}],
"usage": {"prompt_tokens": 10, "total_tokens": 10},
}
def test_together_rerank_defaults_to_together_ai_host(respx_mock: respx.MockRouter, monkeypatch):
"""Regression for the Together host migration: rerank used to hardcode
https://api.together.xyz/v1/rerank. The default must now be api.together.ai."""
monkeypatch.delenv("TOGETHER_AI_API_BASE", raising=False)
mock_route = respx_mock.post("https://api.together.ai/v1/rerank")
mock_route.return_value = httpx.Response(200, json=TOGETHER_RERANK_BODY)
response = litellm.rerank(
model="together_ai/mixedbread-ai/mxbai-rerank-large-v2",
query=MARKER_QUERY,
documents=[MARKER_DOC],
api_key="fake-together-key",
)
assert mock_route.called
assert response.results[0]["relevance_score"] == 0.95
def test_together_rerank_honors_api_base(respx_mock: respx.MockRouter):
"""Regression: a custom api_base was silently ignored by the Together rerank handler."""
mock_route = respx_mock.post("https://custom-together.example/v1/rerank")
mock_route.return_value = httpx.Response(200, json=TOGETHER_RERANK_BODY)
litellm.rerank(
model="together_ai/mixedbread-ai/mxbai-rerank-large-v2",
query=MARKER_QUERY,
documents=[MARKER_DOC],
api_key="fake-together-key",
api_base="https://custom-together.example/v1",
)
assert mock_route.called
assert mock_route.calls[0].request.headers["authorization"] == "Bearer fake-together-key"
@pytest.mark.asyncio
async def test_together_rerank_async_honors_env_api_base(respx_mock: respx.MockRouter, monkeypatch):
"""Regression: TOGETHER_AI_API_BASE was honored by chat but ignored by rerank."""
monkeypatch.setenv("TOGETHER_AI_API_BASE", "https://env-together.example/v1")
monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True")
mock_route = respx_mock.post("https://env-together.example/v1/rerank")
mock_route.return_value = httpx.Response(200, json=TOGETHER_RERANK_BODY)
response = await litellm.arerank(
model="together_ai/mixedbread-ai/mxbai-rerank-large-v2",
query=MARKER_QUERY,
documents=[MARKER_DOC],
api_key="fake-together-key",
)
assert mock_route.called
assert response.results[0]["relevance_score"] == 0.95

13
tests/test_litellm/test_main.py
View file

@ -2944,3 +2944,16 @@ def test_a_stream_that_reported_no_usage_is_still_billed(local_cost_map):
assert cost == pytest.approx(
_priced_at(rebuilt.usage.prompt_tokens, rebuilt.usage.completion_tokens)
)
@pytest.mark.asyncio
async def test_acompletion_resolves_provider_from_api_base():
response = await litellm.acompletion(
model="deepseek-chat",
api_base="https://api.deepseek.com/v1",
api_key="fake-key",
messages=[{"role": "user", "content": "hi"}],
mock_response="resolved",
)
assert response.choices[0].message.content == "resolved"

62
tests/test_litellm/test_router.py
View file

@ -9162,3 +9162,65 @@ def test_model_group_info_reasoning_efforts_ignore_a_mode_the_operator_declared(
assert result is not None
assert result.supported_reasoning_efforts == expected
class TestAddDeploymentApiBaseProviderResolution:
def test_bare_model_with_known_api_base_initializes(self):
router = litellm.Router(
model_list=[
{
"model_name": "groq-pinned",
"litellm_params": {
"model": "llama-3.3-70b-versatile",
"api_base": "https://api.groq.com/openai/v1",
"api_key": "fake-key",
},
},
{
"model_name": "deepseek-pinned",
"litellm_params": {
"model": "deepseek-chat",
"api_base": "https://api.deepseek.com/v1",
"api_key": "fake-key",
},
},
]
)
model_list = router.get_model_list()
assert model_list is not None
assert {m["model_name"] for m in model_list} == {"groq-pinned", "deepseek-pinned"}
def test_bare_model_with_unknown_api_base_still_raises(self):
with pytest.raises(litellm.BadRequestError, match="LLM Provider NOT provided"):
litellm.Router(
model_list=[
{
"model_name": "mystery",
"litellm_params": {
"model": "some-unknown-model",
"api_base": "https://llm.internal.example.com/v1",
"api_key": "fake-key",
},
}
]
)
def test_explicit_custom_llm_provider_beats_api_base_endpoint_match(self):
router = litellm.Router(
model_list=[
{
"model_name": "openai-via-gateway",
"litellm_params": {
"model": "gpt-3.5-turbo",
"custom_llm_provider": "openai",
"api_base": "https://api.groq.com/openai/v1",
"api_key": "fake-key",
},
}
]
)
deployment = router.get_deployment_by_model_group_name("openai-via-gateway")
assert deployment is not None
assert deployment.litellm_params.custom_llm_provider == "openai"

4
type-discipline-budget.json
View file

@ -1,9 +1,9 @@
{
"LIT001": {
"limit": 22805
"limit": 22804
},
"LIT002": {
"limit": 26873
"limit": 26872
},
"LIT003": {
"limit": 269

38
ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_edit.test.tsx
View file

@ -381,6 +381,44 @@ describe("MCPServerEdit (true passthrough warning)", () => {
});
});
describe("MCPServerEdit (OAuth authorize temp payload)", () => {
beforeEach(() => {
vi.clearAllMocks();
});
it("forwards issuer/authorization_url/token_url/registration_url to the temp OAuth session payload", async () => {
// Without these fields the ephemeral server the temp OAuth session endpoint builds has no
// admin-configured OAuth endpoints on it, discovery falls back to (and fails against) the
// plain server url, and Authorize & Fetch Token 400s with "authorization url is not
// configured" even though the saved server (and the visible form) has all four fields filled in.
render(
<MCPServerEdit
mcpServer={{
...interactiveOAuthServer,
issuer: "https://github.com/login/oauth",
authorization_url: "https://github.com/login/oauth/authorize",
token_url: "https://github.com/login/oauth/access_token",
registration_url: "https://github.com/login/oauth/register",
}}
accessToken="access-token"
onCancel={vi.fn()}
onSuccess={vi.fn()}
availableAccessGroups={[]}
/>,
);
await waitFor(() => {
expect(mockOauth.getTemporaryPayload).toBeTruthy();
});
const payload = mockOauth.getTemporaryPayload!();
expect(payload).toBeTruthy();
expect(payload?.issuer).toBe("https://github.com/login/oauth");
expect(payload?.authorization_url).toBe("https://github.com/login/oauth/authorize");
expect(payload?.token_url).toBe("https://github.com/login/oauth/access_token");
expect(payload?.registration_url).toBe("https://github.com/login/oauth/register");
});
});
describe("MCPServerEdit (auth type switch)", () => {
beforeEach(() => {
vi.clearAllMocks();

4
ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_edit.tsx
View file

@ -282,6 +282,10 @@ const MCPServerEdit: React.FC<MCPServerEditProps> = ({
credentials: isClientForwardedTokenMode(values.auth_type)
? preservedAdminCredentials(values.credentials)
: values.credentials,
issuer: values.issuer,
authorization_url: values.authorization_url,
token_url: values.token_url,
registration_url: values.registration_url,
mcp_access_groups: values.mcp_access_groups || mcpServer.mcp_access_groups,
static_headers: staticHeaders,
command: values.command,

63
ui/litellm-dashboard/src/lib/http/schema.d.ts generated vendored
View file

@ -14936,6 +14936,33 @@ export interface paths {
patch?: never;
trace?: never;
};
"/team/{team_id}/member/{user_id}/reset_spend": {
parameters: {
query?: never;
header?: never;
path?: never;
cookie?: never;
};
get?: never;
put?: never;
/**
* Reset Team Member Spend Fn
* @description Reset a team member's tracked spend against their per-member budget.
*
* A member's spend is tracked separately from both their own personal
* budget and the team's own budget (LiteLLM_TeamMembership.spend), so
* neither /user/update nor /team/update can clear it: this is the only
* endpoint that does. The cross-pod spend counter and cached membership
* reads are invalidated so the reset takes effect on the member's next
* request rather than waiting on the membership cache's TTL.
*/
post: operations["reset_team_member_spend_fn_team__team_id__member__user_id__reset_spend_post"];
delete?: never;
options?: never;
head?: never;
patch?: never;
trace?: never;
};
"/team/{team_id}/members/me": {
parameters: {
query?: never;
@ -55085,6 +55112,42 @@ export interface operations {
};
};
};
reset_team_member_spend_fn_team__team_id__member__user_id__reset_spend_post: {
parameters: {
query?: never;
header?: never;
path: {
team_id: string;
user_id: string;
};
cookie?: never;
};
requestBody: {
content: {
"application/json": components["schemas"]["ResetSpendRequest"];
};
};
responses: {
/** @description Successful Response */
200: {
headers: {
[name: string]: unknown;
};
content: {
"application/json": unknown;
};
};
/** @description Validation Error */
422: {
headers: {
[name: string]: unknown;
};
content: {
"application/json": components["schemas"]["HTTPValidationError"];
};
};
};
};
team_member_me_team__team_id__members_me_get: {
parameters: {
query?: never;