scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-15 09:12:18 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 311

Merge remote-tracking branch 'upstream/litellm_internal_staging' into litellm_fix-ssrf-path-traversal

Browse source
This commit is contained in:
Ishaan Jaffer 2026-05-01 10:32:58 -07:00
commit e92d05aea3
No known key found for this signature in database
102 changed files with 6756 additions and 1554 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
.gitignore vendored
View file

@ -90,7 +90,6 @@ test.py
litellm_config.yaml
!.github/observatory/litellm_config.yaml
.cursor
.vscode/launch.json
litellm/proxy/to_delete_loadtest_work/*
update_model_cost_map.py
tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py
@ -100,4 +99,5 @@ STABILIZATION_TODO.md
**/test-results
**/playwright-report
**/*.storageState.json
**/coverage
**/coverage
test-config

4
litellm-proxy-extras/pyproject.toml
View file

@ -1,6 +1,6 @@
[project]
name = "litellm-proxy-extras"
version = "0.4.69"
version = "0.4.70"
description = "Additional files for the LiteLLM Proxy. Reduces the size of the main litellm package."
readme = "README.md"
requires-python = ">=3.9"
@ -26,7 +26,7 @@ required-version = ">=0.10.9"
module-root = ""
[tool.commitizen]
version = "0.4.69"
version = "0.4.70"
version_files = [
"pyproject.toml:^version",
"../pyproject.toml:litellm-proxy-extras==",

3
litellm/caching/caching.py
View file

@ -432,9 +432,10 @@ class Cache:
str: The final hashed cache key with the redis namespace.
"""
dynamic_cache_control: DynamicCacheControl = kwargs.get("cache", {})
metadata = kwargs.get("metadata") or {}
namespace = (
dynamic_cache_control.get("namespace")
or kwargs.get("metadata", {}).get("redis_namespace")
or metadata.get("redis_namespace")
or self.namespace
)
if namespace:

99
litellm/caching/caching_handler.py
View file

@ -87,6 +87,18 @@ class CachingHandlerResponse(BaseModel):
in_memory_cache_obj = InMemoryCache()
def _should_defer_streaming_cache_hit_callbacks(*, kwargs: Dict[str, Any]) -> bool:
"""
When stream=True, do not run success callbacks at cache-hit time.
Cached chat/text completion replay uses CustomStreamWrapper; cached Responses
replay uses CachedResponsesAPIStreamingIterator. Both invoke logging success
handlers when the stream finishes; firing them here too would double-count
spend and callback records.
"""
return kwargs.get("stream", False) is True
class LLMCachingHandler:
def __init__(
self,
@ -99,6 +111,7 @@ class LLMCachingHandler:
self.async_streaming_chunks: List[ModelResponse] = []
self.sync_streaming_chunks: List[ModelResponse] = []
self.request_kwargs = request_kwargs
self.preset_cache_key: Optional[str] = None
self.original_function = original_function
self.start_time = start_time
if litellm.cache is not None and isinstance(litellm.cache.cache, RedisCache):
@ -206,7 +219,7 @@ class LLMCachingHandler:
custom_llm_provider=kwargs.get("custom_llm_provider", None),
args=args,
)
if kwargs.get("stream", False) is False:
if not _should_defer_streaming_cache_hit_callbacks(kwargs=kwargs):
# LOG SUCCESS
self._async_log_cache_hit_on_callbacks(
logging_obj=logging_obj,
@ -215,11 +228,12 @@ class LLMCachingHandler:
end_time=end_time,
cache_hit=cache_hit,
)
cache_key = litellm.cache.get_cache_key(**kwargs)
if (
isinstance(cached_result, BaseModel)
or isinstance(cached_result, CustomStreamWrapper)
) and hasattr(cached_result, "_hidden_params"):
cache_key = (
self.preset_cache_key
or self.request_kwargs.get("cache_key")
or litellm.cache.get_cache_key(**self.request_kwargs)
)
if hasattr(cached_result, "_hidden_params"):
cached_result._hidden_params["cache_key"] = cache_key # type: ignore
return CachingHandlerResponse(cached_result=cached_result)
elif (
@ -265,8 +279,6 @@ class LLMCachingHandler:
kwargs: Dict[str, Any],
args: Optional[Tuple[Any, ...]] = None,
) -> CachingHandlerResponse:
from litellm.utils import CustomStreamWrapper
cached_result: Optional[Any] = None
# Check if caching should be performed BEFORE doing expensive kwargs copy
@ -282,6 +294,11 @@ class LLMCachingHandler:
args,
)
)
if new_kwargs.get("metadata") is None:
new_kwargs.pop("metadata", None)
if new_kwargs.get("stream") is True and "cache_key" not in new_kwargs:
new_kwargs["cache_key"] = litellm.cache.get_cache_key(**new_kwargs)
self.request_kwargs = new_kwargs
print_verbose("Checking Sync Cache")
cached_result = litellm.cache.get_cache(**new_kwargs)
if cached_result is not None:
@ -322,17 +339,19 @@ class LLMCachingHandler:
is_async=False,
)
logging_obj.handle_sync_success_callbacks_for_async_calls(
result=cached_result,
start_time=start_time,
end_time=end_time,
cache_hit=cache_hit,
if not _should_defer_streaming_cache_hit_callbacks(kwargs=kwargs):
logging_obj.handle_sync_success_callbacks_for_async_calls(
result=cached_result,
start_time=start_time,
end_time=end_time,
cache_hit=cache_hit,
)
cache_key = (
self.preset_cache_key
or self.request_kwargs.get("cache_key")
or litellm.cache.get_cache_key(**self.request_kwargs)
)
cache_key = litellm.cache.get_cache_key(**kwargs)
if (
isinstance(cached_result, BaseModel)
or isinstance(cached_result, CustomStreamWrapper)
) and hasattr(cached_result, "_hidden_params"):
if hasattr(cached_result, "_hidden_params"):
cached_result._hidden_params["cache_key"] = cache_key # type: ignore
return CachingHandlerResponse(cached_result=cached_result)
return CachingHandlerResponse(cached_result=cached_result)
@ -686,6 +705,11 @@ class LLMCachingHandler:
args,
)
)
if new_kwargs.get("metadata") is None:
new_kwargs.pop("metadata", None)
if new_kwargs.get("stream") is True and "cache_key" not in new_kwargs:
new_kwargs["cache_key"] = litellm.cache.get_cache_key(**new_kwargs)
self.request_kwargs = new_kwargs
cached_result: Optional[Any] = None
if call_type == CallTypes.aembedding.value:
if isinstance(new_kwargs["input"], str):
@ -710,14 +734,26 @@ class LLMCachingHandler:
if all(result is None for result in cached_result):
cached_result = None
else:
request_kwargs = new_kwargs.copy()
request_cache_key = request_kwargs.pop("cache_key", None)
if litellm.cache._supports_async() is True:
## check if dual cache is supported ##
self.preset_cache_key = (
request_cache_key or litellm.cache.get_cache_key(**request_kwargs)
)
cached_result = await litellm.cache.async_get_cache(
dynamic_cache_object=self.dual_cache, **new_kwargs
dynamic_cache_object=self.dual_cache,
cache_key=self.preset_cache_key,
**request_kwargs,
)
else: # fallback for caches that don't support async
self.preset_cache_key = (
request_cache_key or litellm.cache.get_cache_key(**request_kwargs)
)
cached_result = litellm.cache.get_cache(
dynamic_cache_object=self.dual_cache, **new_kwargs
dynamic_cache_object=self.dual_cache,
cache_key=self.preset_cache_key,
**request_kwargs,
)
return cached_result
@ -825,8 +861,27 @@ class LLMCachingHandler:
elif (call_type == "aresponses" or call_type == "responses") and isinstance(
cached_result, dict
):
# Convert cached dict back to ResponsesAPIResponse object
cached_result = ResponsesAPIResponse(**cached_result)
from litellm.responses.streaming_iterator import (
CachedResponsesAPIStreamingIterator,
)
response_obj = ResponsesAPIResponse(**cached_result)
if (
hasattr(response_obj, "_hidden_params")
and response_obj._hidden_params is not None
and isinstance(response_obj._hidden_params, dict)
):
response_obj._hidden_params["cache_hit"] = True
if kwargs.get("stream", False) is True:
cached_result = CachedResponsesAPIStreamingIterator(
response=response_obj,
logging_obj=logging_obj,
request_data=kwargs,
call_type=call_type,
)
else:
cached_result = response_obj
if (
hasattr(cached_result, "_hidden_params")

24
litellm/caching/dual_cache.py
View file

@ -92,6 +92,25 @@ class DualCache(BaseCache):
if default_redis_ttl is not None:
self.default_redis_ttl = default_redis_ttl
def attach_redis_cache(
self,
redis_cache: Optional[RedisCache] = None,
*,
default_redis_ttl: Optional[float] = None,
) -> None:
"""
Attach a Redis backend if this DualCache does not already have one.
No-op when ``redis_cache`` is None or when Redis was already set (constructor
or a prior attach). Use this for lazy wiring after a shared Redis client exists.
Does not backfill in-memory-only keys to Redis.
"""
if redis_cache is None or self.redis_cache is not None:
return
self.redis_cache = redis_cache
if default_redis_ttl is not None:
self.default_redis_ttl = default_redis_ttl
def set_cache(self, key, value, local_only: bool = False, **kwargs):
# Update both Redis and in-memory cache
try:
@ -392,6 +411,7 @@ class DualCache(BaseCache):
value: float,
parent_otel_span: Optional[Span] = None,
local_only: bool = False,
refresh_ttl: bool = False,
**kwargs,
) -> Optional[float]:
"""
@ -399,6 +419,9 @@ class DualCache(BaseCache):
Value - float - the value you want to increment by
Refresh_ttl - bool - if True, resets the Redis TTL on every write.
Default False preserves window-style semantics.
Returns - the incremented value, or None if no cache backend is
available (in_memory_cache is None and Redis failed/is absent).
"""
@ -415,6 +438,7 @@ class DualCache(BaseCache):
value,
parent_otel_span=parent_otel_span,
ttl=kwargs.get("ttl", None),
refresh_ttl=refresh_ttl,
)
return result

20
litellm/caching/redis_cache.py
View file

@ -551,6 +551,13 @@ class RedisCache(BaseCache):
async def async_set_cache(self, key, value, **kwargs):
from redis.asyncio import Redis
if key is None:
verbose_logger.debug(
"LiteLLM Redis Caching: async set() skipped — key is None, value=%r",
value,
)
return None
start_time = time.time()
try:
_redis_client: Redis = self.init_async_client() # type: ignore
@ -569,8 +576,9 @@ class RedisCache(BaseCache):
)
)
verbose_logger.error(
"LiteLLM Redis Caching: async set() - Got exception from REDIS %s, Writing value=%s",
"LiteLLM Redis Caching: async set() - Got exception from REDIS %s, key=%r, value=%r",
str(e),
key,
value,
)
raise e
@ -824,6 +832,7 @@ class RedisCache(BaseCache):
value: float,
ttl: Optional[int] = None,
parent_otel_span: Optional[Span] = None,
refresh_ttl: bool = False,
) -> float:
from redis.asyncio import Redis
@ -834,11 +843,12 @@ class RedisCache(BaseCache):
try:
result = await _redis_client.incrbyfloat(name=key, amount=value)
if _used_ttl is not None:
# check if key already has ttl, if not -> set ttl
current_ttl = await _redis_client.ttl(key)
if current_ttl == -1:
# Key has no expiration
if refresh_ttl:
await _redis_client.expire(key, _used_ttl)
else:
current_ttl = await _redis_client.ttl(key)
if current_ttl == -1:
await _redis_client.expire(key, _used_ttl)
## LOGGING ##
end_time = time.time()

1
litellm/integrations/prometheus.py
View file

@ -265,6 +265,7 @@ class PrometheusLogger(CustomLogger):
########################################
# LiteLLM Virtual API KEY metrics
########################################
# Remaining MODEL RPM limit for API Key
self.litellm_remaining_api_key_requests_for_model = self._gauge_factory(
"litellm_remaining_api_key_requests_for_model",

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

@ -4582,6 +4582,11 @@ class BedrockConverseMessagesProcessor:
message=cast(ChatCompletionFileObject, element)
)
_parts.append(_part)
elif element["type"] == "document":
_part = BedrockConverseMessagesProcessor._process_document_message(
element
)
_parts.append(_part)
_cache_point_block = (
litellm.AmazonConverseConfig()._get_cache_point_block(
message_block=cast(
@ -4864,6 +4869,44 @@ class BedrockConverseMessagesProcessor:
image_url=cast(str, file_id or file_data), format=format
)
@staticmethod
def _process_document_message(element: dict) -> BedrockContentBlock:
"""Convert a document content block to a Bedrock DocumentBlock.
Handles the Anthropic-style document format:
{"type": "document", "source": {"type": "base64", "media_type": "application/pdf", "data": "..."}}
"""
source = element["source"]
source_type = source.get("type")
if source_type != "base64":
raise ValueError(
f"Bedrock Converse only supports base64-encoded document sources, got '{source_type}'. "
"Please convert the document to base64 before sending to Bedrock."
)
media_type: str = source["media_type"]
data: str = source["data"]
doc_format = BedrockImageProcessor._validate_format(
mime_type=media_type, image_format=media_type.split("/")[1]
)
# Deterministic name using the same hashing pattern as _create_bedrock_block
HASH_SAMPLE_BYTES = 64 * 1024
normalized = "".join(data.split()).encode("utf-8")
sample = normalized[:HASH_SAMPLE_BYTES]
hasher = hashlib.sha256()
hasher.update(sample)
hasher.update(str(len(normalized)).encode("utf-8"))
content_hash = hasher.hexdigest()[:16]
document_name = f"Document_{content_hash}_{doc_format}"
return BedrockContentBlock(
document=BedrockDocumentBlock(
source=BedrockSourceBlock(bytes=data),
format=doc_format,
name=document_name,
)
)
@staticmethod
def add_thinking_blocks_to_assistant_content(
thinking_blocks: List[BedrockContentBlock],
@ -4961,6 +5004,11 @@ def _bedrock_converse_messages_pt( # noqa: PLR0915
)
)
_parts.append(_part)
elif element["type"] == "document":
_part = BedrockConverseMessagesProcessor._process_document_message(
element
)
_parts.append(_part)
_cache_point_block = (
litellm.AmazonConverseConfig()._get_cache_point_block(
message_block=cast(

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

@ -1553,25 +1553,43 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
)
data["output_config"] = output_config
def _transform_response_for_json_mode(
def _resolve_json_mode_non_streaming(
self,
json_mode: Optional[bool],
tool_calls: List[ChatCompletionToolCallChunk],
) -> Optional[LitellmMessage]:
_message: Optional[LitellmMessage] = None
if json_mode is True and len(tool_calls) == 1:
# check if tool name is the default tool name
json_mode_content_str: Optional[str] = None
if (
"name" in tool_calls[0]["function"]
and tool_calls[0]["function"]["name"] == RESPONSE_FORMAT_TOOL_NAME
):
json_mode_content_str = tool_calls[0]["function"].get("arguments")
if json_mode_content_str is not None:
_message = AnthropicConfig._convert_tool_response_to_message(
tool_calls=tool_calls,
)
return _message
) -> Tuple[
Optional[LitellmMessage],
List[ChatCompletionToolCallChunk],
Optional[str],
]:
"""Strip internal response_format tool calls; merge payload into content when mixed with user tools."""
if json_mode is not True or not tool_calls:
return None, tool_calls, None
json_indices = [
i
for i, t in enumerate(tool_calls)
if t.get("function", {}).get("name") == RESPONSE_FORMAT_TOOL_NAME
]
if not json_indices:
return None, tool_calls, None
if len(json_indices) == len(tool_calls):
json_tool = tool_calls[json_indices[0]]
if json_tool.get("function", {}).get("arguments") is None:
return None, tool_calls, None
_message = AnthropicConfig._convert_tool_response_to_message(
tool_calls=[json_tool]
)
return _message, [], None
first_json = tool_calls[json_indices[0]]
json_msg = AnthropicConfig._convert_tool_response_to_message([first_json])
extra_content: Optional[str] = (
json_msg.content if json_msg is not None else None
)
filtered_tools = [t for i, t in enumerate(tool_calls) if i not in json_indices]
return None, filtered_tools, extra_content
def extract_response_content(self, completion_response: dict) -> Tuple[
str,
@ -1931,19 +1949,27 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
tool_calls,
)
json_mode_message, tool_calls_for_message, json_extra_content = (
self._resolve_json_mode_non_streaming(
json_mode=json_mode,
tool_calls=tool_calls,
)
)
merged_text = text_content or ""
if json_extra_content:
merged_text = (
merged_text + json_extra_content if merged_text else json_extra_content
)
_message = litellm.Message(
tool_calls=tool_calls,
content=text_content or None,
tool_calls=tool_calls_for_message,
content=merged_text or None,
provider_specific_fields=provider_specific_fields,
thinking_blocks=thinking_blocks,
reasoning_content=reasoning_content,
)
_message.provider_specific_fields = provider_specific_fields
json_mode_message = self._transform_response_for_json_mode(
json_mode=json_mode,
tool_calls=tool_calls,
)
if json_mode_message is not None:
completion_response["stop_reason"] = "stop"
_message = json_mode_message

1
litellm/llms/base_llm/rerank/transformation.py
View file

@ -33,6 +33,7 @@ class BaseRerankConfig(ABC):
model: str,
optional_rerank_params: Dict,
headers: dict,
litellm_params: Optional[dict] = None,
) -> dict:
return {}

1
litellm/llms/cohere/rerank/transformation.py
View file

@ -111,6 +111,7 @@ class CohereRerankConfig(BaseRerankConfig):
model: str,
optional_rerank_params: Dict,
headers: dict,
litellm_params: Optional[dict] = None,
) -> dict:
if "query" not in optional_rerank_params:
raise ValueError("query is required for Cohere rerank")

1
litellm/llms/cohere/rerank_v2/transformation.py
View file

@ -71,6 +71,7 @@ class CohereRerankV2Config(CohereRerankConfig):
model: str,
optional_rerank_params: Dict,
headers: dict,
litellm_params: Optional[dict] = None,
) -> dict:
if "query" not in optional_rerank_params:
raise ValueError("query is required for Cohere rerank")

26
litellm/llms/custom_httpx/llm_http_handler.py
View file

@ -1007,6 +1007,7 @@ class BaseLLMHTTPHandler:
api_key: Optional[str] = None,
api_base: Optional[str] = None,
client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None,
litellm_params: Optional[Dict[str, Any]] = None,
) -> RerankResponse:
# get config from model, custom llm provider
headers = provider_config.validate_environment(
@ -1026,6 +1027,7 @@ class BaseLLMHTTPHandler:
model=model,
optional_rerank_params=optional_rerank_params,
headers=headers,
litellm_params=litellm_params,
)
## LOGGING
@ -2535,10 +2537,16 @@ class BaseLLMHTTPHandler:
},
)
delete_kwargs: Dict[str, Any] = {
"url": url,
"headers": headers,
"timeout": timeout,
}
if data:
delete_kwargs["json"] = data
try:
response = await async_httpx_client.delete(
url=url, headers=headers, json=data, timeout=timeout
)
response = await async_httpx_client.delete(**delete_kwargs)
except Exception as e:
raise self._handle_error(
@ -2619,10 +2627,16 @@ class BaseLLMHTTPHandler:
},
)
delete_kwargs: Dict[str, Any] = {
"url": url,
"headers": headers,
"timeout": timeout,
}
if data:
delete_kwargs["json"] = data
try:
response = sync_httpx_client.delete(
url=url, headers=headers, json=data, timeout=timeout
)
response = sync_httpx_client.delete(**delete_kwargs)
except Exception as e:
raise self._handle_error(

1
litellm/llms/deepinfra/rerank/transformation.py
View file

@ -132,6 +132,7 @@ class DeepinfraRerankConfig(BaseRerankConfig):
model: str,
optional_rerank_params: Dict,
headers: dict,
litellm_params: Optional[dict] = None,
) -> dict:
# Convert OptionalRerankParams to dict as expected by parent class
if optional_rerank_params is None:

1
litellm/llms/fireworks_ai/rerank/transformation.py
View file

@ -127,6 +127,7 @@ class FireworksAIRerankConfig(FireworksAIMixin, BaseRerankConfig):
model: str,
optional_rerank_params: Dict,
headers: dict,
litellm_params: Optional[dict] = None,
) -> dict:
"""
Transform request to Fireworks AI rerank format

1
litellm/llms/hosted_vllm/rerank/transformation.py
View file

@ -121,6 +121,7 @@ class HostedVLLMRerankConfig(BaseRerankConfig):
model: str,
optional_rerank_params: Dict,
headers: dict,
litellm_params: Optional[dict] = None,
) -> dict:
if "query" not in optional_rerank_params:
raise ValueError("query is required for Hosted VLLM rerank")

1
litellm/llms/huggingface/rerank/transformation.py
View file

@ -146,6 +146,7 @@ class HuggingFaceRerankConfig(BaseRerankConfig):
model: str,
optional_rerank_params: Union[OptionalRerankParams, dict],
headers: dict,
litellm_params: Optional[dict] = None,
) -> dict:
if "query" not in optional_rerank_params:
raise ValueError("query is required for HuggingFace rerank")

6
litellm/llms/jina_ai/rerank/transformation.py
View file

@ -74,7 +74,11 @@ class JinaAIRerankConfig(BaseRerankConfig):
return cleaned_base
def transform_rerank_request(
self, model: str, optional_rerank_params: Dict, headers: Dict
self,
model: str,
optional_rerank_params: Dict,
headers: Dict,
litellm_params: Optional[dict] = None,
) -> Dict:
return {"model": model, **optional_rerank_params}

2
litellm/llms/nvidia_nim/rerank/ranking_transformation.py
View file

@ -66,6 +66,7 @@ class NvidiaNimRankingConfig(NvidiaNimRerankConfig):
model: str,
optional_rerank_params: Dict,
headers: dict,
litellm_params: Optional[dict] = None,
) -> dict:
"""
Transform request, using clean model name without 'ranking/' prefix.
@ -75,4 +76,5 @@ class NvidiaNimRankingConfig(NvidiaNimRerankConfig):
model=clean_model,
optional_rerank_params=optional_rerank_params,
headers=headers,
litellm_params=litellm_params,
)

1
litellm/llms/nvidia_nim/rerank/transformation.py
View file

@ -177,6 +177,7 @@ class NvidiaNimRerankConfig(BaseRerankConfig):
model: str,
optional_rerank_params: Dict,
headers: dict,
litellm_params: Optional[dict] = None,
) -> dict:
"""
Transform request to Nvidia NIM format.

62
litellm/llms/vertex_ai/common_utils.py
View file

@ -27,6 +27,53 @@ class VertexAIError(BaseLLMException):
super().__init__(message=message, status_code=status_code, headers=headers)
def vertex_request_labels_from_litellm_params(
litellm_params: Optional[dict],
) -> Optional[Dict[str, str]]:
"""
Build Vertex/GCP billing labels from LiteLLM user metadata on ``litellm_params``:
``metadata`` (``completion(..., metadata=...)``) or ``litellm_metadata``,
using ``requester_metadata`` string key-value pairs (same convention as Gemini).
``metadata`` is tried first when both are present.
"""
if not litellm_params:
return None
for key in ("metadata", "litellm_metadata"):
if key not in litellm_params:
continue
metadata = litellm_params[key]
if metadata is None or not isinstance(metadata, dict):
continue
if "requester_metadata" not in metadata:
continue
rm = metadata["requester_metadata"]
if not isinstance(rm, dict):
continue
labels = {k: v for k, v in rm.items() if isinstance(v, str)}
if labels:
return labels
return None
def pop_vertex_request_labels(
optional_params: Optional[dict],
litellm_params: Optional[dict],
) -> Optional[Dict[str, str]]:
"""
Resolve labels from optional ``labels`` (Gemini-style) and/or
``litellm_params["metadata"]`` / ``litellm_params["litellm_metadata"]``
(``requester_metadata``). Pops ``labels`` from optional_params when present.
"""
labels: Optional[Dict[str, str]] = None
if optional_params is not None and "labels" in optional_params:
raw = optional_params.pop("labels")
if isinstance(raw, dict):
labels = {k: v for k, v in raw.items() if isinstance(v, str)}
if not labels:
labels = vertex_request_labels_from_litellm_params(litellm_params)
return labels if labels else None
class VertexAIModelRoute(str, Enum):
"""Enum for Vertex AI model routing"""
@ -50,7 +97,7 @@ def get_vertex_ai_model_route(
Determine which handler to use for a Vertex AI model based on the model name.
Args:
model: The model name (e.g., "llama3-405b", "gemini-pro", "gemma/gemma-3-12b-it", "openai/gpt-oss-120b")
model: The model name (e.g., "llama3-405b", "gemini-pro", "gemma/gemma-3-12b-it", "xai/grok-4.1-fast-non-reasoning")
litellm_params: Optional litellm parameters dict that may contain base_model for routing
Returns:
@ -66,7 +113,7 @@ def get_vertex_ai_model_route(
>>> get_vertex_ai_model_route("gemma/gemma-3-12b-it")
VertexAIModelRoute.GEMMA
>>> get_vertex_ai_model_route("openai/gpt-oss-120b")
>>> get_vertex_ai_model_route("xai/grok-4.1-fast-non-reasoning")
VertexAIModelRoute.MODEL_GARDEN
>>> get_vertex_ai_model_route("1234567890", {"api_base": "http://10.96.32.8"})
@ -102,8 +149,11 @@ def get_vertex_ai_model_route(
if "gemma/" in model:
return VertexAIModelRoute.GEMMA
# Check for model garden openai models
if "openai" in model:
# Check for model garden OpenAI-compatible publisher models.
# Examples:
# - openai/gpt-oss-120b-maas
# - xai/grok-4.1-fast-non-reasoning
if "openai" in model or model.startswith("xai/"):
return VertexAIModelRoute.MODEL_GARDEN
# Check for gemini models
@ -209,8 +259,8 @@ def get_vertex_base_model_name(model: str) -> str:
>>> get_vertex_base_model_name("gemma/gemma-3-12b-it")
"gemma-3-12b-it"
>>> get_vertex_base_model_name("openai/gpt-oss-120b")
"gpt-oss-120b"
>>> get_vertex_base_model_name("xai/grok-4.1-fast-non-reasoning")
"grok-4.1-fast-non-reasoning"
>>> get_vertex_base_model_name("1234567890")
"1234567890"

13
litellm/llms/vertex_ai/gemini/transformation.py
View file

@ -24,6 +24,7 @@ from litellm.litellm_core_utils.prompt_templates.factory import (
response_schema_prompt,
)
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
from litellm.llms.vertex_ai.common_utils import pop_vertex_request_labels
from litellm.types.files import (
get_file_mime_type_for_file_type,
get_file_type_from_extension,
@ -714,16 +715,8 @@ def _transform_request_body( # noqa: PLR0915
optional_params.pop("output_config", None)
config_fields = GenerationConfig.__annotations__.keys()
# If the LiteLLM client sends Gemini-supported parameter "labels", add it
# as "labels" field to the request sent to the Gemini backend.
labels: Optional[dict[str, str]] = optional_params.pop("labels", None)
# If the LiteLLM client sends OpenAI-supported parameter "metadata", add it
# as "labels" field to the request sent to the Gemini backend.
if labels is None and "metadata" in litellm_params:
metadata = litellm_params["metadata"]
if metadata is not None and "requester_metadata" in metadata:
rm = metadata["requester_metadata"]
labels = {k: v for k, v in rm.items() if isinstance(v, str)}
# labels: optional explicit param and/or metadata.requester_metadata (OpenAI metadata)
labels = pop_vertex_request_labels(optional_params, litellm_params)
filtered_params = {
k: v

12
litellm/llms/vertex_ai/image_generation/vertex_imagen_transformation.py
View file

@ -7,7 +7,10 @@ import litellm
from litellm.llms.base_llm.image_generation.transformation import (
BaseImageGenerationConfig,
)
from litellm.llms.vertex_ai.common_utils import get_vertex_base_url
from litellm.llms.vertex_ai.common_utils import (
get_vertex_base_url,
pop_vertex_request_labels,
)
from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexLLM
from litellm.secret_managers.main import get_secret_str
from litellm.types.llms.openai import (
@ -203,13 +206,16 @@ class VertexAIImagenImageGenerationConfig(BaseImageGenerationConfig, VertexLLM):
"sampleCount": 1,
}
# Merge with optional params
labels = pop_vertex_request_labels(optional_params, litellm_params)
# Merge with optional params (after popping labels so they are not sent as Imagen parameters)
parameters = {**default_params, **optional_params}
request_body = {
request_body: dict = {
"instances": [{"prompt": prompt}],
"parameters": parameters,
}
if labels:
request_body["labels"] = labels
return request_body

10
litellm/llms/vertex_ai/rerank/transformation.py
View file

@ -11,12 +11,15 @@ import httpx
import litellm
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.base_llm.rerank.transformation import BaseRerankConfig
from litellm.llms.vertex_ai.common_utils import (
vertex_request_labels_from_litellm_params,
)
from litellm.llms.vertex_ai.vertex_llm_base import VertexBase
from litellm.secret_managers.main import get_secret_str
from litellm.types.rerank import (
RerankBilledUnits,
RerankResponse,
RerankResponseMeta,
RerankBilledUnits,
RerankResponseResult,
)
@ -109,6 +112,7 @@ class VertexAIRerankConfig(BaseRerankConfig, VertexBase):
model: str,
optional_rerank_params: Dict,
headers: dict,
litellm_params: Optional[dict] = None,
) -> dict:
"""
Transform the request from Cohere format to Vertex AI Discovery Engine format
@ -145,6 +149,10 @@ class VertexAIRerankConfig(BaseRerankConfig, VertexBase):
# When return_documents is False, we want to ignore record details (return only IDs)
request_data["ignoreRecordDetailsInResponse"] = not return_documents
user_labels = vertex_request_labels_from_litellm_params(litellm_params)
if user_labels:
request_data["userLabels"] = user_labels
return request_data
def transform_rerank_response(

15
litellm/llms/vertex_ai/vertex_embeddings/embedding_handler.py
View file

@ -1,4 +1,4 @@
from typing import Literal, Optional, Union
from typing import Dict, Literal, Optional, Union
import httpx
@ -44,6 +44,7 @@ class VertexEmbedding(VertexBase):
vertex_credentials: Optional[VERTEX_CREDENTIALS_TYPES] = None,
gemini_api_key: Optional[str] = None,
extra_headers: Optional[dict] = None,
litellm_params: Optional[Dict] = None,
) -> EmbeddingResponse:
if aembedding is True:
return self.async_embedding( # type: ignore
@ -61,6 +62,7 @@ class VertexEmbedding(VertexBase):
vertex_credentials=vertex_credentials,
gemini_api_key=gemini_api_key,
extra_headers=extra_headers,
litellm_params=litellm_params,
)
should_use_v1beta1_features = self.is_using_v1beta1_features(
@ -92,7 +94,10 @@ class VertexEmbedding(VertexBase):
headers = self.set_headers(auth_header=auth_header, extra_headers=extra_headers)
vertex_request: VertexEmbeddingRequest = (
litellm.vertexAITextEmbeddingConfig.transform_openai_request_to_vertex_embedding_request(
input=input, optional_params=optional_params, model=model
input=input,
optional_params=optional_params,
model=model,
litellm_params=litellm_params,
)
)
@ -156,6 +161,7 @@ class VertexEmbedding(VertexBase):
gemini_api_key: Optional[str] = None,
extra_headers: Optional[dict] = None,
encoding=None,
litellm_params: Optional[Dict] = None,
) -> EmbeddingResponse:
"""
Async embedding implementation
@ -188,7 +194,10 @@ class VertexEmbedding(VertexBase):
headers = self.set_headers(auth_header=auth_header, extra_headers=extra_headers)
vertex_request: VertexEmbeddingRequest = (
litellm.vertexAITextEmbeddingConfig.transform_openai_request_to_vertex_embedding_request(
input=input, optional_params=optional_params, model=model
input=input,
optional_params=optional_params,
model=model,
litellm_params=litellm_params,
)
)

27
litellm/llms/vertex_ai/vertex_embeddings/transformation.py
View file

@ -3,6 +3,7 @@ from typing import List, Literal, Optional, Union
from pydantic import BaseModel
from litellm.llms.vertex_ai.common_utils import pop_vertex_request_labels
from litellm.types.utils import EmbeddingResponse, Usage
from .types import *
@ -100,7 +101,11 @@ class VertexAITextEmbeddingConfig(BaseModel):
return optional_params
def transform_openai_request_to_vertex_embedding_request(
self, input: Union[list, str], optional_params: dict, model: str
self,
input: Union[list, str],
optional_params: dict,
model: str,
litellm_params: Optional[dict] = None,
) -> VertexEmbeddingRequest:
"""
Transforms an openai request to a vertex embedding request.
@ -108,16 +113,26 @@ class VertexAITextEmbeddingConfig(BaseModel):
# Import here to avoid circular import issues with litellm.__init__
from litellm.llms.vertex_ai.vertex_embeddings.bge import VertexBGEConfig
labels = pop_vertex_request_labels(optional_params, litellm_params)
if model.isdigit():
return self._transform_openai_request_to_fine_tuned_embedding_request(
input, optional_params, model
vertex_request = (
self._transform_openai_request_to_fine_tuned_embedding_request(
input, optional_params, model
)
)
if labels:
vertex_request["labels"] = labels
return vertex_request
if VertexBGEConfig.is_bge_model(model):
return VertexBGEConfig.transform_request(
vertex_request = VertexBGEConfig.transform_request(
input=input, optional_params=optional_params, model=model
)
if labels:
vertex_request["labels"] = labels
return vertex_request
vertex_request: VertexEmbeddingRequest = VertexEmbeddingRequest()
vertex_request = VertexEmbeddingRequest()
vertex_text_embedding_input_list: List[TextEmbeddingInput] = []
task_type: Optional[TaskType] = optional_params.get("task_type")
title = optional_params.get("title")
@ -133,6 +148,8 @@ class VertexAITextEmbeddingConfig(BaseModel):
vertex_request["instances"] = vertex_text_embedding_input_list
vertex_request["parameters"] = EmbeddingParameters(**optional_params)
if labels:
vertex_request["labels"] = labels
return vertex_request

3
litellm/llms/vertex_ai/vertex_embeddings/types.py
View file

@ -3,7 +3,7 @@ Types for Vertex Embeddings Requests
"""
from enum import Enum
from typing import List, Optional, Union
from typing import Dict, List, Optional, Union
from typing_extensions import TypedDict
@ -56,6 +56,7 @@ class VertexEmbeddingRequest(TypedDict, total=False):
List[TextEmbeddingFineTunedInput],
]
parameters: Optional[Union[EmbeddingParameters, TextEmbeddingFineTunedParameters]]
labels: Optional[Dict[str, str]]
# Example usage:

23
litellm/llms/vertex_ai/vertex_model_garden/main.py
View file

@ -27,6 +27,17 @@ from ..common_utils import VertexAIError, get_vertex_base_model_name
from ..vertex_llm_base import VertexBase
def _vertex_model_garden_model_id_in_json_body(model: str) -> bool:
"""
Vertex catalog / publisher models are addressed as publisher/model (e.g.
xai/grok-4.1-fast-reasoning) on the shared OpenAPI URL, with the id in the JSON body.
Deployed Model Garden endpoints are typically a single segment (often numeric)
and use .../endpoints/{ENDPOINT_ID}/chat/completions with an empty model field.
"""
return "/" in model
def create_vertex_url(
vertex_location: str,
vertex_project: str,
@ -34,8 +45,13 @@ def create_vertex_url(
model: str,
api_base: Optional[str] = None,
) -> str:
"""Return the base url for the vertex garden models"""
"""Return the api base for vertex model garden (without /chat/completions)."""
base_url = get_vertex_base_url(vertex_location)
if _vertex_model_garden_model_id_in_json_body(model):
return (
f"{base_url}/v1/projects/{vertex_project}/locations/{vertex_location}"
"/endpoints/openapi"
)
return f"{base_url}/v1beta1/projects/{vertex_project}/locations/{vertex_location}/endpoints/{model}"
@ -129,7 +145,10 @@ class VertexAIModelGardenModels(VertexBase):
vertex_location=vertex_location or "us-central1",
vertex_api_version="v1beta1",
)
model = ""
# Publisher/catalog models: model id must be sent in the JSON body (OpenAPI route).
# Single-segment endpoint ids: model is encoded in the URL path; body model stays empty.
if not _vertex_model_garden_model_id_in_json_body(model):
model = ""
return openai_like_chat_completions.completion(
model=model,
messages=messages,

6
litellm/llms/voyage/rerank/transformation.py
View file

@ -67,7 +67,11 @@ class VoyageRerankConfig(BaseRerankConfig):
return api_base
def transform_rerank_request(
self, model: str, optional_rerank_params: Dict, headers: Dict
self,
model: str,
optional_rerank_params: Dict,
headers: Dict,
litellm_params: Optional[dict] = None,
) -> Dict:
return {"model": model, **optional_rerank_params}

1
litellm/llms/watsonx/rerank/transformation.py
View file

@ -143,6 +143,7 @@ class IBMWatsonXRerankConfig(IBMWatsonXMixin, BaseRerankConfig):
model: str,
optional_rerank_params: Dict,
headers: dict,
litellm_params: Optional[dict] = None,
) -> dict:
"""
Transform request to IBM watsonx.ai rerank format

1
litellm/llms/xai/chat/transformation.py
View file

@ -43,6 +43,7 @@ class XAIChatConfig(OpenAIGPTConfig):
"logprobs",
"max_tokens",
"n",
"parallel_tool_calls",
"presence_penalty",
"response_format",
"seed",

1
litellm/main.py
View file

@ -5311,6 +5311,7 @@ def embedding( # noqa: PLR0915
api_key=api_key,
api_base=api_base,
client=client,
litellm_params=litellm_params_dict,
)
elif custom_llm_provider == "oobabooga":
response = oobabooga.embedding(

66
litellm/model_prices_and_context_window_backup.json
View file

@ -33337,6 +33337,72 @@
"source": "https://console.cloud.google.com/vertex-ai/publishers/openai/model-garden/gpt-oss-120b-maas",
"supports_reasoning": true
},
"vertex_ai/xai/grok-4.1-fast-non-reasoning": {
"cache_read_input_token_cost": 5e-08,
"input_cost_per_token": 2e-07,
"litellm_provider": "vertex_ai",
"max_input_tokens": 2000000,
"max_output_tokens": 2000000,
"max_tokens": 2000000,
"mode": "chat",
"output_cost_per_token": 5e-07,
"source": "https://docs.x.ai/docs/models (Vertex AI Model Garden)",
"supports_function_calling": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true,
"supports_web_search": true
},
"vertex_ai/xai/grok-4.1-fast-reasoning": {
"cache_read_input_token_cost": 5e-08,
"input_cost_per_token": 2e-07,
"litellm_provider": "vertex_ai",
"max_input_tokens": 2000000,
"max_output_tokens": 2000000,
"max_tokens": 2000000,
"mode": "chat",
"output_cost_per_token": 5e-07,
"source": "https://docs.x.ai/docs/models (Vertex AI Model Garden)",
"supports_function_calling": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true,
"supports_web_search": true
},
"vertex_ai/xai/grok-4.20-non-reasoning": {
"cache_read_input_token_cost": 2e-07,
"input_cost_per_token": 2e-06,
"litellm_provider": "vertex_ai",
"max_input_tokens": 2000000,
"max_output_tokens": 2000000,
"max_tokens": 2000000,
"mode": "chat",
"output_cost_per_token": 6e-06,
"source": "https://docs.x.ai/docs/models (Vertex AI Model Garden)",
"supports_function_calling": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true,
"supports_web_search": true
},
"vertex_ai/xai/grok-4.20-reasoning": {
"cache_read_input_token_cost": 2e-07,
"input_cost_per_token": 2e-06,
"litellm_provider": "vertex_ai",
"max_input_tokens": 2000000,
"max_output_tokens": 2000000,
"max_tokens": 2000000,
"mode": "chat",
"output_cost_per_token": 6e-06,
"source": "https://docs.x.ai/docs/models (Vertex AI Model Garden)",
"supports_function_calling": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true,
"supports_web_search": true
},
"vertex_ai/qwen/qwen3-235b-a22b-instruct-2507-maas": {
"input_cost_per_token": 2.5e-07,
"litellm_provider": "vertex_ai-qwen_models",

69
litellm/proxy/_experimental/mcp_server/mcp_server_manager.py
View file

@ -169,6 +169,37 @@ def _deserialize_json_dict(data: Any) -> Optional[Dict[str, str]]:
class MCPServerManager:
_STDIO_ENV_TEMPLATE_PATTERN = re.compile(r"^\$\{(X-[^}]+)\}$")
@staticmethod
def _resolve_oauth2_flow(
*,
auth_type: Optional[MCPAuthType],
oauth2_flow: Optional[str],
token_url: Optional[str],
authorization_url: Optional[str],
client_id: Optional[str],
client_secret: Optional[str],
) -> Optional[Literal["client_credentials", "authorization_code"]]:
"""Infer oauth2_flow for legacy records that omit the field.
DB rows created before oauth2_flow support may have OAuth2 client
credentials + token_url but a null oauth2_flow. Treat these as M2M,
unless authorization_url is present (interactive OAuth).
"""
if oauth2_flow in ("client_credentials", "authorization_code"):
return cast(
Literal["client_credentials", "authorization_code"], oauth2_flow
)
if oauth2_flow:
# Ignore unknown/untyped values and continue legacy inference.
return None
if auth_type != MCPAuth.oauth2:
return None
if authorization_url:
return None
if token_url and client_id and client_secret:
return "client_credentials"
return None
def __init__(self):
self.registry: Dict[str, MCPServer] = {}
self.config_mcp_servers: Dict[str, MCPServer] = {}
@ -342,7 +373,14 @@ class MCPServerManager:
# oauth specific fields
client_id=server_config.get("client_id", None),
client_secret=server_config.get("client_secret", None),
oauth2_flow=server_config.get("oauth2_flow", None),
oauth2_flow=self._resolve_oauth2_flow(
auth_type=auth_type,
oauth2_flow=server_config.get("oauth2_flow", None),
token_url=resolved_token_url,
authorization_url=resolved_authorization_url,
client_id=server_config.get("client_id", None),
client_secret=server_config.get("client_secret", None),
),
scopes=resolved_scopes,
authorization_url=resolved_authorization_url,
token_url=resolved_token_url,
@ -679,7 +717,17 @@ class MCPServerManager:
client_id=client_id_value or getattr(mcp_server, "client_id", None),
client_secret=client_secret_value
or getattr(mcp_server, "client_secret", None),
oauth2_flow=getattr(mcp_server, "oauth2_flow", None),
oauth2_flow=self._resolve_oauth2_flow(
auth_type=auth_type,
oauth2_flow=getattr(mcp_server, "oauth2_flow", None),
token_url=mcp_server.token_url
or getattr(mcp_oauth_metadata, "token_url", None),
authorization_url=mcp_server.authorization_url
or getattr(mcp_oauth_metadata, "authorization_url", None),
client_id=client_id_value or getattr(mcp_server, "client_id", None),
client_secret=client_secret_value
or getattr(mcp_server, "client_secret", None),
),
scopes=resolved_scopes,
authorization_url=mcp_server.authorization_url
or getattr(mcp_oauth_metadata, "authorization_url", None),
@ -2426,7 +2474,7 @@ class MCPServerManager:
)
)
async def _call_regular_mcp_tool(
async def _call_regular_mcp_tool( # noqa: PLR0915
self,
mcp_server: MCPServer,
original_tool_name: str,
@ -2489,7 +2537,11 @@ class MCPServerManager:
# oauth2 headers
extra_headers: Optional[Dict[str, str]] = None
if mcp_server.auth_type == MCPAuth.oauth2:
extra_headers = oauth2_headers
if mcp_server.has_client_credentials:
# For M2M OAuth servers, Authorization must come from token fetch.
extra_headers = None
else:
extra_headers = oauth2_headers
if mcp_server.extra_headers and raw_headers:
if extra_headers is None:
@ -2501,6 +2553,11 @@ class MCPServerManager:
for header in mcp_server.extra_headers:
if not isinstance(header, str):
continue
if (
mcp_server.has_client_credentials
and header.lower() == "authorization"
):
continue
header_value = normalized_raw_headers.get(header.lower())
if header_value is None:
continue
@ -2536,6 +2593,10 @@ class MCPServerManager:
)
extra_headers.update(hook_extra_headers)
# Reset to None if no headers were actually added
if extra_headers is not None and len(extra_headers) == 0:
extra_headers = None
stdio_env = self._build_stdio_env(mcp_server, raw_headers)
client = await self._create_mcp_client(

48
litellm/proxy/_experimental/mcp_server/server.py
View file

@ -153,6 +153,7 @@ if MCP_AVAILABLE:
MCPAuthenticatedUser,
)
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
MCPServerManager,
global_mcp_server_manager,
)
from litellm.proxy._experimental.mcp_server.openapi_to_mcp_generator import (
@ -900,6 +901,20 @@ if MCP_AVAILABLE:
allowed_mcp_server_id
)
if mcp_server is not None:
# Apply oauth2_flow resolution for legacy DB rows where it may be NULL
resolved_flow = MCPServerManager._resolve_oauth2_flow(
auth_type=mcp_server.auth_type,
oauth2_flow=mcp_server.oauth2_flow,
token_url=mcp_server.token_url,
authorization_url=mcp_server.authorization_url,
client_id=mcp_server.client_id,
client_secret=mcp_server.client_secret,
)
if resolved_flow and resolved_flow != mcp_server.oauth2_flow:
# Create a new instance with the resolved flow for this request
mcp_server = mcp_server.model_copy(
update={"oauth2_flow": resolved_flow}
)
allowed_mcp_servers.append(mcp_server)
if mcp_servers is not None:
@ -1100,8 +1115,13 @@ if MCP_AVAILABLE:
extra_headers: Optional[Dict[str, str]] = None
if server.auth_type == MCPAuth.oauth2:
# Copy to avoid mutating the original dict (important for parallel fetching)
extra_headers = oauth2_headers.copy() if oauth2_headers else None
# For OAuth2 M2M servers, upstream Authorization must come from
# client_credentials token fetch, never from caller headers.
if server.has_client_credentials:
extra_headers = None
else:
# Copy to avoid mutating the original dict (important for parallel fetching)
extra_headers = oauth2_headers.copy() if oauth2_headers else None
if server.extra_headers and raw_headers:
if extra_headers is None:
@ -1114,11 +1134,17 @@ if MCP_AVAILABLE:
for header in server.extra_headers:
if not isinstance(header, str):
continue
if server.has_client_credentials and header.lower() == "authorization":
continue
header_value = normalized_raw_headers.get(header.lower())
if header_value is None:
continue
extra_headers[header] = header_value
# Reset to None if no headers were actually added
if extra_headers is not None and len(extra_headers) == 0:
extra_headers = None
if server_auth_header is None:
server_auth_header = mcp_auth_header
@ -1377,11 +1403,19 @@ if MCP_AVAILABLE:
spend_meta["per_server_tool_counts"] = per_server_tool_counts
end_time = datetime.now()
await litellm_logging_obj.async_success_handler(
result=all_tools,
start_time=list_tools_start_time,
end_time=end_time,
)
try:
await litellm_logging_obj.async_success_handler(
result=all_tools,
start_time=list_tools_start_time,
end_time=end_time,
)
except Exception as log_exc:
# list_tools responses must not be dropped due to non-blocking
# observability/serialization failures.
verbose_logger.warning(
"MCP list_tools success logging failed (continuing): %s",
log_exc,
)
verbose_logger.info(
f"Successfully fetched {len(all_tools)} tools total from all MCP servers"

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

@ -12,14 +12,13 @@ Run checks for:
import asyncio
import re
import time
from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Union, cast
from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Type, Union, cast
from fastapi import HTTPException, Request, status
from pydantic import BaseModel
import litellm
from litellm._logging import verbose_proxy_logger
from litellm.caching.caching import DualCache
from litellm.caching.dual_cache import LimitedSizeOrderedDict
from litellm.constants import (
CLI_JWT_EXPIRATION_HOURS,
@ -66,6 +65,8 @@ from litellm.proxy.guardrails.tool_name_extraction import (
TOOL_CAPABLE_CALL_TYPES,
extract_request_tool_names,
)
from litellm.proxy.common_utils.cache_pydantic_utils import CacheCodec
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.proxy.route_llm_request import route_request
from litellm.proxy.utils import PrismaClient, ProxyLogging, log_db_metrics
from litellm.router import Router
@ -852,7 +853,7 @@ def get_actual_routes(allowed_routes: list) -> list:
async def get_default_end_user_budget(
prisma_client: Optional[PrismaClient],
user_api_key_cache: DualCache,
user_api_key_cache: UserApiKeyCache,
parent_otel_span: Optional[Span] = None,
) -> Optional[LiteLLM_BudgetTable]:
"""
@ -875,9 +876,12 @@ async def get_default_end_user_budget(
cache_key = f"default_end_user_budget:{litellm.max_end_user_budget_id}"
# Check cache first
cached_budget = await user_api_key_cache.async_get_cache(key=cache_key)
cached_budget = await user_api_key_cache.async_get_cache(
key=cache_key,
model_type=LiteLLM_BudgetTable,
)
if cached_budget is not None:
return LiteLLM_BudgetTable(**cached_budget)
return cached_budget
# Fetch from database
try:
@ -891,14 +895,16 @@ async def get_default_end_user_budget(
)
return None
_budget_obj = LiteLLM_BudgetTable(**budget_record.dict())
# Cache the budget for 60 seconds
await user_api_key_cache.async_set_cache(
key=cache_key,
value=budget_record.dict(),
value=_budget_obj,
model_type=LiteLLM_BudgetTable,
ttl=DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL,
)
return LiteLLM_BudgetTable(**budget_record.dict())
return _budget_obj
except Exception as e:
verbose_proxy_logger.error(f"Error fetching default end user budget: {str(e)}")
@ -909,7 +915,7 @@ async def get_default_end_user_budget(
async def get_team_member_default_budget(
budget_id: str,
prisma_client: Optional[PrismaClient],
user_api_key_cache: DualCache,
user_api_key_cache: UserApiKeyCache,
) -> Optional[LiteLLM_BudgetTable]:
"""
Fetches the team-level default per-member budget referenced by team.metadata["team_member_budget_id"].
@ -966,7 +972,7 @@ async def get_team_member_default_budget(
async def _apply_default_budget_to_end_user(
end_user_obj: LiteLLM_EndUserTable,
prisma_client: PrismaClient,
user_api_key_cache: DualCache,
user_api_key_cache: UserApiKeyCache,
parent_otel_span: Optional[Span] = None,
) -> LiteLLM_EndUserTable:
"""
@ -1039,7 +1045,7 @@ def _check_end_user_budget(
async def get_end_user_object(
end_user_id: Optional[str],
prisma_client: Optional[PrismaClient],
user_api_key_cache: DualCache,
user_api_key_cache: UserApiKeyCache,
route: str,
parent_otel_span: Optional[Span] = None,
proxy_logging_obj: Optional[ProxyLogging] = None,
@ -1070,10 +1076,12 @@ async def get_end_user_object(
_key = "end_user_id:{}".format(end_user_id)
# Check cache first
cached_user_obj = await user_api_key_cache.async_get_cache(key=_key)
cached_user_obj = await user_api_key_cache.async_get_cache(
key=_key,
model_type=LiteLLM_EndUserTable,
)
if cached_user_obj is not None:
return_obj = LiteLLM_EndUserTable(**cached_user_obj)
return_obj = cached_user_obj
# Apply default budget if needed
return_obj = await _apply_default_budget_to_end_user(
end_user_obj=return_obj,
@ -1108,9 +1116,11 @@ async def get_end_user_object(
parent_otel_span=parent_otel_span,
)
# Save to cache (always store as dict for consistency)
# Save to cache
await user_api_key_cache.async_set_cache(
key="end_user_id:{}".format(end_user_id), value=_response.dict()
key="end_user_id:{}".format(end_user_id),
value=_response,
model_type=LiteLLM_EndUserTable,
)
# Check budget limits
@ -1128,7 +1138,7 @@ async def get_end_user_object(
async def get_tag_objects_batch(
tag_names: List[str],
prisma_client: Optional[PrismaClient],
user_api_key_cache: DualCache,
user_api_key_cache: UserApiKeyCache,
parent_otel_span: Optional[Span] = None,
proxy_logging_obj: Optional[ProxyLogging] = None,
) -> Dict[str, LiteLLM_TagTable]:
@ -1161,12 +1171,12 @@ async def get_tag_objects_batch(
# Try to get all tags from cache first
for tag_name in tag_names:
cache_key = f"tag:{tag_name}"
cached_tag = await user_api_key_cache.async_get_cache(key=cache_key)
cached_tag = await user_api_key_cache.async_get_cache(
key=cache_key,
model_type=LiteLLM_TagTable,
)
if cached_tag is not None:
if isinstance(cached_tag, dict):
tag_objects[tag_name] = LiteLLM_TagTable(**cached_tag)
else:
tag_objects[tag_name] = cached_tag
tag_objects[tag_name] = cached_tag
else:
uncached_tags.append(tag_name)
@ -1182,11 +1192,13 @@ async def get_tag_objects_batch(
for db_tag in db_tags:
tag_name = db_tag.tag_name
cache_key = f"tag:{tag_name}"
# Cache with default TTL (same as end_user objects)
_tag_obj = LiteLLM_TagTable(**db_tag.dict())
await user_api_key_cache.async_set_cache(
key=cache_key, value=db_tag.dict()
key=cache_key,
value=_tag_obj,
model_type=LiteLLM_TagTable,
)
tag_objects[tag_name] = LiteLLM_TagTable(**db_tag.dict())
tag_objects[tag_name] = _tag_obj
except Exception as e:
verbose_proxy_logger.debug(f"Error batch fetching tags from database: {e}")
@ -1197,7 +1209,7 @@ async def get_tag_objects_batch(
async def get_tag_object(
tag_name: Optional[str],
prisma_client: Optional[PrismaClient],
user_api_key_cache: DualCache,
user_api_key_cache: UserApiKeyCache,
parent_otel_span: Optional[Span] = None,
proxy_logging_obj: Optional[ProxyLogging] = None,
) -> Optional[LiteLLM_TagTable]:
@ -1236,7 +1248,7 @@ async def get_team_membership(
user_id: str,
team_id: str,
prisma_client: Optional[PrismaClient],
user_api_key_cache: DualCache,
user_api_key_cache: UserApiKeyCache,
parent_otel_span: Optional[Span] = None,
proxy_logging_obj: Optional[ProxyLogging] = None,
) -> Optional["LiteLLM_TeamMembership"]:
@ -1256,9 +1268,12 @@ async def get_team_membership(
_key = "team_membership:{}:{}".format(user_id, team_id)
# check if in cache
cached_membership_obj = await user_api_key_cache.async_get_cache(key=_key)
cached_membership_obj = await user_api_key_cache.async_get_cache(
key=_key,
model_type=LiteLLM_TeamMembership,
)
if cached_membership_obj is not None:
return LiteLLM_TeamMembership(**cached_membership_obj)
return cached_membership_obj
# else, check db
try:
@ -1270,10 +1285,12 @@ async def get_team_membership(
if response is None:
return None
# save the team membership object to cache (store as dict)
await user_api_key_cache.async_set_cache(key=_key, value=response.dict())
_response = LiteLLM_TeamMembership(**response.dict())
await user_api_key_cache.async_set_cache(
key=_key,
value=_response,
model_type=LiteLLM_TeamMembership,
)
return _response
except Exception:
@ -1441,7 +1458,7 @@ async def _get_fuzzy_user_object(
async def get_user_object(
user_id: Optional[str],
prisma_client: Optional[PrismaClient],
user_api_key_cache: DualCache,
user_api_key_cache: UserApiKeyCache,
user_id_upsert: bool,
parent_otel_span: Optional[Span] = None,
proxy_logging_obj: Optional[ProxyLogging] = None,
@ -1460,12 +1477,12 @@ async def get_user_object(
# check if in cache
if not check_db_only:
cached_user_obj = await user_api_key_cache.async_get_cache(key=user_id)
cached_user_obj = await user_api_key_cache.async_get_cache(
key=user_id,
model_type=LiteLLM_UserTable,
)
if cached_user_obj is not None:
if isinstance(cached_user_obj, dict):
return LiteLLM_UserTable(**cached_user_obj)
elif isinstance(cached_user_obj, LiteLLM_UserTable):
return cached_user_obj
return cached_user_obj
# else, check db
if prisma_client is None:
raise Exception("No db connected")
@ -1527,7 +1544,8 @@ async def get_user_object(
# save the user object to cache
await user_api_key_cache.async_set_cache(
key=user_id,
value=response_dict,
value=_response,
model_type=LiteLLM_UserTable,
ttl=DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL,
)
@ -1548,13 +1566,21 @@ async def get_user_object(
async def _cache_management_object(
key: str,
value: BaseModel,
user_api_key_cache: DualCache,
value: Union[BaseModel, Dict[str, Any]],
user_api_key_cache: UserApiKeyCache,
proxy_logging_obj: Optional[ProxyLogging],
*,
model_type: Type[BaseModel],
):
"""
Persist management objects via ``UserApiKeyCache`` (in-memory + optional Redis).
``UserApiKeyCache`` serializes with ``model_type`` so Redis and in-memory stay aligned.
"""
await user_api_key_cache.async_set_cache(
key=key,
value=value,
model_type=model_type,
ttl=DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL,
)
@ -1562,7 +1588,7 @@ async def _cache_management_object(
async def _cache_team_object(
team_id: str,
team_table: LiteLLM_TeamTableCachedObj,
user_api_key_cache: DualCache,
user_api_key_cache: UserApiKeyCache,
proxy_logging_obj: Optional[ProxyLogging],
):
key = "team_id:{}".format(team_id)
@ -1575,13 +1601,14 @@ async def _cache_team_object(
value=team_table,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
model_type=LiteLLM_TeamTableCachedObj,
)
async def _cache_key_object(
hashed_token: str,
user_api_key_obj: UserAPIKeyAuth,
user_api_key_cache: DualCache,
user_api_key_cache: UserApiKeyCache,
proxy_logging_obj: Optional[ProxyLogging],
):
key = hashed_token
@ -1594,12 +1621,13 @@ async def _cache_key_object(
value=user_api_key_obj,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
model_type=UserAPIKeyAuth,
)
async def _delete_cache_key_object(
hashed_token: str,
user_api_key_cache: DualCache,
user_api_key_cache: UserApiKeyCache,
proxy_logging_obj: Optional[ProxyLogging],
):
key = hashed_token
@ -1647,7 +1675,7 @@ async def _get_team_object_from_db(team_id: str, prisma_client: PrismaClient):
async def _get_team_object_from_user_api_key_cache(
team_id: str,
prisma_client: PrismaClient,
user_api_key_cache: DualCache,
user_api_key_cache: UserApiKeyCache,
last_db_access_time: LimitedSizeOrderedDict,
db_cache_expiry: int,
proxy_logging_obj: Optional[ProxyLogging],
@ -1708,38 +1736,38 @@ async def _get_team_object_from_user_api_key_cache(
async def _get_team_object_from_cache(
key: str,
proxy_logging_obj: Optional[ProxyLogging],
user_api_key_cache: DualCache,
user_api_key_cache: UserApiKeyCache,
parent_otel_span: Optional[Span],
) -> Optional[LiteLLM_TeamTableCachedObj]:
cached_team_obj: Optional[LiteLLM_TeamTableCachedObj] = None
## CHECK REDIS CACHE ##
## INTERNAL USAGE CACHE (plain DualCache) — checked before UserApiKeyCache stores ##
if (
proxy_logging_obj is not None
and proxy_logging_obj.internal_usage_cache.dual_cache
):
cached_team_obj = (
cached_raw = (
await proxy_logging_obj.internal_usage_cache.dual_cache.async_get_cache(
key=key, parent_otel_span=parent_otel_span
)
)
if cached_raw is not None:
from_internal = CacheCodec.deserialize(
cached_raw, LiteLLM_TeamTableCachedObj
)
if from_internal is not None:
return from_internal
if cached_team_obj is None:
cached_team_obj = await user_api_key_cache.async_get_cache(key=key)
if cached_team_obj is not None:
if isinstance(cached_team_obj, dict):
return LiteLLM_TeamTableCachedObj(**cached_team_obj)
elif isinstance(cached_team_obj, LiteLLM_TeamTableCachedObj):
return cached_team_obj
return None
decoded = await user_api_key_cache.async_get_cache(
key=key,
parent_otel_span=parent_otel_span,
model_type=LiteLLM_TeamTableCachedObj,
)
return decoded
async def get_team_object(
team_id: str,
prisma_client: Optional[PrismaClient],
user_api_key_cache: DualCache,
user_api_key_cache: UserApiKeyCache,
parent_otel_span: Optional[Span] = None,
proxy_logging_obj: Optional[ProxyLogging] = None,
check_cache_only: Optional[bool] = None,
@ -1805,20 +1833,21 @@ async def get_team_object(
async def _cache_access_object(
access_group_id: str,
access_group_table: LiteLLM_AccessGroupTable,
user_api_key_cache: DualCache,
user_api_key_cache: UserApiKeyCache,
proxy_logging_obj: Optional[ProxyLogging] = None,
):
key = "access_group_id:{}".format(access_group_id)
await user_api_key_cache.async_set_cache(
key=key,
value=access_group_table,
model_type=LiteLLM_AccessGroupTable,
ttl=DEFAULT_ACCESS_GROUP_CACHE_TTL,
)
async def _delete_cache_access_object(
access_group_id: str,
user_api_key_cache: DualCache,
user_api_key_cache: UserApiKeyCache,
proxy_logging_obj: Optional[ProxyLogging] = None,
):
key = "access_group_id:{}".format(access_group_id)
@ -1836,7 +1865,7 @@ async def _delete_cache_access_object(
async def get_access_object(
access_group_id: str,
prisma_client: Optional[PrismaClient],
user_api_key_cache: DualCache,
user_api_key_cache: UserApiKeyCache,
proxy_logging_obj: Optional[ProxyLogging] = None,
) -> LiteLLM_AccessGroupTable:
"""
@ -1858,13 +1887,12 @@ async def get_access_object(
key = "access_group_id:{}".format(access_group_id)
# Always check cache first
cached_access_obj = await user_api_key_cache.async_get_cache(key=key)
cached_access_obj = await user_api_key_cache.async_get_cache(
key=key,
model_type=LiteLLM_AccessGroupTable,
)
if cached_access_obj is not None:
if isinstance(cached_access_obj, dict):
return LiteLLM_AccessGroupTable(**cached_access_obj)
elif isinstance(cached_access_obj, LiteLLM_AccessGroupTable):
return cached_access_obj
return cached_access_obj
# Not in cache - fetch from DB
try:
@ -1910,7 +1938,7 @@ async def get_access_object(
async def get_team_object_by_alias(
team_alias: str,
prisma_client: Optional[PrismaClient],
user_api_key_cache: DualCache,
user_api_key_cache: UserApiKeyCache,
parent_otel_span: Optional["Span"] = None,
proxy_logging_obj: Optional[ProxyLogging] = None,
) -> LiteLLM_TeamTableCachedObj:
@ -1992,6 +2020,7 @@ async def get_team_object_by_alias(
await user_api_key_cache.async_set_cache(
key=cache_key,
value=team_obj,
model_type=LiteLLM_TeamTableCachedObj,
ttl=DEFAULT_IN_MEMORY_TTL,
)
# Also cache by team_id for consistency
@ -1999,6 +2028,7 @@ async def get_team_object_by_alias(
await user_api_key_cache.async_set_cache(
key=team_id_cache_key,
value=team_obj,
model_type=LiteLLM_TeamTableCachedObj,
ttl=DEFAULT_IN_MEMORY_TTL,
)
@ -2020,7 +2050,7 @@ async def get_team_object_by_alias(
async def get_org_object_by_alias(
org_alias: str,
prisma_client: Optional[PrismaClient],
user_api_key_cache: DualCache,
user_api_key_cache: UserApiKeyCache,
parent_otel_span: Optional["Span"] = None,
proxy_logging_obj: Optional[ProxyLogging] = None,
) -> Optional[LiteLLM_OrganizationTable]:
@ -2047,12 +2077,12 @@ async def get_org_object_by_alias(
# Check cache first (keyed by alias)
cache_key = "org_alias:{}".format(org_alias)
cached_org_obj = await user_api_key_cache.async_get_cache(key=cache_key)
cached_org_obj = await user_api_key_cache.async_get_cache(
key=cache_key,
model_type=LiteLLM_OrganizationTable,
)
if cached_org_obj is not None:
if isinstance(cached_org_obj, dict):
return LiteLLM_OrganizationTable(**cached_org_obj)
elif isinstance(cached_org_obj, LiteLLM_OrganizationTable):
return cached_org_obj
return cached_org_obj
# Query database by organization_alias
try:
@ -2082,13 +2112,15 @@ async def get_org_object_by_alias(
# Cache the result
await user_api_key_cache.async_set_cache(
key=cache_key,
value=org_obj.model_dump(),
value=org_obj,
model_type=LiteLLM_OrganizationTable,
ttl=DEFAULT_IN_MEMORY_TTL,
)
# Also cache by org_id for consistency
await user_api_key_cache.async_set_cache(
key="org_id:{}".format(org_obj.organization_id),
value=org_obj.model_dump(),
value=org_obj,
model_type=LiteLLM_OrganizationTable,
ttl=DEFAULT_IN_MEMORY_TTL,
)
@ -2291,7 +2323,7 @@ async def get_jwt_key_mapping_object(
async def get_key_object(
hashed_token: str,
prisma_client: Optional[PrismaClient],
user_api_key_cache: DualCache,
user_api_key_cache: UserApiKeyCache,
parent_otel_span: Optional[Span] = None,
proxy_logging_obj: Optional[ProxyLogging] = None,
check_cache_only: Optional[bool] = None,
@ -2309,15 +2341,14 @@ async def get_key_object(
# check if in cache
key = hashed_token
cached_key_obj: Optional[UserAPIKeyAuth] = await user_api_key_cache.async_get_cache(
key=key
# Same flow as before: use cache only when we have a hit we can turn into UserAPIKeyAuth
# (dict from Redis / model_dump, or UserAPIKeyAuth from in-memory). Otherwise fall through to DB.
user_api_key_auth = await user_api_key_cache.async_get_cache(
key=key,
model_type=UserAPIKeyAuth,
)
if cached_key_obj is not None:
if isinstance(cached_key_obj, dict):
return UserAPIKeyAuth(**cached_key_obj)
elif isinstance(cached_key_obj, UserAPIKeyAuth):
return cached_key_obj
if user_api_key_auth is not None:
return user_api_key_auth
if check_cache_only:
raise Exception(
@ -2374,7 +2405,7 @@ async def get_key_object(
async def get_object_permission(
object_permission_id: str,
prisma_client: Optional[PrismaClient],
user_api_key_cache: DualCache,
user_api_key_cache: UserApiKeyCache,
parent_otel_span: Optional[Span] = None,
proxy_logging_obj: Optional[ProxyLogging] = None,
) -> Optional[LiteLLM_ObjectPermissionTable]:
@ -2390,12 +2421,12 @@ async def get_object_permission(
# check if in cache
key = "object_permission_id:{}".format(object_permission_id)
cached_obj_permission = await user_api_key_cache.async_get_cache(key=key)
if cached_obj_permission is not None:
if isinstance(cached_obj_permission, dict):
return LiteLLM_ObjectPermissionTable(**cached_obj_permission)
elif isinstance(cached_obj_permission, LiteLLM_ObjectPermissionTable):
return cached_obj_permission
deserialized_perm = await user_api_key_cache.async_get_cache(
key=key,
model_type=LiteLLM_ObjectPermissionTable,
)
if deserialized_perm is not None:
return deserialized_perm
# else, check db
try:
@ -2406,14 +2437,15 @@ async def get_object_permission(
if response is None:
return None
# save the object permission to cache
_perm_obj = LiteLLM_ObjectPermissionTable(**response.dict())
await user_api_key_cache.async_set_cache(
key=key,
value=response.model_dump(),
value=_perm_obj,
model_type=LiteLLM_ObjectPermissionTable,
ttl=DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL,
)
return LiteLLM_ObjectPermissionTable(**response.dict())
return _perm_obj
except Exception:
return None
@ -2422,7 +2454,7 @@ async def get_object_permission(
async def get_managed_vector_store_rows_by_uuids(
uuids: List[str],
prisma_client: Optional[PrismaClient],
user_api_key_cache: DualCache,
user_api_key_cache: UserApiKeyCache,
parent_otel_span: Optional[Span] = None,
proxy_logging_obj: Optional[ProxyLogging] = None,
) -> List[LiteLLM_ManagedVectorStoresTable]:
@ -2442,14 +2474,12 @@ async def get_managed_vector_store_rows_by_uuids(
for uuid in uuids:
key = "managed_vector_store_id:{}".format(uuid)
cached = await user_api_key_cache.async_get_cache(key=key)
if cached is not None:
if isinstance(cached, dict):
result.append(LiteLLM_ManagedVectorStoresTable(**cached))
elif isinstance(cached, LiteLLM_ManagedVectorStoresTable):
result.append(cached)
else:
cache_misses.append(uuid)
deserialized_vs = await user_api_key_cache.async_get_cache(
key=key,
model_type=LiteLLM_ManagedVectorStoresTable,
)
if deserialized_vs is not None:
result.append(deserialized_vs)
else:
cache_misses.append(uuid)
@ -2475,7 +2505,8 @@ async def get_managed_vector_store_rows_by_uuids(
key = "managed_vector_store_id:{}".format(cached_obj.vector_store_id)
await user_api_key_cache.async_set_cache(
key=key,
value=row_dict,
value=cached_obj,
model_type=LiteLLM_ManagedVectorStoresTable,
ttl=DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL,
)
result.append(cached_obj)
@ -2487,7 +2518,7 @@ async def get_managed_vector_store_rows_by_uuids(
async def get_org_object(
org_id: str,
prisma_client: Optional[PrismaClient],
user_api_key_cache: DualCache,
user_api_key_cache: UserApiKeyCache,
parent_otel_span: Optional[Span] = None,
proxy_logging_obj: Optional[ProxyLogging] = None,
include_budget_table: bool = False,
@ -2518,12 +2549,12 @@ async def get_org_object(
cache_key = "org_id:{}:with_budget".format(org_id)
# check if in cache
cached_org_obj = user_api_key_cache.async_get_cache(key=cache_key)
if cached_org_obj is not None:
if isinstance(cached_org_obj, dict):
return LiteLLM_OrganizationTable(**cached_org_obj)
elif isinstance(cached_org_obj, LiteLLM_OrganizationTable):
return cached_org_obj
deserialized_org = await user_api_key_cache.async_get_cache(
key=cache_key,
model_type=LiteLLM_OrganizationTable,
)
if deserialized_org is not None:
return deserialized_org
# else, check db
try:
query_kwargs: Dict[str, Any] = {"where": {"organization_id": org_id}}
@ -2537,16 +2568,16 @@ async def get_org_object(
if response is None:
raise Exception
_org_obj = LiteLLM_OrganizationTable(**response.model_dump())
# Cache the result
await user_api_key_cache.async_set_cache(
key=cache_key,
value=(
response.model_dump() if hasattr(response, "model_dump") else response
),
value=_org_obj,
model_type=LiteLLM_OrganizationTable,
ttl=DEFAULT_IN_MEMORY_TTL,
)
return response
return _org_obj
except Exception:
raise Exception(
f"Organization doesn't exist in db. Organization={org_id}. Create organization via `/organization/new` call."
@ -2559,7 +2590,7 @@ async def _get_resources_from_access_groups(
"access_model_names", "access_mcp_server_ids", "access_agent_ids"
],
prisma_client: Optional[PrismaClient] = None,
user_api_key_cache: Optional[DualCache] = None,
user_api_key_cache: Optional[UserApiKeyCache] = None,
proxy_logging_obj: Optional[ProxyLogging] = None,
) -> List[str]:
"""
@ -2617,7 +2648,7 @@ async def _get_resources_from_access_groups(
async def _get_models_from_access_groups(
access_group_ids: List[str],
prisma_client: Optional[PrismaClient] = None,
user_api_key_cache: Optional[DualCache] = None,
user_api_key_cache: Optional[UserApiKeyCache] = None,
proxy_logging_obj: Optional[ProxyLogging] = None,
) -> List[str]:
"""
@ -2636,7 +2667,7 @@ async def _get_models_from_access_groups(
async def _get_mcp_server_ids_from_access_groups(
access_group_ids: List[str],
prisma_client: Optional[PrismaClient] = None,
user_api_key_cache: Optional[DualCache] = None,
user_api_key_cache: Optional[UserApiKeyCache] = None,
proxy_logging_obj: Optional[ProxyLogging] = None,
) -> List[str]:
"""
@ -2655,7 +2686,7 @@ async def _get_mcp_server_ids_from_access_groups(
async def _get_agent_ids_from_access_groups(
access_group_ids: List[str],
prisma_client: Optional[PrismaClient] = None,
user_api_key_cache: Optional[DualCache] = None,
user_api_key_cache: Optional[UserApiKeyCache] = None,
proxy_logging_obj: Optional[ProxyLogging] = None,
) -> List[str]:
"""
@ -3379,7 +3410,7 @@ async def _check_team_member_budget(
user_object: Optional[LiteLLM_UserTable],
valid_token: Optional[UserAPIKeyAuth],
prisma_client: Optional[PrismaClient],
user_api_key_cache: DualCache,
user_api_key_cache: UserApiKeyCache,
proxy_logging_obj: ProxyLogging,
):
"""Check if team member is over their max budget within the team."""
@ -3447,7 +3478,7 @@ async def _check_team_member_model_access(
valid_token: UserAPIKeyAuth,
llm_router: Optional[Router],
prisma_client: Optional["PrismaClient"],
user_api_key_cache: DualCache,
user_api_key_cache: UserApiKeyCache,
proxy_logging_obj: ProxyLogging,
) -> None:
"""
@ -3754,7 +3785,7 @@ async def _project_soft_budget_check(
async def get_project_object(
project_id: str,
prisma_client: Optional[PrismaClient],
user_api_key_cache: DualCache,
user_api_key_cache: UserApiKeyCache,
proxy_logging_obj: Optional[ProxyLogging] = None,
) -> Optional[LiteLLM_ProjectTableCachedObj]:
"""
@ -3769,12 +3800,12 @@ async def get_project_object(
# Check cache first
cache_key = "project_id:{}".format(project_id)
cached_obj = await user_api_key_cache.async_get_cache(key=cache_key)
if cached_obj is not None:
if isinstance(cached_obj, dict):
return LiteLLM_ProjectTableCachedObj(**cached_obj)
elif isinstance(cached_obj, LiteLLM_ProjectTableCachedObj):
return cached_obj
deserialized_project = await user_api_key_cache.async_get_cache(
key=cache_key,
model_type=LiteLLM_ProjectTableCachedObj,
)
if deserialized_project is not None:
return deserialized_project
# Fetch from DB
project_row = await prisma_client.db.litellm_projecttable.find_unique(
@ -3793,6 +3824,7 @@ async def get_project_object(
value=project_obj,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
model_type=LiteLLM_ProjectTableCachedObj,
)
return project_obj
@ -3802,7 +3834,7 @@ async def _organization_max_budget_check(
valid_token: Optional[UserAPIKeyAuth],
team_object: Optional[LiteLLM_TeamTable],
prisma_client: Optional[PrismaClient],
user_api_key_cache: DualCache,
user_api_key_cache: UserApiKeyCache,
proxy_logging_obj: ProxyLogging,
):
"""
@ -3896,7 +3928,7 @@ async def _organization_max_budget_check(
async def _tag_max_budget_check(
request_body: dict,
prisma_client: Optional[PrismaClient],
user_api_key_cache: DualCache,
user_api_key_cache: UserApiKeyCache,
proxy_logging_obj: ProxyLogging,
valid_token: Optional[UserAPIKeyAuth],
):

28
litellm/proxy/auth/handle_jwt.py
View file

@ -6,6 +6,8 @@ Currently only supports admin.
JWT token must have 'litellm_proxy_admin' in scope.
"""
from __future__ import annotations
import fnmatch
import hashlib
import os
@ -20,7 +22,6 @@ import jwt
from jwt.api_jwk import PyJWK
from litellm._logging import verbose_proxy_logger
from litellm.caching.caching import DualCache
from litellm.constants import DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL
from litellm.litellm_core_utils.dot_notation_indexing import get_nested_value
from litellm.llms.custom_httpx.httpx_handler import HTTPHandler
@ -46,6 +47,7 @@ from litellm.proxy._types import (
)
from litellm.proxy.auth.auth_checks import can_team_access_model
from litellm.proxy.auth.route_checks import RouteChecks
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.proxy.utils import PrismaClient, ProxyLogging
from .auth_checks import (
@ -73,7 +75,7 @@ class JWTHandler:
"""
prisma_client: Optional[PrismaClient]
user_api_key_cache: DualCache
user_api_key_cache: UserApiKeyCache
# Supported algos: https://pyjwt.readthedocs.io/en/stable/algorithms.html
# "Warning: Make sure not to mix symmetric and asymmetric algorithms that interpret
# the key in different ways (e.g. HS* and RS*)."
@ -99,7 +101,7 @@ class JWTHandler:
def update_environment(
self,
prisma_client: Optional[PrismaClient],
user_api_key_cache: DualCache,
user_api_key_cache: UserApiKeyCache,
litellm_jwtauth: LiteLLM_JWTAuth,
leeway: int = 0,
) -> None:
@ -952,7 +954,7 @@ class JWTAuthManager:
jwt_handler: JWTHandler,
jwt_valid_token: dict,
prisma_client: Optional[PrismaClient],
user_api_key_cache: DualCache,
user_api_key_cache: UserApiKeyCache,
parent_otel_span: Optional[Span],
proxy_logging_obj: ProxyLogging,
) -> Tuple[Optional[str], Optional[LiteLLM_TeamTable]]:
@ -1045,7 +1047,7 @@ class JWTAuthManager:
route: str,
jwt_handler: JWTHandler,
prisma_client: Optional[PrismaClient],
user_api_key_cache: DualCache,
user_api_key_cache: UserApiKeyCache,
parent_otel_span: Optional[Span],
proxy_logging_obj: ProxyLogging,
) -> Tuple[Optional[str], Optional[LiteLLM_TeamTable]]:
@ -1133,7 +1135,7 @@ class JWTAuthManager:
valid_user_email: Optional[bool],
jwt_handler: JWTHandler,
prisma_client: Optional[PrismaClient],
user_api_key_cache: DualCache,
user_api_key_cache: UserApiKeyCache,
parent_otel_span: Optional[Span],
proxy_logging_obj: ProxyLogging,
route: str,
@ -1349,7 +1351,7 @@ class JWTAuthManager:
jwt_valid_token: dict,
user_object: Optional[LiteLLM_UserTable],
prisma_client: Optional[PrismaClient],
user_api_key_cache: Optional[DualCache] = None,
user_api_key_cache: Optional[UserApiKeyCache] = None,
) -> None:
"""
Sync user role and team memberships with JWT claims
@ -1377,7 +1379,8 @@ class JWTAuthManager:
if user_api_key_cache is not None:
await user_api_key_cache.async_set_cache(
key=user_object.user_id,
value=user_object.model_dump(),
value=user_object,
model_type=LiteLLM_UserTable,
ttl=DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL,
)
@ -1400,7 +1403,8 @@ class JWTAuthManager:
if user_api_key_cache is not None:
await user_api_key_cache.async_set_cache(
key=user_object.user_id,
value=user_object.model_dump(),
value=user_object,
model_type=LiteLLM_UserTable,
ttl=DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL,
)
return None
@ -1412,7 +1416,7 @@ class JWTAuthManager:
request_headers: Optional[dict],
jwt_handler: JWTHandler,
prisma_client: Optional[PrismaClient],
user_api_key_cache: DualCache,
user_api_key_cache: UserApiKeyCache,
parent_otel_span: Optional[Span],
proxy_logging_obj: ProxyLogging,
) -> None:
@ -1456,7 +1460,7 @@ class JWTAuthManager:
user_object: Optional[LiteLLM_UserTable],
user_id: Optional[str],
prisma_client: Optional[PrismaClient],
user_api_key_cache: DualCache,
user_api_key_cache: UserApiKeyCache,
parent_otel_span: Optional[Span],
proxy_logging_obj: ProxyLogging,
team_id_upsert: Optional[bool],
@ -1514,7 +1518,7 @@ class JWTAuthManager:
general_settings: dict,
route: str,
prisma_client: Optional[PrismaClient],
user_api_key_cache: DualCache,
user_api_key_cache: UserApiKeyCache,
parent_otel_span: Optional[Span],
proxy_logging_obj: ProxyLogging,
request_headers: Optional[dict] = None,

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

@ -20,7 +20,6 @@ from fastapi.security.api_key import APIKeyHeader
import litellm
from litellm._logging import verbose_logger, verbose_proxy_logger
from litellm._service_logger import ServiceLogging
from litellm.caching import DualCache
from litellm.constants import LITELLM_PROXY_MASTER_KEY_ALIAS
from litellm.litellm_core_utils.dd_tracing import tracer
from litellm.litellm_core_utils.dot_notation_indexing import get_nested_value
@ -60,6 +59,7 @@ from litellm.proxy.auth.oauth2_check import Oauth2Handler
from litellm.proxy.auth.oauth2_proxy_hook import handle_oauth2_proxy_request
from litellm.proxy.auth.route_checks import RouteChecks
from litellm.proxy.common_utils.cache_coordinator import EventDrivenCacheCoordinator
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.proxy.common_utils.http_parsing_utils import (
_read_request_body,
_safe_get_request_headers,
@ -329,7 +329,7 @@ _global_spend_coordinator = EventDrivenCacheCoordinator(log_prefix="[GLOBAL SPEN
async def _fetch_global_spend_with_event_coordination(
cache_key: str,
user_api_key_cache: DualCache,
user_api_key_cache: UserApiKeyCache,
prisma_client: PrismaClient,
) -> Optional[float]:
"""
@ -345,14 +345,14 @@ async def _fetch_global_spend_with_event_coordination(
return await _global_spend_coordinator.get_or_load(
cache_key=cache_key,
cache=user_api_key_cache,
cache=user_api_key_cache, # pyright: ignore[reportArgumentType]
load_fn=_load_global_spend,
)
async def get_global_proxy_spend(
litellm_proxy_admin_name: str,
user_api_key_cache: DualCache,
user_api_key_cache: UserApiKeyCache,
prisma_client: Optional[PrismaClient],
token: str,
proxy_logging_obj: ProxyLogging,
@ -510,7 +510,7 @@ async def _resolve_jwt_to_virtual_key(
jwt_claims: dict,
jwt_handler: JWTHandler,
prisma_client: Optional[PrismaClient],
user_api_key_cache: DualCache,
user_api_key_cache: UserApiKeyCache,
parent_otel_span: Optional[Span],
proxy_logging_obj: ProxyLogging,
) -> Optional[UserAPIKeyAuth]:
@ -1112,9 +1112,7 @@ async def _user_api_key_auth_builder( # noqa: PLR0915
is_master_key_valid = False
## VALIDATE MASTER KEY ##
try:
assert isinstance(master_key, str)
except Exception:
if not isinstance(master_key, str):
raise HTTPException(
status_code=500,
detail={
@ -1184,11 +1182,15 @@ async def _user_api_key_auth_builder( # noqa: PLR0915
if len(api_key) > 8
else "****"
)
assert api_key.startswith(
"sk-"
), "LiteLLM Virtual Key expected. Received={}, expected to start with 'sk-'.".format(
_masked_key
) # prevent token hashes from being used
if not api_key.startswith("sk-"):
raise HTTPException(
status_code=status.HTTP_401_UNAUTHORIZED,
detail=(
"LiteLLM Virtual Key expected. Received={}, expected to start with 'sk-'.".format(
_masked_key
)
),
) # prevent token hashes from being used
else:
verbose_logger.warning(
"litellm.proxy.proxy_server.user_api_key_auth(): Warning - Key is not a string. Got type={}".format(
@ -1296,7 +1298,8 @@ async def _user_api_key_auth_builder( # noqa: PLR0915
_cache_key = f"{valid_token.team_id}_{valid_token.user_id}"
team_member_info = await user_api_key_cache.async_get_cache(
key=_cache_key
key=_cache_key,
model_type=LiteLLM_TeamMembership,
)
if team_member_info is None:
# read from DB
@ -1304,18 +1307,23 @@ async def _user_api_key_auth_builder( # noqa: PLR0915
_team_id = valid_token.team_id
if _user_id is not None and _team_id is not None:
team_member_info = await prisma_client.db.litellm_teammembership.find_first(
_db_member = await prisma_client.db.litellm_teammembership.find_first(
where={
"user_id": _user_id,
"team_id": _team_id,
}, # type: ignore
include={"litellm_budget_table": True},
)
await user_api_key_cache.async_set_cache(
key=_cache_key,
value=team_member_info,
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
@ -1462,9 +1470,13 @@ async def _user_api_key_auth_builder( # noqa: PLR0915
else:
valid_token.team_object_permission = None
await user_api_key_cache.async_set_cache(
key=valid_token.team_id, value=_team_obj
) # save team table in cache - used for tpm/rpm limiting - tpm_rpm_limiter.py
# Only cache when the key is a real team_id (non-team keys must not use key=None).
if valid_token.team_id is not None and _team_obj is not None:
await user_api_key_cache.async_set_cache(
key=valid_token.team_id,
value=_team_obj,
model_type=LiteLLM_TeamTableCachedObj,
) # save team table in cache - used for tpm/rpm limiting - tpm_rpm_limiter.py
# Fetch project object if key belongs to a project
_project_obj = None

10
litellm/proxy/common_request_processing.py
View file

@ -744,6 +744,11 @@ class ProxyBaseLLMRequestProcessing:
"aingest",
"aretrieve_container",
"adelete_container",
"aupload_container_file",
"alist_container_files",
"aretrieve_container_file",
"adelete_container_file",
"aretrieve_container_file_content",
"acreate_skill",
"alist_skills",
"aget_skill",
@ -1001,6 +1006,11 @@ class ProxyBaseLLMRequestProcessing:
"aingest",
"aretrieve_container",
"adelete_container",
"aupload_container_file",
"alist_container_files",
"aretrieve_container_file",
"adelete_container_file",
"aretrieve_container_file_content",
"acreate_skill",
"alist_skills",
"aget_skill",

25
litellm/proxy/common_utils/cache_coordinator.py
View file

@ -20,11 +20,27 @@ T = TypeVar("T")
class AsyncCacheProtocol(Protocol):
"""Protocol for cache backends used by EventDrivenCacheCoordinator."""
"""Protocol for cache backends used by EventDrivenCacheCoordinator.
async def async_get_cache(self, key: str, **kwargs: Any) -> Any: ...
Matches ``DualCache`` / ``UserApiKeyCache`` call shapes (explicit optional params
before ``**kwargs``), not only ``(key, **kwargs)``, so overloads validate.
"""
async def async_set_cache(self, key: str, value: Any, **kwargs: Any) -> Any: ...
async def async_get_cache(
self,
key: str,
parent_otel_span: Any = None,
local_only: bool = False,
**kwargs: Any,
) -> Any: ...
async def async_set_cache(
self,
key: str,
value: Any,
local_only: bool = False,
**kwargs: Any,
) -> Any: ...
class EventDrivenCacheCoordinator:
@ -36,6 +52,9 @@ class EventDrivenCacheCoordinator:
- Other requests: wait for the signal, then read from cache.
Create one instance per resource (e.g. one for global spend, one for feature flags).
Args:
log_prefix: Prefix for debug log messages.
"""
def __init__(self, log_prefix: str = "[CACHE]"):

93
litellm/proxy/common_utils/cache_pydantic_utils.py Normal file
View file

@ -0,0 +1,93 @@
"""
DualCache presents a single API for reads and writes, but the two backends behave
differently: the in-memory layer can store arbitrary Python objects (including live
``BaseModel`` instances), while Redis persists strings and therefore needs JSON-safe
payloads (``json.dumps`` on the Redis side).
Call sites therefore see cache ``value`` / ``cached`` as effectively ``Any``: the same
key may deserialize to a model on one process (memory hit) or to a ``dict`` after a
Redis round-trip. ``CacheCodec`` centralizes encode/decode at that boundary:
``CacheCodec.serialize`` before ``set``, ``CacheCodec.deserialize`` after ``get``
when you need a typed ``BaseModel``.
``dataclasses`` are not supported: only ``dict`` and Pydantic ``BaseModel`` inputs
are encoded; pass a Pydantic model or convert with e.g. ``dataclasses.asdict`` first.
"""
from __future__ import annotations
from typing import Any, Optional, Type, TypeVar
from pydantic import BaseModel, ValidationError
from litellm._logging import verbose_proxy_logger
T = TypeVar("T", bound=BaseModel)
class CacheCodec:
"""
Encode/decode Pydantic models for DualCache (memory vs Redis safe payloads).
Dataclasses are not supported yet (only ``dict`` and ``BaseModel``).
Use ``serialize`` with ``model_type`` when writing so the same schema is used
as on read (``deserialize``). Pass ``model_type`` whenever you know it
(validates ``dict`` payloads and normalizes ``BaseModel`` instances).
"""
@staticmethod
def serialize(value: Any, model_type: Optional[Type[T]] = None) -> Any:
"""
Encode a value for DualCache / Redis (``json.dumps``-safe).
If ``model_type`` is set, the payload is validated with that model, then
``model_dump(mode="json", exclude_none=True)`` — symmetric with ``deserialize``.
If the value is already an instance of ``model_type`` (or a subclass),
``model_validate`` is skipped to avoid an unnecessary Pydantic copy — the
value is dumped directly.
If ``model_type`` is omitted, any ``BaseModel`` is dumped as above; other
values (e.g. plain ``dict``) are returned unchanged.
"""
if model_type is not None:
if isinstance(value, model_type):
# Already the right type: dump directly, skip re-validation copy.
return value.model_dump(mode="json", exclude_none=True)
if isinstance(value, (dict, BaseModel)):
return model_type.model_validate(value).model_dump(
mode="json", exclude_none=True
)
return value
if isinstance(value, BaseModel):
return value.model_dump(mode="json", exclude_none=True)
return value
@staticmethod
def deserialize(cached: Any, model_type: Type[T]) -> Optional[T]:
"""
Decode a cache entry to ``model_type``.
- ``None`` → ``None``
- Already an instance of ``model_type`` (including subclasses) → returned as-is
- ``dict`` → ``model_type.model_validate(...)``; on ``ValidationError``,
logs a warning and returns ``None`` (treat as cache miss; avoids serving
malformed or schema-drifted entries)
- Any other type → ``None`` (caller should treat as cache miss or log)
"""
if cached is None:
return None
if isinstance(cached, model_type):
return cached
if isinstance(cached, dict):
try:
return model_type.model_validate(cached)
except ValidationError as e:
verbose_proxy_logger.warning(
"CacheCodec.deserialize: validation failed for %s (%s)",
model_type.__name__,
e,
)
return None
return None

4
litellm/proxy/common_utils/expired_ui_session_key_cleanup_manager.py
View file

@ -8,7 +8,7 @@ from datetime import datetime, timezone
from typing import Any, Dict, List, Optional
from litellm._logging import verbose_proxy_logger
from litellm.caching import DualCache
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.constants import (
EXPIRED_UI_SESSION_KEY_CLEANUP_JOB_NAME,
LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_BATCH_SIZE,
@ -31,7 +31,7 @@ class ExpiredUISessionKeyCleanupManager:
def __init__(
self,
prisma_client: PrismaClient,
user_api_key_cache: DualCache,
user_api_key_cache: UserApiKeyCache,
pod_lock_manager=None,
):
self.prisma_client = prisma_client

120
litellm/proxy/common_utils/reset_budget_job.py
View file

@ -52,6 +52,37 @@ class ResetBudgetJob:
### RESET MULTI-WINDOW BUDGETS ###
await self.reset_budget_windows()
@staticmethod
async def _invalidate_spend_counter(counter_key: str) -> None:
"""Zero a spend counter so a DB-row reset takes effect immediately.
Call AFTER the DB write commits. Clearing Redis before the DB
commit opens a window where get_current_spend reads 0 from Redis
while the DB still holds the pre-reset value, allowing bypass.
"""
try:
from litellm.proxy.proxy_server import spend_counter_cache
spend_counter_cache.in_memory_cache.set_cache(
key=counter_key, value=0.0, ttl=60
)
if spend_counter_cache.redis_cache is not None:
try:
await spend_counter_cache.redis_cache.async_set_cache(
key=counter_key, value=0.0, ttl=60
)
except Exception as redis_err:
verbose_proxy_logger.warning(
"Failed to reset spend counter %s in Redis: %s. "
"Budget may be over-enforced until counter expires.",
counter_key,
redis_err,
)
except Exception as e:
verbose_proxy_logger.warning(
"Failed to reset spend counter %s: %s", counter_key, e
)
async def reset_budget_for_litellm_team_members(
self, budgets_to_reset: List[LiteLLM_BudgetTableFull]
):
@ -64,46 +95,30 @@ class ResetBudgetJob:
if budget.budget_id is not None
]
# Reset spend counters for affected team members.
# Reset Redis directly so a transient failure doesn't leave stale
# counters that get_current_spend would read as authoritative.
try:
from litellm.proxy.proxy_server import spend_counter_cache
memberships = await self.prisma_client.db.litellm_teammembership.find_many(
where={"budget_id": {"in": budget_ids}}
)
for m in memberships:
counter_key = f"spend:team_member:{m.user_id}:{m.team_id}"
# Always reset in-memory
spend_counter_cache.in_memory_cache.set_cache(
key=counter_key, value=0.0
)
# Explicitly reset Redis with warning on failure
if spend_counter_cache.redis_cache is not None:
try:
await spend_counter_cache.redis_cache.async_set_cache(
key=counter_key, value=0.0
)
except Exception as redis_err:
verbose_proxy_logger.warning(
"Failed to reset team member spend counter in Redis %s: %s. "
"Budget may be over-enforced until counter expires.",
counter_key,
redis_err,
)
except Exception as e:
memberships = []
verbose_proxy_logger.warning(
"Failed to reset team member spend counters: %s", e
"Failed to fetch team memberships for counter invalidation: %s", e
)
return await self.prisma_client.db.litellm_teammembership.update_many(
update_result = await self.prisma_client.db.litellm_teammembership.update_many(
where={"budget_id": {"in": budget_ids}},
data={
"spend": 0,
},
)
for m in memberships:
await self._invalidate_spend_counter(
f"spend:team_member:{m.user_id}:{m.team_id}"
)
return update_result
async def reset_budget_for_keys_linked_to_budgets(
self, budgets_to_reset: List[LiteLLM_BudgetTableFull]
):
@ -126,17 +141,36 @@ class ResetBudgetJob:
if not budget_ids:
return
return await self.prisma_client.db.litellm_verificationtoken.update_many(
where={
"budget_id": {"in": budget_ids},
"budget_duration": None, # only keys without their own reset schedule
"spend": {"gt": 0}, # only reset keys that have accumulated spend
},
data={
"spend": 0,
},
where_clause: dict = {
"budget_id": {"in": budget_ids},
"budget_duration": None, # only keys without their own reset schedule
"spend": {"gt": 0}, # only reset keys that have accumulated spend
}
try:
keys = await self.prisma_client.db.litellm_verificationtoken.find_many(
where=where_clause
)
except Exception as e:
keys = []
verbose_proxy_logger.warning(
"Failed to fetch keys for counter invalidation: %s", e
)
update_result = (
await self.prisma_client.db.litellm_verificationtoken.update_many(
where=where_clause,
data={
"spend": 0,
},
)
)
for k in keys:
await self._invalidate_spend_counter(f"spend:key:{k.token}")
return update_result
async def reset_budget_for_litellm_budget_table(self):
"""
Resets the budget for all LiteLLM End-Users (Customers), and Team Members if their budget has expired
@ -365,6 +399,10 @@ class ResetBudgetJob:
data_list=updated_keys,
table_name="key",
)
for k in updated_keys:
token = getattr(k, "token", None)
if token:
await self._invalidate_spend_counter(f"spend:key:{token}")
end_time = time.time()
if len(failed_keys) > 0: # If any keys failed to reset
@ -450,6 +488,12 @@ class ResetBudgetJob:
data_list=updated_users,
table_name="user",
)
for u in updated_users:
user_id = getattr(u, "user_id", None)
if user_id:
await self._invalidate_spend_counter(
f"spend:user:{user_id}"
)
end_time = time.time()
if len(failed_users) > 0: # If any users failed to reset
@ -541,6 +585,12 @@ class ResetBudgetJob:
data_list=updated_teams,
table_name="team",
)
for t in updated_teams:
team_id = getattr(t, "team_id", None)
if team_id:
await self._invalidate_spend_counter(
f"spend:team:{team_id}"
)
end_time = time.time()
if len(failed_teams) > 0: # If any teams failed to reset

162
litellm/proxy/common_utils/user_api_key_cache.py Normal file
View file

@ -0,0 +1,162 @@
from __future__ import annotations
from typing import Any, Optional, Type, TypeVar, Union, cast, overload
from pydantic import BaseModel
from litellm._logging import verbose_proxy_logger
from litellm.caching.dual_cache import DualCache
from litellm.proxy.common_utils.cache_pydantic_utils import CacheCodec
T = TypeVar("T", bound=BaseModel)
class UserApiKeyCache(DualCache):
"""
DualCache wrapper for UserAPIKeyAuth-like payloads.
Stores a Redis-safe JSON payload in BOTH in-memory and Redis to avoid
"memory returns BaseModel, Redis returns dict" format drift.
When ``model_type`` is provided:
- writes are serialized via ``CacheCodec.serialize(..., model_type=...)``
- reads are deserialized via ``CacheCodec.deserialize(..., model_type)``
and return ``Optional[T]``: the model on success, ``None`` on cache miss
**or** if the cached payload fails validation (schema drift). On
validation failure after a cache hit, an error line is emitted via
``verbose_proxy_logger``.
When ``model_type`` is omitted, the interface behaves like ``DualCache``:
raw cached payload is returned (dict/str/etc.).
``async_set_cache_pipeline`` applies the same untyped Codec pass as omitting
``model_type`` on ``async_set_cache`` (so ``BaseModel`` rows are dumped before Redis).
``get_cache`` / ``async_get_cache`` overloads and implementations must be contiguous
(no other methods in between) so mypy resolves ``@overload`` + implementation correctly.
"""
@overload
def get_cache(
self,
key: Any,
parent_otel_span: Any = None,
local_only: bool = False,
*,
model_type: Type[T],
**kwargs: Any,
) -> Optional[T]: ...
@overload
def get_cache(
self,
key: Any,
parent_otel_span: Any = None,
local_only: bool = False,
**kwargs: Any,
) -> Any: ...
def get_cache( # type: ignore[override]
self,
key,
parent_otel_span=None,
local_only: bool = False,
model_type: Optional[Type[BaseModel]] = None,
**kwargs,
) -> Union[Any, Optional[BaseModel]]:
if model_type is None and "model_type" in kwargs:
model_type = cast(Optional[Type[BaseModel]], kwargs.pop("model_type", None))
cached = super().get_cache(
key=key, parent_otel_span=parent_otel_span, local_only=local_only, **kwargs
)
if model_type is None:
return cached
if cached is None:
return None
decoded = CacheCodec.deserialize(cached, model_type=model_type)
if decoded is None:
verbose_proxy_logger.error(
"UserApiKeyCache.get_cache failed to deserialize cached value for "
"key=%r model_type=%s",
key,
getattr(model_type, "__name__", str(model_type)),
)
return None
return decoded
@overload
async def async_get_cache(
self,
key: Any,
parent_otel_span: Any = None,
local_only: bool = False,
*,
model_type: Type[T],
**kwargs: Any,
) -> Optional[T]: ...
@overload
async def async_get_cache(
self,
key: Any,
parent_otel_span: Any = None,
local_only: bool = False,
**kwargs: Any,
) -> Any: ...
async def async_get_cache( # type: ignore[override]
self,
key,
parent_otel_span=None,
local_only: bool = False,
model_type: Optional[Type[BaseModel]] = None,
**kwargs,
) -> Union[Any, Optional[BaseModel]]:
if model_type is None and "model_type" in kwargs:
model_type = cast(Optional[Type[BaseModel]], kwargs.pop("model_type", None))
cached = await super().async_get_cache(
key=key, parent_otel_span=parent_otel_span, local_only=local_only, **kwargs
)
if model_type is None:
return cached
if cached is None:
return None
decoded = CacheCodec.deserialize(cached, model_type=model_type)
if decoded is None:
verbose_proxy_logger.error(
"UserApiKeyCache.async_get_cache failed to deserialize cached value for "
"key=%r model_type=%s",
key,
getattr(model_type, "__name__", str(model_type)),
)
return None
return decoded
def set_cache(self, key, value, local_only: bool = False, **kwargs): # type: ignore[override]
model_type = cast(Optional[Type[BaseModel]], kwargs.pop("model_type", None))
payload = CacheCodec.serialize(value, model_type=model_type)
return super().set_cache(
key=key, value=payload, local_only=local_only, **kwargs
)
async def async_set_cache(self, key, value, local_only: bool = False, **kwargs): # type: ignore[override]
model_type = cast(Optional[Type[BaseModel]], kwargs.pop("model_type", None))
payload = CacheCodec.serialize(value, model_type=model_type)
return await super().async_set_cache(
key=key, value=payload, local_only=local_only, **kwargs
)
async def async_set_cache_pipeline( # type: ignore[override]
self, cache_list: list, local_only: bool = False, **kwargs
) -> None:
"""
Batch writes with the same Codec boundary as ``async_set_cache`` without
``model_type``: ``BaseModel`` values become JSON-safe dicts; dicts/scalars unchanged.
"""
normalized = [
(key, CacheCodec.serialize(value, model_type=None))
for key, value in cache_list
]
return await super().async_set_cache_pipeline(
cache_list=normalized, local_only=local_only, **kwargs
)

125
litellm/proxy/container_endpoints/handler_factory.py
View file

@ -19,7 +19,6 @@ from litellm.proxy.common_utils.openai_endpoint_utils import (
get_custom_llm_provider_from_request_headers,
get_custom_llm_provider_from_request_query,
)
from litellm.responses.utils import ResponsesAPIRequestUtils
def _load_endpoints_config() -> Dict:
@ -64,10 +63,12 @@ def _create_handler_for_path_params(
request: Request,
container_id: str,
file_id: str,
fastapi_response: Response,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
return await _process_binary_request(
request=request,
fastapi_response=fastapi_response,
container_id=container_id,
file_id=file_id,
user_api_key_dict=user_api_key_dict,
@ -152,63 +153,61 @@ def _create_handler_for_path_params(
async def _process_binary_request(
request: Request,
fastapi_response: Response,
container_id: str,
file_id: str,
user_api_key_dict: UserAPIKeyAuth,
):
"""
Process binary content requests using the proper transformation pattern.
Process binary content requests through the standard proxy/router pipeline.
This uses the provider config transformations and llm_http_handler
to maintain consistency with the established pattern.
The router owns managed container ID decoding and deployment selection. This
handler only adapts the byte response to FastAPI.
"""
from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler
from litellm.types.router import GenericLiteLLMParams
from litellm.proxy.proxy_server import (
general_settings,
llm_router,
proxy_config,
proxy_logging_obj,
select_data_generator,
user_api_base,
user_max_tokens,
user_model,
user_request_timeout,
user_temperature,
version,
)
# Extract custom_llm_provider
custom_llm_provider = (
get_custom_llm_provider_from_request_headers(request=request)
or get_custom_llm_provider_from_request_query(request=request)
or "openai"
)
# Build litellm_params - credentials are resolved by provider config from env
litellm_params = GenericLiteLLMParams()
# Decode container ID and extract provider info
decoded = ResponsesAPIRequestUtils._decode_container_id(container_id)
original_container_id = decoded.get("response_id", container_id)
# If container ID has encoded provider info and user didn't explicitly set provider, use it
decoded_provider = decoded.get("custom_llm_provider")
if decoded_provider and custom_llm_provider == "openai":
custom_llm_provider = decoded_provider
# Get the provider config
container_provider_config = _get_container_provider_config(custom_llm_provider)
# Create logging object
logging_obj = Logging(
model="container-file-content",
messages=[],
stream=False,
call_type="container_file_content",
start_time=None,
litellm_call_id="",
function_id="",
)
# Use the HTTP handler to make the request
handler = BaseLLMHTTPHandler()
data: Dict[str, Any] = {
"container_id": container_id,
"file_id": file_id,
"custom_llm_provider": custom_llm_provider,
}
processor = ProxyBaseLLMRequestProcessing(data=data)
try:
content = await handler.async_container_file_content_handler(
container_id=original_container_id, # Use decoded original ID
file_id=file_id,
container_provider_config=container_provider_config,
litellm_params=litellm_params,
logging_obj=logging_obj,
content = await processor.base_process_llm_request(
request=request,
fastapi_response=fastapi_response,
user_api_key_dict=user_api_key_dict,
route_type="aretrieve_container_file_content",
proxy_logging_obj=proxy_logging_obj,
llm_router=llm_router,
general_settings=general_settings,
proxy_config=proxy_config,
select_data_generator=select_data_generator,
model=None,
user_model=user_model,
user_temperature=user_temperature,
user_request_timeout=user_request_timeout,
user_max_tokens=user_max_tokens,
user_api_base=user_api_base,
version=version,
)
# Determine content type based on common file extensions in the file_id
@ -229,13 +228,25 @@ async def _process_binary_request(
elif ".pdf" in file_id_lower:
content_type = "application/pdf"
if not isinstance(content, bytes):
raise TypeError(
"aretrieve_container_file_content expected bytes, got "
f"{type(content).__name__}"
)
return Response(
content=content,
headers=dict(fastapi_response.headers),
media_type=content_type,
)
except Exception as e:
raise e
raise await processor._handle_llm_api_exception(
e=e,
user_api_key_dict=user_api_key_dict,
proxy_logging_obj=proxy_logging_obj,
version=version,
)
async def _process_multipart_upload_request(
@ -284,16 +295,7 @@ async def _process_multipart_upload_request(
or "openai"
)
# Decode container ID and extract provider info
decoded = ResponsesAPIRequestUtils._decode_container_id(container_id)
original_container_id = decoded.get("response_id", container_id)
# If container ID has encoded provider info and user didn't explicitly set provider, use it
decoded_provider = decoded.get("custom_llm_provider")
if decoded_provider and custom_llm_provider == "openai":
custom_llm_provider = decoded_provider
data["container_id"] = original_container_id # Use decoded original ID
data["container_id"] = container_id
data["custom_llm_provider"] = custom_llm_provider
processor = ProxyBaseLLMRequestProcessing(data=data)
@ -359,21 +361,6 @@ async def _process_request(
or "openai"
)
# Decode container_id if present in path_params
if "container_id" in path_params:
decoded = ResponsesAPIRequestUtils._decode_container_id(
path_params["container_id"]
)
original_container_id = decoded.get("response_id", path_params["container_id"])
# If container ID has encoded provider info and user didn't explicitly set provider, use it
decoded_provider = decoded.get("custom_llm_provider")
if decoded_provider and custom_llm_provider == "openai":
custom_llm_provider = decoded_provider
# Update path_params with decoded original ID
data["container_id"] = original_container_id
data["custom_llm_provider"] = custom_llm_provider
processor = ProxyBaseLLMRequestProcessing(data=data)

14
litellm/proxy/db/spend_counter_reseed.py
View file

@ -129,7 +129,9 @@ class SpendCounterReseed:
"""
lock = await SpendCounterReseed._get_lock(counter_key)
async with lock:
# Re-check after acquiring the lock - another waiter may have warmed it.
# Re-check after acquiring the lock. Skip in-memory on a clean
# Redis miss - in-memory is per-pod-stale.
redis_clean_miss = False
if spend_counter_cache.redis_cache is not None:
try:
val = await spend_counter_cache.redis_cache.async_get_cache(
@ -137,11 +139,13 @@ class SpendCounterReseed:
)
if val is not None:
return float(val)
redis_clean_miss = True
except Exception:
pass
val = spend_counter_cache.in_memory_cache.get_cache(key=counter_key)
if val is not None:
return float(val)
if not redis_clean_miss:
val = spend_counter_cache.in_memory_cache.get_cache(key=counter_key)
if val is not None:
return float(val)
db_spend = await SpendCounterReseed.from_db(prisma_client, counter_key)
if db_spend is None:
@ -149,7 +153,7 @@ class SpendCounterReseed:
# Warm even when 0 so subsequent reads hit cache, not DB.
try:
await spend_counter_cache.async_increment_cache(
key=counter_key, value=db_spend
key=counter_key, value=db_spend, refresh_ttl=True
)
except Exception:
verbose_proxy_logger.exception(

171
litellm/proxy/google_endpoints/endpoints.py
View file

@ -1,10 +1,6 @@
from datetime import datetime
from fastapi import APIRouter, Depends, Request, Response
from fastapi.responses import ORJSONResponse
from fastapi import APIRouter, Depends, HTTPException, Request, Response
from fastapi.responses import ORJSONResponse, StreamingResponse
import litellm
from litellm._uuid import uuid
from litellm.proxy._types import *
from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth, user_api_key_auth
from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
@ -30,12 +26,17 @@ async def google_generate_content(
fastapi_response: Response,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request
from litellm.proxy.proxy_server import (
general_settings,
llm_router,
proxy_config,
proxy_logging_obj,
select_data_generator,
user_api_base,
user_max_tokens,
user_model,
user_request_timeout,
user_temperature,
version,
)
@ -43,48 +44,33 @@ async def google_generate_content(
if "model" not in data:
data["model"] = model_name
# Extract generationConfig and pass it as config parameter
generation_config = data.pop("generationConfig", None)
if generation_config:
data["config"] = generation_config
# Add user authentication metadata for cost tracking
data = await add_litellm_data_to_request(
data=data,
request=request,
user_api_key_dict=user_api_key_dict,
proxy_config=proxy_config,
general_settings=general_settings,
version=version,
)
# Create logging object with full request metadata so callbacks (e.g. S3) get user/trace_id
data["litellm_call_id"] = request.headers.get(
"x-litellm-call-id", str(uuid.uuid4())
)
logging_obj, data = litellm.utils.function_setup(
original_function="agenerate_content",
rules_obj=litellm.utils.Rules(),
start_time=datetime.now(),
**data,
)
data["litellm_logging_obj"] = logging_obj
# call router
if llm_router is None:
raise HTTPException(status_code=500, detail="Router not initialized")
response = await llm_router.agenerate_content(**data)
success_headers = await ProxyBaseLLMRequestProcessing.build_litellm_proxy_success_headers_from_llm_response(
response=response,
request_data=data,
request=request,
user_api_key_dict=user_api_key_dict,
logging_obj=logging_obj,
version=version,
proxy_logging_obj=proxy_logging_obj,
)
fastapi_response.headers.update(success_headers)
return response
processor = ProxyBaseLLMRequestProcessing(data=data)
try:
return await processor.base_process_llm_request(
request=request,
fastapi_response=fastapi_response,
user_api_key_dict=user_api_key_dict,
route_type="agenerate_content",
proxy_logging_obj=proxy_logging_obj,
llm_router=llm_router,
general_settings=general_settings,
proxy_config=proxy_config,
select_data_generator=select_data_generator,
model=model_name,
user_model=user_model,
user_temperature=user_temperature,
user_request_timeout=user_request_timeout,
user_max_tokens=user_max_tokens,
user_api_base=user_api_base,
version=version,
)
except Exception as e:
raise await processor._handle_llm_api_exception(
e=e,
user_api_key_dict=user_api_key_dict,
proxy_logging_obj=proxy_logging_obj,
version=version,
)
@router.post(
@ -101,73 +87,52 @@ async def google_stream_generate_content(
fastapi_response: Response,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request
from litellm.proxy.proxy_server import (
general_settings,
llm_router,
proxy_config,
proxy_logging_obj,
select_data_generator,
user_api_base,
user_max_tokens,
user_model,
user_request_timeout,
user_temperature,
version,
)
data = await _read_request_body(request=request)
if "model" not in data:
data["model"] = model_name
data["stream"] = True
data["stream"] = True # enforce streaming for this endpoint
# Extract generationConfig and pass it as config parameter
generation_config = data.pop("generationConfig", None)
if generation_config:
data["config"] = generation_config
# Add user authentication metadata for cost tracking
data = await add_litellm_data_to_request(
data=data,
request=request,
user_api_key_dict=user_api_key_dict,
proxy_config=proxy_config,
general_settings=general_settings,
version=version,
)
# Create logging object with full request metadata so streaming END callbacks (e.g. S3) get user/trace_id
data["litellm_call_id"] = request.headers.get(
"x-litellm-call-id", str(uuid.uuid4())
)
logging_obj, data = litellm.utils.function_setup(
original_function="agenerate_content_stream",
rules_obj=litellm.utils.Rules(),
start_time=datetime.now(),
**data,
)
data["litellm_logging_obj"] = logging_obj
# call router
if llm_router is None:
raise HTTPException(status_code=500, detail="Router not initialized")
response = await llm_router.agenerate_content_stream(**data)
success_headers = await ProxyBaseLLMRequestProcessing.build_litellm_proxy_success_headers_from_llm_response(
response=response,
request_data=data,
request=request,
user_api_key_dict=user_api_key_dict,
logging_obj=logging_obj,
version=version,
proxy_logging_obj=proxy_logging_obj,
)
# Check if response is an async iterator (streaming response)
if response is not None and hasattr(response, "__aiter__"):
return StreamingResponse(
content=response,
media_type="text/event-stream",
headers=success_headers,
processor = ProxyBaseLLMRequestProcessing(data=data)
try:
return await processor.base_process_llm_request(
request=request,
fastapi_response=fastapi_response,
user_api_key_dict=user_api_key_dict,
route_type="agenerate_content_stream",
proxy_logging_obj=proxy_logging_obj,
llm_router=llm_router,
general_settings=general_settings,
proxy_config=proxy_config,
select_data_generator=select_data_generator,
model=model_name,
user_model=user_model,
user_temperature=user_temperature,
user_request_timeout=user_request_timeout,
user_max_tokens=user_max_tokens,
user_api_base=user_api_base,
version=version,
)
except Exception as e:
raise await processor._handle_llm_api_exception(
e=e,
user_api_key_dict=user_api_key_dict,
proxy_logging_obj=proxy_logging_obj,
version=version,
)
fastapi_response.headers.update(success_headers)
return response
@router.post(

20
litellm/proxy/management_endpoints/access_group_endpoints.py
View file

@ -236,13 +236,12 @@ async def _patch_key_caches_add_access_group(
) -> None:
"""Patch cached key objects to include access_group_id."""
for token in key_tokens:
cached_key = await user_api_key_cache.async_get_cache(key=token)
cached_key = await user_api_key_cache.async_get_cache(
key=token,
model_type=UserAPIKeyAuth,
)
if cached_key is None:
continue
if isinstance(cached_key, dict):
cached_key = UserAPIKeyAuth(**cached_key)
if not isinstance(cached_key, UserAPIKeyAuth):
continue
if cached_key.access_group_ids is None:
cached_key.access_group_ids = [access_group_id]
elif access_group_id not in cached_key.access_group_ids:
@ -267,12 +266,11 @@ async def _patch_key_caches_remove_access_group(
) -> None:
"""Patch cached key objects to remove access_group_id."""
for token in key_tokens:
cached_key = await user_api_key_cache.async_get_cache(key=token)
if cached_key is None:
continue
if isinstance(cached_key, dict):
cached_key = UserAPIKeyAuth(**cached_key)
if isinstance(cached_key, UserAPIKeyAuth) and cached_key.access_group_ids:
cached_key = await user_api_key_cache.async_get_cache(
key=token,
model_type=UserAPIKeyAuth,
)
if cached_key is not None and cached_key.access_group_ids:
cached_key.access_group_ids = [
ag for ag in cached_key.access_group_ids if ag != access_group_id
]

20
litellm/proxy/management_endpoints/key_management_endpoints.py
View file

@ -27,7 +27,7 @@ from fastapi import APIRouter, Depends, Header, HTTPException, Query, Request, s
import litellm
from litellm._logging import verbose_proxy_logger
from litellm._uuid import uuid
from litellm.caching import DualCache
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.constants import (
LENGTH_OF_LITELLM_GENERATED_KEY,
LITELLM_PROXY_ADMIN_NAME,
@ -1059,7 +1059,7 @@ async def _check_project_key_limits(
project_id: str,
data: Union[GenerateKeyRequest, UpdateKeyRequest],
prisma_client: PrismaClient,
user_api_key_cache: DualCache,
user_api_key_cache: UserApiKeyCache,
) -> None:
"""
Validate that key's models and budget respect its project's limits.
@ -1834,7 +1834,7 @@ async def _process_single_key_update(
user_api_key_dict: UserAPIKeyAuth,
litellm_changed_by: Optional[str],
prisma_client: Optional[PrismaClient],
user_api_key_cache: DualCache,
user_api_key_cache: UserApiKeyCache,
proxy_logging_obj: Any,
llm_router: Optional[Router],
user_custom_key_update: Optional[Callable] = None,
@ -3298,7 +3298,7 @@ async def _team_key_deletion_check(
user_api_key_dict: UserAPIKeyAuth,
key_info: LiteLLM_VerificationToken,
prisma_client: PrismaClient,
user_api_key_cache: DualCache,
user_api_key_cache: UserApiKeyCache,
):
is_team_key = _is_team_key(data=key_info)
@ -3341,7 +3341,7 @@ async def _team_key_deletion_check(
async def can_modify_verification_token(
key_info: LiteLLM_VerificationToken,
user_api_key_cache: DualCache,
user_api_key_cache: UserApiKeyCache,
user_api_key_dict: UserAPIKeyAuth,
prisma_client: PrismaClient,
) -> bool:
@ -3415,7 +3415,7 @@ async def can_modify_verification_token(
async def delete_verification_tokens(
tokens: List,
user_api_key_cache: DualCache,
user_api_key_cache: UserApiKeyCache,
user_api_key_dict: UserAPIKeyAuth,
litellm_changed_by: Optional[str] = None,
) -> Tuple[Optional[Dict], List[LiteLLM_VerificationToken]]:
@ -3605,7 +3605,7 @@ async def _persist_deleted_verification_tokens(
async def delete_key_aliases(
key_aliases: List[str],
user_api_key_cache: DualCache,
user_api_key_cache: UserApiKeyCache,
prisma_client: PrismaClient,
user_api_key_dict: UserAPIKeyAuth,
litellm_changed_by: Optional[str] = None,
@ -3862,7 +3862,7 @@ async def _execute_virtual_key_regeneration(
data: Optional[RegenerateKeyRequest],
user_api_key_dict: UserAPIKeyAuth,
litellm_changed_by: Optional[str],
user_api_key_cache: DualCache,
user_api_key_cache: UserApiKeyCache,
proxy_logging_obj: ProxyLogging,
) -> GenerateKeyResponse:
"""Generate new token, update DB, invalidate cache, and return response."""
@ -4152,7 +4152,7 @@ async def _check_proxy_or_team_admin_for_key(
key_in_db: LiteLLM_VerificationToken,
user_api_key_dict: UserAPIKeyAuth,
prisma_client: PrismaClient,
user_api_key_cache: DualCache,
user_api_key_cache: UserApiKeyCache,
) -> None:
if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value:
return
@ -5173,7 +5173,7 @@ async def _check_key_admin_access(
user_api_key_dict: UserAPIKeyAuth,
hashed_token: str,
prisma_client: Any,
user_api_key_cache: DualCache,
user_api_key_cache: UserApiKeyCache,
route: str,
) -> None:
"""

21
litellm/proxy/management_endpoints/ui_sso.py
View file

@ -39,9 +39,9 @@ from fastapi import APIRouter, Depends, Header, HTTPException, Request, status
from fastapi.responses import RedirectResponse
import litellm
from litellm.caching.dual_cache import DualCache
from litellm._logging import verbose_proxy_logger
from litellm._uuid import uuid
from litellm.caching import DualCache
from litellm.constants import (
CLI_SSO_SESSION_CACHE_KEY_PREFIX,
CLI_SSO_SESSION_TTL_SECONDS,
@ -75,7 +75,11 @@ from litellm.proxy._types import (
UserAPIKeyAuth,
)
from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken, get_user_object
from litellm.proxy.auth.auth_utils import _get_request_ip_address, _has_user_setup_sso
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.proxy.auth.auth_utils import (
_get_request_ip_address,
_has_user_setup_sso,
)
from litellm.proxy.auth.handle_jwt import JWTHandler
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.common_utils.admin_ui_utils import (
@ -1301,7 +1305,7 @@ async def get_existing_user_info_from_db(
user_id: Optional[str],
user_email: Optional[str],
prisma_client: PrismaClient,
user_api_key_cache: DualCache,
user_api_key_cache: UserApiKeyCache,
proxy_logging_obj: ProxyLogging,
) -> Optional[LiteLLM_UserTable]:
try:
@ -1325,7 +1329,7 @@ async def get_existing_user_info_from_db(
async def get_user_info_from_db(
result: Union[CustomOpenID, OpenID, dict],
prisma_client: PrismaClient,
user_api_key_cache: DualCache,
user_api_key_cache: UserApiKeyCache,
proxy_logging_obj: ProxyLogging,
user_email: Optional[str],
user_defined_values: Optional[SSOUserDefinedValues],
@ -1445,7 +1449,7 @@ async def _sync_user_role_from_jwt_role_map(
received_response: Optional[dict],
user_info: Optional[Union[LiteLLM_UserTable, NewUserResponse]],
prisma_client: PrismaClient,
user_api_key_cache: DualCache,
user_api_key_cache: UserApiKeyCache,
user_defined_values: Optional[SSOUserDefinedValues],
) -> None:
"""
@ -1484,11 +1488,8 @@ async def _sync_user_role_from_jwt_role_map(
user_info.user_role = mapped_role.value
await user_api_key_cache.async_set_cache(
key=user_info.user_id,
value=(
user_info.model_dump()
if hasattr(user_info, "model_dump")
else dict(user_info)
),
value=user_info,
model_type=LiteLLM_UserTable,
)

4
litellm/proxy/management_helpers/team_member_permission_checks.py
View file

@ -1,6 +1,5 @@
from typing import List, Optional
from litellm.caching import DualCache
from litellm.proxy._types import (
KeyManagementRoutes,
LiteLLM_TeamTableCachedObj,
@ -12,6 +11,7 @@ from litellm.proxy._types import (
ProxyException,
UserAPIKeyAuth,
)
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.proxy.auth.auth_checks import get_team_object
from litellm.proxy.auth.route_checks import RouteChecks
from litellm.proxy.utils import PrismaClient
@ -65,7 +65,7 @@ class TeamMemberPermissionChecks:
user_api_key_dict: UserAPIKeyAuth,
route: KeyManagementRoutes,
prisma_client: PrismaClient,
user_api_key_cache: DualCache,
user_api_key_cache: UserApiKeyCache,
existing_key_row: LiteLLM_VerificationToken,
):
"""

396
litellm/proxy/proxy_server.py
View file

@ -78,8 +78,11 @@ from litellm.proxy._types import (
InvitationNew,
InvitationUpdate,
Litellm_EntityType,
LiteLLM_EndUserTable,
LiteLLM_JWTAuth,
LiteLLM_TagTable,
LiteLLM_TeamTable,
LiteLLM_TeamTableCachedObj,
LiteLLM_UserTable,
LitellmUserRoles,
PassThroughGenericEndpoint,
@ -94,6 +97,7 @@ from litellm.proxy._types import (
UI_TEAM_ID,
UserAPIKeyAuth,
)
from litellm.proxy.common_utils.cache_pydantic_utils import CacheCodec
from litellm.proxy.common_utils.callback_utils import (
normalize_callback_names,
process_callback,
@ -206,6 +210,7 @@ from litellm import Router
from litellm._logging import verbose_proxy_logger, verbose_router_logger
from litellm.caching.caching import DualCache, RedisCache
from litellm.caching.redis_cluster_cache import RedisClusterCache
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.constants import (
_REALTIME_BODY_CACHE_SIZE,
APSCHEDULER_COALESCE,
@ -1612,7 +1617,7 @@ prisma_client: Optional[PrismaClient] = None
shared_aiohttp_session: Optional["ClientSession"] = (
None # Global shared session for connection reuse
)
user_api_key_cache = DualCache(
user_api_key_cache: UserApiKeyCache = UserApiKeyCache(
default_in_memory_ttl=UserAPIKeyCacheTTLEnum.in_memory_cache_ttl.value
)
spend_counter_cache = DualCache(
@ -1798,12 +1803,16 @@ async def get_current_spend(counter_key: str, fallback_spend: float) -> float:
3. Reseed from authoritative DB spend (counter expired, cross-pod stale)
4. Caller-supplied fallback (DB unavailable, cold start)
"""
# 1. Try Redis first (cross-pod authoritative)
# 1. Redis first (cross-pod authoritative). On clean miss, skip
# in-memory: per-pod in-memory only has this pod's writes, so it
# would mask cross-pod increments.
redis_clean_miss = False
if spend_counter_cache.redis_cache is not None:
try:
val = await spend_counter_cache.redis_cache.async_get_cache(key=counter_key)
if val is not None:
return float(val)
redis_clean_miss = True
except Exception as e:
verbose_proxy_logger.debug(
"get_current_spend: Redis read failed for %s, falling back to in-memory: %s",
@ -1811,10 +1820,11 @@ async def get_current_spend(counter_key: str, fallback_spend: float) -> float:
e,
)
# 2. Fall back to in-memory counter (single-instance or Redis failure)
val = spend_counter_cache.in_memory_cache.get_cache(key=counter_key)
if val is not None:
return float(val)
# 2. In-memory only when Redis is unreachable.
if not redis_clean_miss:
val = spend_counter_cache.in_memory_cache.get_cache(key=counter_key)
if val is not None:
return float(val)
# 3. Reseed from DB - fallback_spend lags cross-pod, would allow bypass.
db_spend = await SpendCounterReseed.coalesced(
@ -1976,10 +1986,12 @@ async def _init_and_increment_spend_counter(
base_spend = getattr(source, "spend", 0.0) or 0.0
if base_spend > 0:
await spend_counter_cache.async_increment_cache(
key=counter_key, value=base_spend
key=counter_key, value=base_spend, refresh_ttl=True
)
await spend_counter_cache.async_increment_cache(key=counter_key, value=increment)
await spend_counter_cache.async_increment_cache(
key=counter_key, value=increment, refresh_ttl=True
)
async def update_cache( # noqa: PLR0915
@ -2007,14 +2019,16 @@ async def update_cache( # noqa: PLR0915
else:
hashed_token = token
verbose_proxy_logger.debug("_update_key_cache: hashed_token=%s", hashed_token)
existing_spend_obj: LiteLLM_VerificationTokenView = await user_api_key_cache.async_get_cache(key=hashed_token) # type: ignore
existing_spend_obj = await user_api_key_cache.async_get_cache(
key=hashed_token, model_type=UserAPIKeyAuth
)
verbose_proxy_logger.debug(
f"_update_key_cache: existing_spend_obj={existing_spend_obj}"
)
if existing_spend_obj is None:
return
else:
existing_spend = existing_spend_obj.spend
existing_spend = existing_spend_obj.spend or 0.0
# Calculate the new cost by adding the existing cost and response_cost
new_spend = existing_spend + response_cost
@ -2072,41 +2086,48 @@ async def update_cache( # noqa: PLR0915
existing_team_member_spend + response_cost
)
# Update the cost column for the given token
# Existing spend_obj is mutated; UserApiKeyCache.async_set_cache_pipeline turns
# BaseModel values into dicts for Redis (same Codec path as async_set_cache).
existing_spend_obj.spend = new_spend
values_to_update_in_cache.append((hashed_token, existing_spend_obj))
### UPDATE USER SPEND ###
async def _update_user_cache():
## UPDATE CACHE FOR USER ID + GLOBAL PROXY
if response_cost is None:
return
user_ids = [user_id]
try:
for _id in user_ids:
# Fetch the existing cost for the given user
if _id is None:
continue
existing_spend_obj = await user_api_key_cache.async_get_cache(key=_id)
if existing_spend_obj is None:
cached_user = await user_api_key_cache.async_get_cache(key=_id)
if cached_user is None:
# do nothing if there is no cache value
return
existing_spend_obj = CacheCodec.deserialize(
cached_user, LiteLLM_UserTable
)
if existing_spend_obj is None:
return
verbose_proxy_logger.debug(
f"_update_user_db: existing spend: {existing_spend_obj}; response_cost: {response_cost}"
)
if isinstance(existing_spend_obj, dict):
existing_spend = existing_spend_obj["spend"]
else:
existing_spend = existing_spend_obj.spend
existing_spend = existing_spend_obj.spend or 0.0
# Calculate the new cost by adding the existing cost and response_cost
new_spend = existing_spend + response_cost
# Update the cost column for the given user
if isinstance(existing_spend_obj, dict):
existing_spend_obj["spend"] = new_spend
values_to_update_in_cache.append((_id, existing_spend_obj))
else:
existing_spend_obj.spend = new_spend
values_to_update_in_cache.append((_id, existing_spend_obj.json()))
existing_spend_obj.spend = new_spend
values_to_update_in_cache.append(
(
_id,
CacheCodec.serialize(
existing_spend_obj, model_type=LiteLLM_UserTable
),
)
)
## UPDATE GLOBAL PROXY ##
global_proxy_spend = await user_api_key_cache.async_get_cache(
key="{}:spend".format(litellm_proxy_admin_name)
@ -2138,31 +2159,33 @@ async def update_cache( # noqa: PLR0915
_id = "end_user_id:{}".format(end_user_id)
try:
# Fetch the existing cost for the given user
existing_spend_obj = await user_api_key_cache.async_get_cache(key=_id)
if existing_spend_obj is None:
cached_end_user = await user_api_key_cache.async_get_cache(key=_id)
if cached_end_user is None:
# if user does not exist in LiteLLM_UserTable, create a new user
# do nothing if end-user not in api key cache
return
existing_spend_obj = CacheCodec.deserialize(
cached_end_user, LiteLLM_EndUserTable
)
if existing_spend_obj is None:
return
verbose_proxy_logger.debug(
f"_update_end_user_db: existing spend: {existing_spend_obj}; response_cost: {response_cost}"
)
if existing_spend_obj is None:
existing_spend = 0
else:
if isinstance(existing_spend_obj, dict):
existing_spend = existing_spend_obj["spend"]
else:
existing_spend = existing_spend_obj.spend
existing_spend = existing_spend_obj.spend or 0.0
# Calculate the new cost by adding the existing cost and response_cost
new_spend = existing_spend + response_cost
# Update the cost column for the given user
if isinstance(existing_spend_obj, dict):
existing_spend_obj["spend"] = new_spend
values_to_update_in_cache.append((_id, existing_spend_obj))
else:
existing_spend_obj.spend = new_spend
values_to_update_in_cache.append((_id, existing_spend_obj.json()))
existing_spend_obj.spend = new_spend
values_to_update_in_cache.append(
(
_id,
CacheCodec.serialize(
existing_spend_obj, model_type=LiteLLM_EndUserTable
),
)
)
except Exception as e:
verbose_proxy_logger.warning(
"Spend tracking - failed to update end user spend in cache. "
@ -2181,36 +2204,32 @@ async def update_cache( # noqa: PLR0915
_id = "team_id:{}".format(team_id)
try:
# Fetch the existing cost for the given user
existing_spend_obj: Optional[LiteLLM_TeamTable] = (
await user_api_key_cache.async_get_cache(key=_id)
cached_team = await user_api_key_cache.async_get_cache(key=_id)
if cached_team is None:
# do nothing if team not in api key cache
return
existing_spend_obj: Optional[LiteLLM_TeamTableCachedObj] = (
CacheCodec.deserialize(cached_team, LiteLLM_TeamTableCachedObj)
)
if existing_spend_obj is None:
# do nothing if team not in api key cache
return
verbose_proxy_logger.debug(
f"_update_team_db: existing spend: {existing_spend_obj}; response_cost: {response_cost}"
)
if existing_spend_obj is None:
existing_spend: Optional[float] = 0.0
else:
if isinstance(existing_spend_obj, dict):
existing_spend = existing_spend_obj["spend"]
else:
existing_spend = existing_spend_obj.spend
if existing_spend is None:
existing_spend = 0.0
existing_spend: float = existing_spend_obj.spend or 0.0
# Calculate the new cost by adding the existing cost and response_cost
new_spend = existing_spend + response_cost
# Update the cost column for the given user
if isinstance(existing_spend_obj, dict):
existing_spend_obj["spend"] = new_spend
values_to_update_in_cache.append((_id, existing_spend_obj))
else:
existing_spend_obj.spend = new_spend
values_to_update_in_cache.append((_id, existing_spend_obj))
existing_spend_obj.spend = new_spend
values_to_update_in_cache.append(
(
_id,
CacheCodec.serialize(
existing_spend_obj, model_type=LiteLLM_TeamTableCachedObj
),
)
)
except Exception as e:
verbose_proxy_logger.warning(
"Spend tracking - failed to update team spend in cache. "
@ -2237,32 +2256,32 @@ async def update_cache( # noqa: PLR0915
cache_key = f"tag:{tag_name}"
# Fetch the existing tag object from cache
existing_tag_obj = await user_api_key_cache.async_get_cache(
key=cache_key
)
if existing_tag_obj is None:
cached_tag = await user_api_key_cache.async_get_cache(key=cache_key)
if cached_tag is None:
# do nothing if tag not in api key cache
continue
existing_tag_obj = CacheCodec.deserialize(cached_tag, LiteLLM_TagTable)
if existing_tag_obj is None:
continue
verbose_proxy_logger.debug(
f"_update_tag_cache: existing spend for tag={tag_name}: {existing_tag_obj}; response_cost: {response_cost}"
)
if isinstance(existing_tag_obj, dict):
existing_spend = existing_tag_obj.get("spend", 0) or 0
else:
existing_spend = getattr(existing_tag_obj, "spend", 0) or 0
existing_spend = existing_tag_obj.spend or 0.0
# Calculate the new cost by adding the existing cost and response_cost
new_spend = existing_spend + response_cost
# Update the spend column for the given tag
if isinstance(existing_tag_obj, dict):
existing_tag_obj["spend"] = new_spend
values_to_update_in_cache.append((cache_key, existing_tag_obj))
else:
existing_tag_obj.spend = new_spend
values_to_update_in_cache.append((cache_key, existing_tag_obj))
existing_tag_obj.spend = new_spend
values_to_update_in_cache.append(
(
cache_key,
CacheCodec.serialize(
existing_tag_obj, model_type=LiteLLM_TagTable
),
)
)
except Exception as e:
verbose_proxy_logger.warning(
"Spend tracking - failed to update tag spend in cache. "
@ -2930,8 +2949,9 @@ class ProxyConfig:
def _init_cache(
self,
cache_params: dict,
enable_redis_auth_cache: bool = False,
):
global redis_usage_cache, llm_router
global redis_usage_cache, llm_router, general_settings
from litellm import Cache
if "default_in_memory_ttl" in cache_params:
@ -2947,7 +2967,29 @@ class ProxyConfig:
):
## INIT PROXY REDIS USAGE CLIENT ##
redis_usage_cache = litellm.cache.cache
spend_counter_cache.redis_cache = redis_usage_cache
spend_counter_cache.attach_redis_cache(
redis_usage_cache,
default_redis_ttl=litellm.default_redis_ttl,
)
# Note: PKCE verifier storage uses redis_usage_cache directly (not
# user_api_key_cache) to avoid routing all API-key lookups through Redis.
if enable_redis_auth_cache is True:
user_api_key_cache.attach_redis_cache(
redis_usage_cache,
default_redis_ttl=litellm.default_redis_ttl,
)
verbose_proxy_logger.info(
"enable_redis_auth_cache=True: attached Redis to "
"user_api_key_cache — virtual-key lookups are now "
"shared across all proxy workers."
)
else:
verbose_proxy_logger.info(
"enable_redis_auth_cache is not set: user_api_key_cache "
"remains in-memory only (per-worker). Set "
"litellm_settings.enable_redis_auth_cache: true to share "
"the auth cache across workers and reduce DB load."
)
litellm_config_cache.redis_cache = redis_usage_cache
# Note: PKCE verifier storage uses redis_usage_cache directly (not
# user_api_key_cache) to avoid routing all API-key lookups through Redis.
@ -3273,7 +3315,13 @@ class ProxyConfig:
cache_params[key] = get_secret(value)
## to pass a complete url, or set ssl=True, etc. just set it as `os.environ[REDIS_URL] = <your-redis-url>`, _redis.py checks for REDIS specific environment variables
self._init_cache(cache_params=cache_params)
self._init_cache(
cache_params=cache_params,
enable_redis_auth_cache=litellm_settings.get(
"enable_redis_auth_cache", False
)
is True,
)
if litellm.cache is not None:
verbose_proxy_logger.debug(
f"{blue_color_code}Set Cache on LiteLLM Proxy{reset_color_code}"
@ -3544,21 +3592,23 @@ class ProxyConfig:
verbose_proxy_logger.critical(
"LITELLM_MASTER_KEY is not set! All requests will be treated as INTERNAL_USER with no admin access. Set LITELLM_MASTER_KEY for production use."
)
### USER API KEY CACHE IN-MEMORY TTL ###
### USER API KEY CACHE TTL (in-memory + Redis when Redis auth sharing is enabled) ###
user_api_key_cache_ttl = general_settings.get(
"user_api_key_cache_ttl", None
)
if user_api_key_cache_ttl is not None:
ttl = float(user_api_key_cache_ttl)
# Mirror TTL on Redis as well when ``litellm_settings.enable_redis_auth_cache``
# attaches Redis to ``user_api_key_cache``; otherwise DualCache misses in
# memory fall back to a key that outlasts ``user_api_key_cache_ttl``.
user_api_key_cache.update_cache_ttl(
default_in_memory_ttl=float(user_api_key_cache_ttl),
default_redis_ttl=None, # user_api_key_cache uses in-memory TTL only; Redis not configured for key lookups
default_in_memory_ttl=ttl,
default_redis_ttl=ttl,
)
### PKCE MULTI-INSTANCE PREREQUISITE CHECK ###
# PKCE verifiers are stored in redis_usage_cache when available so they can
# be read back by any instance (not just the one that started the auth flow).
# user_api_key_cache is intentionally left in-memory-only to avoid routing
# all API-key lookups through Redis.
use_pkce = os.getenv("GENERIC_CLIENT_USE_PKCE", "false").lower() == "true"
if use_pkce and redis_usage_cache is None:
global _pkce_no_redis_warning_emitted
@ -6287,7 +6337,7 @@ class ProxyStartupEvent:
cls,
general_settings: dict,
prisma_client: Optional[PrismaClient],
user_api_key_cache: DualCache,
user_api_key_cache: UserApiKeyCache,
):
"""Initialize JWT auth on startup"""
if general_settings.get("litellm_jwtauth", None) is not None:
@ -6336,7 +6386,7 @@ class ProxyStartupEvent:
async def _warm_global_spend_cache(
cls,
litellm_proxy_admin_name: str,
user_api_key_cache: DualCache,
user_api_key_cache: UserApiKeyCache,
prisma_client: PrismaClient,
) -> None:
"""Warm global spend cache once at startup to reduce impact of first wave of requests."""
@ -6976,7 +7026,7 @@ class ProxyStartupEvent:
cls,
database_url: Optional[str],
proxy_logging_obj: ProxyLogging,
user_api_key_cache: DualCache,
user_api_key_cache: UserApiKeyCache,
) -> Optional[PrismaClient]:
"""
- Sets up prisma client
@ -10278,6 +10328,101 @@ def _paginate_models_response(
}
def _team_models_resolve_to_names(
team_models: List[str], access_groups: Dict[str, Any]
) -> List[str]:
"""Expand team model entries (including access group names) to concrete model names."""
resolved: List[str] = []
for name in team_models:
if name in access_groups:
resolved.extend(access_groups[name])
else:
resolved.append(name)
return resolved
async def _load_team_object_for_model_filter(
team_id: str, prisma_client: PrismaClient
) -> Optional[LiteLLM_TeamTable]:
"""Load team row from DB; returns None if missing or on error."""
try:
team_db_object = await prisma_client.db.litellm_teamtable.find_unique(
where={"team_id": team_id}
)
if team_db_object is None:
verbose_proxy_logger.warning(f"Team {team_id} not found in database")
return None
return LiteLLM_TeamTable(**team_db_object.model_dump())
except Exception as e:
verbose_proxy_logger.exception(f"Error fetching team {team_id}: {str(e)}")
return None
async def _gather_team_accessible_model_ids(
team_object: LiteLLM_TeamTable,
team_id: str,
prisma_client: PrismaClient,
llm_router: Router,
) -> Set[str]:
"""Collect model IDs the team can use from router config and DB."""
team_accessible_model_ids: Set[str] = set()
access_groups = llm_router.get_model_access_groups() if llm_router else {}
if (
not team_object.models
or SpecialModelNames.all_proxy_models.value in team_object.models
):
model_list = llm_router.get_model_list() if llm_router else []
if model_list is not None:
for model in model_list:
model_id = model.get("model_info", {}).get("id", None)
if model_id is None:
continue
team_model_id = model.get("model_info", {}).get("team_id", None)
if team_model_id is None or team_model_id == team_id:
team_accessible_model_ids.add(model_id)
else:
resolved_model_names: Set[str] = set()
for model_name in team_object.models:
if model_name in access_groups:
resolved_model_names.update(access_groups[model_name])
else:
resolved_model_names.add(model_name)
for model_name in resolved_model_names:
_models = (
llm_router.get_model_list(model_name=model_name, team_id=team_id)
if llm_router
else []
)
if _models is not None:
for model in _models:
model_id = model.get("model_info", {}).get("id", None)
if model_id is not None:
team_accessible_model_ids.add(model_id)
try:
if (
team_object.models
and SpecialModelNames.all_proxy_models.value not in team_object.models
):
_resolved_names = _team_models_resolve_to_names(
team_object.models, access_groups
)
db_models = await prisma_client.db.litellm_proxymodeltable.find_many(
where={"model_name": {"in": _resolved_names}}
)
for db_model in db_models:
if db_model.model_id:
team_accessible_model_ids.add(db_model.model_id)
except Exception as e:
verbose_proxy_logger.debug(
f"Error querying database models for team {team_id}: {str(e)}"
)
return team_accessible_model_ids
async def _filter_models_by_team_id(
all_models: List[Dict[str, Any]],
team_id: str,
@ -10300,78 +10445,13 @@ async def _filter_models_by_team_id(
Returns:
Filtered list of models
"""
# Get team from database
try:
team_db_object = await prisma_client.db.litellm_teamtable.find_unique(
where={"team_id": team_id}
)
if team_db_object is None:
verbose_proxy_logger.warning(f"Team {team_id} not found in database")
# If team doesn't exist, return empty list
return []
team_object = LiteLLM_TeamTable(**team_db_object.model_dump())
except Exception as e:
verbose_proxy_logger.exception(f"Error fetching team {team_id}: {str(e)}")
team_object = await _load_team_object_for_model_filter(team_id, prisma_client)
if team_object is None:
return []
# Get models accessible to this team (similar to _add_team_models_to_all_models)
team_accessible_model_ids: Set[str] = set()
if (
not team_object.models # empty list = all model access
or SpecialModelNames.all_proxy_models.value in team_object.models
):
# Team has access to all models
model_list = llm_router.get_model_list() if llm_router else []
if model_list is not None:
for model in model_list:
model_id = model.get("model_info", {}).get("id", None)
if model_id is None:
continue
# if team model id set, check if team id matches
team_model_id = model.get("model_info", {}).get("team_id", None)
can_add_model = False
if team_model_id is None:
can_add_model = True
elif team_model_id == team_id:
can_add_model = True
if can_add_model:
team_accessible_model_ids.add(model_id)
else:
# Team has access to specific models
for model_name in team_object.models:
_models = (
llm_router.get_model_list(model_name=model_name, team_id=team_id)
if llm_router
else []
)
if _models is not None:
for model in _models:
model_id = model.get("model_info", {}).get("id", None)
if model_id is not None:
team_accessible_model_ids.add(model_id)
# Also search database for models accessible to this team
# This complements the config search done above
try:
if (
team_object.models
and SpecialModelNames.all_proxy_models.value not in team_object.models
):
# Team has specific models - check database for those model names
db_models = await prisma_client.db.litellm_proxymodeltable.find_many(
where={"model_name": {"in": team_object.models}}
)
for db_model in db_models:
model_id = db_model.model_id
if model_id:
team_accessible_model_ids.add(model_id)
except Exception as e:
verbose_proxy_logger.debug(
f"Error querying database models for team {team_id}: {str(e)}"
)
team_accessible_model_ids = await _gather_team_accessible_model_ids(
team_object, team_id, prisma_client, llm_router
)
# Filter models based on direct_access or access_via_team_ids
# Models are already enriched with these fields before this function is called

5
litellm/proxy/utils.py
View file

@ -101,6 +101,7 @@ from litellm.proxy._types import (
UserAPIKeyAuth,
)
from litellm.proxy.auth.route_checks import RouteChecks
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.proxy.db.create_views import (
create_missing_views,
should_create_missing_views,
@ -340,7 +341,7 @@ class ProxyLogging:
def __init__(
self,
user_api_key_cache: DualCache,
user_api_key_cache: UserApiKeyCache,
premium_user: bool = False,
):
## INITIALIZE LITELLM CALLBACKS ##
@ -5715,7 +5716,7 @@ async def get_available_models_for_user(
include_model_access_groups: bool = False,
only_model_access_groups: bool = False,
return_wildcard_routes: bool = False,
user_api_key_cache: Optional["DualCache"] = None,
user_api_key_cache: Optional["UserApiKeyCache"] = None,
) -> List[str]:
"""
Get the list of models available to a user based on their API key and team permissions.

29
litellm/rerank_api/main.py
View file

@ -163,19 +163,21 @@ def rerank( # noqa: PLR0915
model_response = RerankResponse()
rerank_litellm_params = {
"litellm_call_id": litellm_call_id,
"proxy_server_request": proxy_server_request,
"model_info": model_info,
"preset_cache_key": None,
"stream_response": {},
**optional_params.model_dump(exclude_unset=True),
}
litellm_logging_obj.update_from_kwargs(
kwargs=kwargs,
model=model,
user=user,
optional_params=dict(optional_rerank_params),
litellm_params={
"litellm_call_id": litellm_call_id,
"proxy_server_request": proxy_server_request,
"model_info": model_info,
"preset_cache_key": None,
"stream_response": {},
**optional_params.model_dump(exclude_unset=True),
},
litellm_params=dict(rerank_litellm_params),
custom_llm_provider=_custom_llm_provider,
)
@ -214,6 +216,7 @@ def rerank( # noqa: PLR0915
headers=headers or litellm.headers or {},
client=client,
model_response=model_response,
litellm_params=rerank_litellm_params,
)
elif _custom_llm_provider == litellm.LlmProviders.AZURE_AI:
api_base = (
@ -235,6 +238,7 @@ def rerank( # noqa: PLR0915
headers=headers or litellm.headers or {},
client=client,
model_response=model_response,
litellm_params=rerank_litellm_params,
)
elif _custom_llm_provider == litellm.LlmProviders.INFINITY:
# Implement Infinity rerank logic
@ -265,6 +269,7 @@ def rerank( # noqa: PLR0915
headers=headers or litellm.headers or {},
client=client,
model_response=model_response,
litellm_params=rerank_litellm_params,
)
elif _custom_llm_provider == litellm.LlmProviders.TOGETHER_AI:
# Implement Together AI rerank logic
@ -318,6 +323,7 @@ def rerank( # noqa: PLR0915
headers=headers or litellm.headers or {},
client=client,
model_response=model_response,
litellm_params=rerank_litellm_params,
)
elif _custom_llm_provider == litellm.LlmProviders.NVIDIA_NIM:
if dynamic_api_key is None:
@ -346,6 +352,7 @@ def rerank( # noqa: PLR0915
headers=headers or litellm.headers or {},
client=client,
model_response=model_response,
litellm_params=rerank_litellm_params,
)
elif _custom_llm_provider == litellm.LlmProviders.BEDROCK:
api_base = (
@ -409,6 +416,7 @@ def rerank( # noqa: PLR0915
headers=headers or litellm.headers or {},
client=client,
model_response=model_response,
litellm_params=rerank_litellm_params,
)
elif _custom_llm_provider == litellm.LlmProviders.DEEPINFRA:
@ -442,6 +450,7 @@ def rerank( # noqa: PLR0915
headers=headers or litellm.headers or {},
client=client,
model_response=model_response,
litellm_params=rerank_litellm_params,
)
elif _custom_llm_provider == litellm.LlmProviders.FIREWORKS_AI:
api_key = (
@ -472,6 +481,7 @@ def rerank( # noqa: PLR0915
headers=headers or litellm.headers or {},
client=client,
model_response=model_response,
litellm_params=rerank_litellm_params,
)
elif _custom_llm_provider == litellm.LlmProviders.VOYAGE:
api_key = (
@ -500,6 +510,7 @@ def rerank( # noqa: PLR0915
headers=headers or litellm.headers or {},
client=client,
model_response=model_response,
litellm_params=rerank_litellm_params,
)
elif _custom_llm_provider == litellm.LlmProviders.WATSONX:
credentials = IBMWatsonXMixin.get_watsonx_credentials(
@ -527,6 +538,7 @@ def rerank( # noqa: PLR0915
headers=headers or litellm.headers or {},
client=client,
model_response=model_response,
litellm_params=rerank_litellm_params,
)
else:
# Generic handler for all providers that use base_llm_http_handler
@ -559,6 +571,7 @@ def rerank( # noqa: PLR0915
headers=headers or litellm.headers or {},
client=client,
model_response=model_response,
litellm_params=rerank_litellm_params,
)
# Placeholder return

739
litellm/responses/streaming_iterator.py
View file

@ -1,9 +1,12 @@
from __future__ import annotations
import asyncio
import json
import time
import traceback
from datetime import datetime
from typing import Any, Dict, List, Optional
from functools import lru_cache
from typing import Any, Dict, List, Literal, Optional
import httpx
@ -22,19 +25,26 @@ from litellm.litellm_core_utils.llm_response_utils.response_metadata import (
from litellm.litellm_core_utils.thread_pool_executor import executor
from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig
from litellm.responses.utils import ResponsesAPIRequestUtils
from litellm.types.llms.openai import (
OutputTextDeltaEvent,
ResponseAPIUsage,
ResponseCompletedEvent,
ResponsesAPIRequestParams,
ResponsesAPIResponse,
ResponsesAPIStreamEvents,
ResponsesAPIStreamingResponse,
)
from litellm.types.llms.openai import ResponsesAPIStreamEvents
from litellm.types.utils import CallTypes
from litellm.utils import CustomStreamWrapper, async_post_call_success_deployment_hook
@lru_cache(maxsize=1)
def _get_openai_response_types():
from litellm.types.llms import openai as openai_types
return openai_types
def _log_background_task_failure(task: "asyncio.Task[Any]", *, task_name: str) -> None:
if task.cancelled():
return
exception = task.exception()
if exception is not None:
verbose_logger.error("%s failed: %s", task_name, exception)
class BaseResponsesAPIStreamingIterator:
"""
Base class for streaming iterators that process responses from the Responses API.
@ -46,7 +56,7 @@ class BaseResponsesAPIStreamingIterator:
self,
response: httpx.Response,
model: str,
responses_api_provider_config: BaseResponsesAPIConfig,
responses_api_provider_config: Optional[BaseResponsesAPIConfig],
logging_obj: LiteLLMLoggingObj,
litellm_metadata: Optional[Dict[str, Any]] = None,
custom_llm_provider: Optional[str] = None,
@ -58,9 +68,13 @@ class BaseResponsesAPIStreamingIterator:
self.logging_obj = logging_obj
self.finished = False
self.responses_api_provider_config = responses_api_provider_config
self.completed_response: Optional[ResponsesAPIStreamingResponse] = None
self.completed_response: Optional[Any] = None
self.start_time = getattr(logging_obj, "start_time", datetime.now())
self._failure_handled = False # Track if failure handler has been called
self._completed_response_cached = False
self._completed_response_logged = False
self._completed_response_cache_hit: Optional[bool] = None
self._persist_completed_response_before_logging = True
self._stream_created_time: float = time.time()
# track request context for hooks
@ -101,7 +115,7 @@ class BaseResponsesAPIStreamingIterator:
llm_provider=self.custom_llm_provider or "",
)
def _process_chunk(self, chunk) -> Optional[ResponsesAPIStreamingResponse]:
def _process_chunk(self, chunk) -> Optional[Any]:
"""Process a single chunk of data from the stream"""
if not chunk:
return None
@ -122,6 +136,10 @@ class BaseResponsesAPIStreamingIterator:
# Format as ResponsesAPIStreamingResponse
if isinstance(parsed_chunk, dict):
if self.responses_api_provider_config is None:
raise ValueError(
"responses_api_provider_config is required to process live streaming chunks"
)
openai_responses_api_chunk = (
self.responses_api_provider_config.transform_streaming_response(
model=self.model,
@ -195,10 +213,11 @@ class BaseResponsesAPIStreamingIterator:
if self.litellm_metadata and self.litellm_metadata.get(
"encrypted_content_affinity_enabled"
):
openai_types = _get_openai_response_types()
event_type = getattr(openai_responses_api_chunk, "type", None)
if event_type in (
ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED,
ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE,
openai_types.ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED,
openai_types.ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE,
):
item = getattr(openai_responses_api_chunk, "item", None)
if item:
@ -219,10 +238,11 @@ class BaseResponsesAPIStreamingIterator:
# Store the completed response (also for incomplete/failed so logging still fires)
_chunk_type = getattr(openai_responses_api_chunk, "type", None)
openai_types = _get_openai_response_types()
if openai_responses_api_chunk and _chunk_type in (
ResponsesAPIStreamEvents.RESPONSE_COMPLETED,
ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE,
ResponsesAPIStreamEvents.RESPONSE_FAILED,
openai_types.ResponsesAPIStreamEvents.RESPONSE_COMPLETED,
openai_types.ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE,
openai_types.ResponsesAPIStreamEvents.RESPONSE_FAILED,
):
self.completed_response = openai_responses_api_chunk
# Add cost to usage object if include_cost_in_streaming_usage is True
@ -230,11 +250,11 @@ class BaseResponsesAPIStreamingIterator:
litellm.include_cost_in_streaming_usage
and self.logging_obj is not None
):
response_obj: Optional[ResponsesAPIResponse] = getattr(
response_obj: Optional[Any] = getattr(
openai_responses_api_chunk, "response", None
)
if response_obj:
usage_obj: Optional[ResponseAPIUsage] = getattr(
usage_obj: Optional[Any] = getattr(
response_obj, "usage", None
)
if usage_obj is not None:
@ -247,9 +267,13 @@ class BaseResponsesAPIStreamingIterator:
if cost is not None:
setattr(usage_obj, "cost", cost)
except Exception:
# Best-effort usage cost annotation should not break stream replay.
pass
if _chunk_type == ResponsesAPIStreamEvents.RESPONSE_FAILED:
if (
_chunk_type
== openai_types.ResponsesAPIStreamEvents.RESPONSE_FAILED
):
self._handle_logging_failed_response()
else:
self._handle_logging_completed_response()
@ -266,6 +290,59 @@ class BaseResponsesAPIStreamingIterator:
self._handle_failure(e)
raise
def _log_completed_response(self, *, is_async: bool) -> None:
if self._completed_response_logged:
return
self._completed_response_logged = True
if self._persist_completed_response_before_logging:
self._persist_completed_response_to_cache(is_async=is_async)
# Create a copy for logging to avoid modifying the response object that will be returned to the user
# The logging handlers may transform usage from Responses API format (input_tokens/output_tokens)
# to chat completion format (prompt_tokens/completion_tokens) for internal logging
# Use model_dump + model_validate instead of deepcopy to avoid pickle errors with
# Pydantic ValidatorIterator when response contains tool_choice with allowed_tools (fixes #17192)
logging_response = self.completed_response
if self.completed_response is not None and hasattr(
self.completed_response, "model_dump"
):
try:
logging_response = type(self.completed_response).model_validate(
self.completed_response.model_dump()
)
except Exception:
# Fallback to original if serialization fails
pass
end_time = datetime.now()
if is_async:
asyncio.create_task(
self.logging_obj.async_success_handler(
result=logging_response,
start_time=self.start_time,
end_time=end_time,
cache_hit=self._completed_response_cache_hit,
)
)
else:
run_async_function(
async_function=self.logging_obj.async_success_handler,
result=logging_response,
start_time=self.start_time,
end_time=end_time,
cache_hit=self._completed_response_cache_hit,
)
executor.submit(
self.logging_obj.success_handler,
result=logging_response,
cache_hit=self._completed_response_cache_hit,
start_time=self.start_time,
end_time=end_time,
)
self._run_post_success_hooks(end_time=end_time)
def _handle_logging_completed_response(self):
"""Base implementation - should be overridden by subclasses"""
pass
@ -296,6 +373,88 @@ class BaseResponsesAPIStreamingIterator:
)
self._handle_failure(exception)
def _get_completed_response_object(self) -> Optional[Any]:
openai_types = _get_openai_response_types()
completed_response = self.completed_response
if isinstance(completed_response, openai_types.ResponsesAPIResponse):
return completed_response
response_obj = getattr(completed_response, "response", None)
if isinstance(response_obj, openai_types.ResponsesAPIResponse):
return response_obj
return None
def _persist_completed_response_to_cache(self, *, is_async: bool) -> None:
if self._completed_response_cached:
return
completed_response = self.completed_response
openai_types = _get_openai_response_types()
if (
getattr(completed_response, "type", None)
!= openai_types.ResponsesAPIStreamEvents.RESPONSE_COMPLETED
):
return
response_obj = self._get_completed_response_object()
if response_obj is None:
return
caching_handler = getattr(self.logging_obj, "_llm_caching_handler", None)
if caching_handler is None:
return
request_kwargs = getattr(caching_handler, "request_kwargs", None)
if (
not isinstance(request_kwargs, dict)
or request_kwargs.get("stream") is not True
):
return
request_kwargs = request_kwargs.copy()
preset_cache_key = getattr(caching_handler, "preset_cache_key", None)
request_cache_key = request_kwargs.pop("cache_key", None)
if preset_cache_key is None:
preset_cache_key = request_cache_key
if request_kwargs.get("metadata") is None:
request_kwargs.pop("metadata", None)
request_kwargs.pop("custom_llm_provider", None)
if preset_cache_key is not None:
request_kwargs["cache_key"] = preset_cache_key
if not caching_handler._should_store_result_in_cache(
original_function=caching_handler.original_function,
kwargs=request_kwargs,
):
return
if litellm.cache is None:
return
cached_response = response_obj.model_dump_json()
if is_async:
cache_write_task = asyncio.create_task(
litellm.cache.async_add_cache(
cached_response,
dynamic_cache_object=getattr(caching_handler, "dual_cache", None),
**request_kwargs,
)
)
cache_write_task.add_done_callback(
lambda task: _log_background_task_failure(
task,
task_name="Responses stream cache write",
)
)
else:
litellm.cache.add_cache(
cached_response,
dynamic_cache_object=getattr(caching_handler, "dual_cache", None),
**request_kwargs,
)
self._completed_response_cached = True
async def _call_post_streaming_deployment_hook(self, chunk):
"""
Allow callbacks to modify streaming chunks before returning (parity with chat).
@ -480,7 +639,7 @@ class ResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator):
def __aiter__(self):
return self
async def __anext__(self) -> ResponsesAPIStreamingResponse:
async def __anext__(self) -> Any:
try:
self._check_max_streaming_duration()
while True:
@ -520,40 +679,7 @@ class ResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator):
def _handle_logging_completed_response(self):
"""Handle logging for completed responses in async context"""
# Create a copy for logging to avoid modifying the response object that will be returned to the user
# The logging handlers may transform usage from Responses API format (input_tokens/output_tokens)
# to chat completion format (prompt_tokens/completion_tokens) for internal logging
# Use model_dump + model_validate instead of deepcopy to avoid pickle errors with
# Pydantic ValidatorIterator when response contains tool_choice with allowed_tools (fixes #17192)
logging_response = self.completed_response
if self.completed_response is not None and hasattr(
self.completed_response, "model_dump"
):
try:
logging_response = type(self.completed_response).model_validate(
self.completed_response.model_dump()
)
except Exception:
# Fallback to original if serialization fails
pass
asyncio.create_task(
self.logging_obj.async_success_handler(
result=logging_response,
start_time=self.start_time,
end_time=datetime.now(),
cache_hit=None,
)
)
executor.submit(
self.logging_obj.success_handler,
result=logging_response,
cache_hit=None,
start_time=self.start_time,
end_time=datetime.now(),
)
self._run_post_success_hooks(end_time=datetime.now())
self._log_completed_response(is_async=True)
class SyncResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator):
@ -627,39 +753,7 @@ class SyncResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator):
def _handle_logging_completed_response(self):
"""Handle logging for completed responses in sync context"""
# Create a copy for logging to avoid modifying the response object that will be returned to the user
# The logging handlers may transform usage from Responses API format (input_tokens/output_tokens)
# to chat completion format (prompt_tokens/completion_tokens) for internal logging
# Use model_dump + model_validate instead of deepcopy to avoid pickle errors with
# Pydantic ValidatorIterator when response contains tool_choice with allowed_tools (fixes #17192)
logging_response = self.completed_response
if self.completed_response is not None and hasattr(
self.completed_response, "model_dump"
):
try:
logging_response = type(self.completed_response).model_validate(
self.completed_response.model_dump()
)
except Exception:
# Fallback to original if serialization fails
pass
run_async_function(
async_function=self.logging_obj.async_success_handler,
result=logging_response,
start_time=self.start_time,
end_time=datetime.now(),
cache_hit=None,
)
executor.submit(
self.logging_obj.success_handler,
result=logging_response,
cache_hit=None,
start_time=self.start_time,
end_time=datetime.now(),
)
self._run_post_success_hooks(end_time=datetime.now())
self._log_completed_response(is_async=False)
class MockResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator):
@ -683,90 +777,441 @@ class MockResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator):
request_data: Optional[Dict[str, Any]] = None,
call_type: Optional[str] = None,
):
super().__init__(
response=response,
transformed = responses_api_provider_config.transform_response_api_response(
model=model,
responses_api_provider_config=responses_api_provider_config,
raw_response=response,
logging_obj=logging_obj,
)
super().__init__(
response=httpx.Response(200),
model=model,
responses_api_provider_config=None,
logging_obj=logging_obj,
litellm_metadata=litellm_metadata,
custom_llm_provider=custom_llm_provider,
request_data=request_data,
call_type=call_type,
)
self._set_events_from_response(transformed=transformed, logging_obj=logging_obj)
# one-time transform
transformed = (
self.responses_api_provider_config.transform_response_api_response(
model=self.model,
raw_response=response,
logging_obj=logging_obj,
)
def _set_events_from_response(
self,
transformed: Any,
logging_obj: LiteLLMLoggingObj,
) -> None:
self._events = _build_synthetic_response_events(
transformed=transformed,
logging_obj=logging_obj,
chunk_size=self.CHUNK_SIZE,
)
full_text = self._collect_text(transformed)
# build a list of 5‑char delta events
deltas = [
OutputTextDeltaEvent(
type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA,
delta=full_text[i : i + self.CHUNK_SIZE],
item_id=transformed.id,
output_index=0,
content_index=0,
)
for i in range(0, len(full_text), self.CHUNK_SIZE)
]
# Add cost to usage object if include_cost_in_streaming_usage is True
if litellm.include_cost_in_streaming_usage and logging_obj is not None:
usage_obj: Optional[ResponseAPIUsage] = getattr(transformed, "usage", None)
if usage_obj is not None:
try:
cost: Optional[float] = logging_obj._response_cost_calculator(
result=transformed
)
if cost is not None:
setattr(usage_obj, "cost", cost)
except Exception:
# If cost calculation fails, continue without cost
pass
# append the completed event
self._events = deltas + [
ResponseCompletedEvent(
type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED,
response=transformed,
)
]
self._idx = 0
self.completed_response = self._events[-1]
def __aiter__(self):
return self
async def __anext__(self) -> ResponsesAPIStreamingResponse:
async def __anext__(self) -> Any:
if self._idx >= len(self._events):
raise StopAsyncIteration
evt = self._events[self._idx]
self._idx += 1
openai_types = _get_openai_response_types()
if (
getattr(evt, "type", None)
== openai_types.ResponsesAPIStreamEvents.RESPONSE_COMPLETED
):
self.completed_response = evt
self._log_completed_response(is_async=True)
return evt
def __iter__(self):
return self
def __next__(self) -> ResponsesAPIStreamingResponse:
def __next__(self) -> Any:
if self._idx >= len(self._events):
raise StopIteration
evt = self._events[self._idx]
self._idx += 1
openai_types = _get_openai_response_types()
if (
getattr(evt, "type", None)
== openai_types.ResponsesAPIStreamEvents.RESPONSE_COMPLETED
):
self.completed_response = evt
self._log_completed_response(is_async=False)
return evt
def _collect_text(self, resp: ResponsesAPIResponse) -> str:
out = ""
for out_item in resp.output:
item_type = getattr(out_item, "type", None)
if item_type == "message":
for c in getattr(out_item, "content", []):
out += c.text
return out
class CachedResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator):
def __init__(
self,
response: Any,
logging_obj: LiteLLMLoggingObj,
request_data: Optional[Dict[str, Any]] = None,
call_type: Optional[str] = None,
):
BaseResponsesAPIStreamingIterator.__init__(
self,
response=httpx.Response(200),
model=getattr(response, "model", ""),
responses_api_provider_config=None,
logging_obj=logging_obj,
litellm_metadata=None,
custom_llm_provider="cached_response",
request_data=request_data,
call_type=call_type,
)
self._completed_response_cache_hit = True
self._persist_completed_response_before_logging = False
self._events: List[Any] = []
self._idx = 0
self._set_events_from_response(transformed=response, logging_obj=logging_obj)
def _set_events_from_response(
self,
transformed: Any,
logging_obj: LiteLLMLoggingObj,
) -> None:
self._events = _build_synthetic_response_events(
transformed=transformed,
logging_obj=logging_obj,
chunk_size=MockResponsesAPIStreamingIterator.CHUNK_SIZE,
)
self._idx = 0
self.completed_response = self._events[-1]
def __aiter__(self):
return self
async def __anext__(self) -> Any:
if self._idx >= len(self._events):
raise StopAsyncIteration
evt = self._events[self._idx]
self._idx += 1
openai_types = _get_openai_response_types()
if (
getattr(evt, "type", None)
== openai_types.ResponsesAPIStreamEvents.RESPONSE_COMPLETED
):
self.completed_response = evt
self._log_completed_response(is_async=True)
return evt
def __iter__(self):
return self
def __next__(self) -> Any:
if self._idx >= len(self._events):
raise StopIteration
evt = self._events[self._idx]
self._idx += 1
openai_types = _get_openai_response_types()
if (
getattr(evt, "type", None)
== openai_types.ResponsesAPIStreamEvents.RESPONSE_COMPLETED
):
self.completed_response = evt
self._log_completed_response(is_async=False)
return evt
def _dump_response_object(obj: Any) -> Dict[str, Any]:
if hasattr(obj, "model_dump"):
return obj.model_dump()
if isinstance(obj, dict):
return obj
return {}
def _build_response_status_event(
event_type: Literal[
"response.created",
"response.in_progress",
],
transformed: Any,
) -> Any:
openai_types = _get_openai_response_types()
in_progress_response = transformed.model_copy(
deep=True,
update={"status": "in_progress", "output": []},
)
if event_type == openai_types.ResponsesAPIStreamEvents.RESPONSE_CREATED:
return openai_types.ResponseCreatedEvent(
type=event_type, response=in_progress_response
)
return openai_types.ResponseInProgressEvent(
type=event_type, response=in_progress_response
)
def _build_content_part_done_event(
*,
item_id: str,
output_index: int,
content_index: int,
part_payload: Dict[str, Any],
) -> Optional[Any]:
openai_types = _get_openai_response_types()
part_type = part_payload.get("type")
part: Any
if part_type == "output_text":
annotations = [
openai_types.BaseLiteLLMOpenAIResponseObject(**annotation)
for annotation in part_payload.get("annotations", []) or []
]
part = openai_types.ContentPartDonePartOutputText(
type="output_text",
text=str(part_payload.get("text") or ""),
annotations=annotations,
logprobs=part_payload.get("logprobs"),
)
elif part_type == "refusal":
part = openai_types.ContentPartDonePartRefusal(
type="refusal",
refusal=str(part_payload.get("refusal") or ""),
)
elif part_type == "reasoning_text":
part = openai_types.ContentPartDonePartReasoningText(
type="reasoning_text",
reasoning=str(part_payload.get("reasoning") or ""),
)
else:
return None
return openai_types.ContentPartDoneEvent(
type=openai_types.ResponsesAPIStreamEvents.CONTENT_PART_DONE,
item_id=item_id,
output_index=output_index,
content_index=content_index,
part=part,
)
def _add_text_like_part_events(
*,
events: List[Any],
item_id: str,
output_index: int,
content_index: int,
part_payload: Dict[str, Any],
chunk_size: int,
) -> None:
openai_types = _get_openai_response_types()
part_type = part_payload.get("type")
if part_type == "output_text":
text = str(part_payload.get("text") or "")
for i in range(0, len(text), chunk_size):
events.append(
openai_types.OutputTextDeltaEvent(
type=openai_types.ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA,
item_id=item_id,
output_index=output_index,
content_index=content_index,
delta=text[i : i + chunk_size],
)
)
for annotation_index, annotation in enumerate(
part_payload.get("annotations", []) or []
):
events.append(
openai_types.OutputTextAnnotationAddedEvent(
type=openai_types.ResponsesAPIStreamEvents.OUTPUT_TEXT_ANNOTATION_ADDED,
item_id=item_id,
output_index=output_index,
content_index=content_index,
annotation_index=annotation_index,
annotation=annotation,
)
)
events.append(
openai_types.OutputTextDoneEvent(
type=openai_types.ResponsesAPIStreamEvents.OUTPUT_TEXT_DONE,
item_id=item_id,
output_index=output_index,
content_index=content_index,
text=text,
)
)
elif part_type == "refusal":
refusal = str(part_payload.get("refusal") or "")
for i in range(0, len(refusal), chunk_size):
events.append(
openai_types.RefusalDeltaEvent(
type=openai_types.ResponsesAPIStreamEvents.REFUSAL_DELTA,
item_id=item_id,
output_index=output_index,
content_index=content_index,
delta=refusal[i : i + chunk_size],
)
)
events.append(
openai_types.RefusalDoneEvent(
type=openai_types.ResponsesAPIStreamEvents.REFUSAL_DONE,
item_id=item_id,
output_index=output_index,
content_index=content_index,
refusal=refusal,
)
)
def _build_synthetic_response_events(
*,
transformed: Any,
logging_obj: LiteLLMLoggingObj,
chunk_size: int,
) -> List[Any]:
openai_types = _get_openai_response_types()
if litellm.include_cost_in_streaming_usage and logging_obj is not None:
usage_obj: Optional[Any] = getattr(transformed, "usage", None)
if usage_obj is not None:
try:
cost: Optional[float] = logging_obj._response_cost_calculator(
result=transformed
)
if cost is not None:
setattr(usage_obj, "cost", cost)
except Exception:
pass
events: List[Any] = [
_build_response_status_event(
openai_types.ResponsesAPIStreamEvents.RESPONSE_CREATED, transformed
),
_build_response_status_event(
openai_types.ResponsesAPIStreamEvents.RESPONSE_IN_PROGRESS, transformed
),
]
sequence_number = 0
for output_index, output_item in enumerate(
getattr(transformed, "output", []) or []
):
output_item_payload = _dump_response_object(output_item)
item_id = str(output_item_payload.get("id") or transformed.id)
item_type = output_item_payload.get("type")
events.append(
openai_types.OutputItemAddedEvent(
type=openai_types.ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED,
output_index=output_index,
item=openai_types.BaseLiteLLMOpenAIResponseObject(
**output_item_payload
),
)
)
if item_type == "message":
for content_index, part in enumerate(
output_item_payload.get("content", []) or []
):
part_payload = _dump_response_object(part)
events.append(
openai_types.ContentPartAddedEvent(
type=openai_types.ResponsesAPIStreamEvents.CONTENT_PART_ADDED,
item_id=item_id,
output_index=output_index,
content_index=content_index,
part=openai_types.BaseLiteLLMOpenAIResponseObject(
**part_payload
),
)
)
_add_text_like_part_events(
events=events,
item_id=item_id,
output_index=output_index,
content_index=content_index,
part_payload=part_payload,
chunk_size=chunk_size,
)
done_event = _build_content_part_done_event(
item_id=item_id,
output_index=output_index,
content_index=content_index,
part_payload=part_payload,
)
if done_event is not None:
events.append(done_event)
elif item_type == "function_call":
arguments = str(output_item_payload.get("arguments") or "")
for i in range(0, len(arguments), chunk_size):
events.append(
openai_types.FunctionCallArgumentsDeltaEvent(
type=openai_types.ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DELTA,
item_id=item_id,
output_index=output_index,
delta=arguments[i : i + chunk_size],
)
)
events.append(
openai_types.FunctionCallArgumentsDoneEvent(
type=openai_types.ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DONE,
item_id=item_id,
output_index=output_index,
arguments=arguments,
)
)
elif item_type == "reasoning":
for summary_index, summary in enumerate(
output_item_payload.get("summary", []) or []
):
summary_payload = _dump_response_object(summary)
summary_text = str(summary_payload.get("text") or "")
for i in range(0, len(summary_text), chunk_size):
events.append(
openai_types.ReasoningSummaryTextDeltaEvent(
type=openai_types.ResponsesAPIStreamEvents.REASONING_SUMMARY_TEXT_DELTA,
item_id=item_id,
output_index=output_index,
summary_index=summary_index,
delta=summary_text[i : i + chunk_size],
)
)
sequence_number += 1
events.append(
openai_types.ReasoningSummaryTextDoneEvent(
type=openai_types.ResponsesAPIStreamEvents.REASONING_SUMMARY_TEXT_DONE,
item_id=item_id,
output_index=output_index,
sequence_number=sequence_number,
summary_index=summary_index,
text=summary_text,
)
)
sequence_number += 1
events.append(
openai_types.ReasoningSummaryPartDoneEvent(
type=openai_types.ResponsesAPIStreamEvents.REASONING_SUMMARY_PART_DONE,
item_id=item_id,
output_index=output_index,
sequence_number=sequence_number,
summary_index=summary_index,
part=openai_types.BaseLiteLLMOpenAIResponseObject(
**summary_payload
),
)
)
sequence_number += 1
events.append(
openai_types.OutputItemDoneEvent(
type=openai_types.ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE,
output_index=output_index,
sequence_number=sequence_number,
item=openai_types.BaseLiteLLMOpenAIResponseObject(
**output_item_payload
),
)
)
events.append(
openai_types.ResponseCompletedEvent(
type=openai_types.ResponsesAPIStreamEvents.RESPONSE_COMPLETED,
response=transformed,
)
)
return events
# ---------------------------------------------------------------------------
@ -951,8 +1396,8 @@ class ResponsesWebSocketStreaming:
# ---------------------------------------------------------------------------
_RESPONSE_CREATE_PARAMS: frozenset = (
ResponsesAPIRequestParams.__required_keys__
| ResponsesAPIRequestParams.__optional_keys__
_get_openai_response_types().ResponsesAPIRequestParams.__required_keys__
| _get_openai_response_types().ResponsesAPIRequestParams.__optional_keys__
)
_MANAGED_WS_SKIP_KWARGS: frozenset = frozenset(
@ -1085,7 +1530,7 @@ class ManagedResponsesWebSocketHandler:
@staticmethod
def _extract_output_messages(
completed_event: Dict[str, Any]
completed_event: Dict[str, Any],
) -> List[Dict[str, Any]]:
"""
Convert the output items in a ``response.completed`` event into

27
litellm/router.py
View file

@ -5261,11 +5261,34 @@ class Router:
"""
Initialize the Containers API endpoints on the router.
Container operations don't need model-based routing, so we call the
original function directly with the custom_llm_provider.
LiteLLM-managed container IDs (``cntr_...``) encode ``model_id`` and provider
metadata. When present, decode the ID, replace ``container_id`` with the
upstream value, and route through ``_ageneric_api_call_with_fallbacks`` so
deployment credentials (e.g. regional ``api_base`` for Azure) match
:meth:`_init_responses_api_endpoints`. Otherwise call the handler directly.
"""
if custom_llm_provider and "custom_llm_provider" not in kwargs:
kwargs["custom_llm_provider"] = custom_llm_provider
from litellm.responses.utils import ResponsesAPIRequestUtils
container_id = kwargs.get("container_id")
if isinstance(container_id, str):
decoded = ResponsesAPIRequestUtils._decode_container_id(container_id)
original_id = decoded.get("response_id", container_id)
if original_id != container_id:
kwargs["container_id"] = original_id
decoded_provider = decoded.get("custom_llm_provider")
if decoded_provider and kwargs.get("custom_llm_provider") == "openai":
kwargs["custom_llm_provider"] = decoded_provider
model_id = decoded.get("model_id")
if model_id:
kwargs["model"] = model_id
return await self._ageneric_api_call_with_fallbacks(
original_function=original_function,
**kwargs,
)
return await original_function(**kwargs)
async def _init_responses_api_endpoints(

3
litellm/router_strategy/tag_based_routing.py
View file

@ -106,7 +106,8 @@ def _match_deployment(
# check either didn't run (no request tags) or failed (step 1 returned
# None). Block the regex path so it cannot circumvent the operator's
# strict-tag policy.
strict_tag_check_failed = not match_any and bool(deployment_tags)
deployment_has_plain_tags = deployment_tags is not None and len(deployment_tags) > 0
strict_tag_check_failed = not match_any and deployment_has_plain_tags
if deployment_tag_regex and header_strings and not strict_tag_check_failed:
regex_match = _is_valid_deployment_tag_regex(
deployment_tag_regex, header_strings

27
litellm/router_utils/common_utils.py
View file

@ -23,21 +23,28 @@ def add_model_file_id_mappings(
healthy_deployments: Union[List[Dict], Dict], responses: List["OpenAIFileObject"]
) -> dict:
"""
Create a mapping of model name to file id
Create a mapping of model id to file id
{
"model_id": "file_id",
"model_id": "file_id",
}
`healthy_deployments` may be either a list of deployment dicts (multiple
matched deployments) or a single deployment dict (when the router resolved
a specific deployment, e.g. because the requested model matched a
`model_info.id`). Both shapes must be handled by extracting
`model_info.id` from each deployment.
"""
model_file_id_mapping = {}
if isinstance(healthy_deployments, list):
for deployment, response in zip(healthy_deployments, responses):
model_file_id_mapping[deployment.get("model_info", {}).get("id")] = (
response.id
)
elif isinstance(healthy_deployments, dict):
for model_id, file_id in healthy_deployments.items():
model_file_id_mapping[model_id] = file_id
model_file_id_mapping: Dict[str, str] = {}
deployments_list: List[Dict] = (
healthy_deployments
if isinstance(healthy_deployments, list)
else [healthy_deployments]
)
for deployment, response in zip(deployments_list, responses):
model_id = deployment.get("model_info", {}).get("id")
if model_id is not None:
model_file_id_mapping[model_id] = response.id
return model_file_id_mapping

5
litellm/types/llms/openai.py
View file

@ -1482,6 +1482,7 @@ class ReasoningSummaryTextDeltaEvent(BaseLiteLLMOpenAIResponseObject):
type: Literal[ResponsesAPIStreamEvents.REASONING_SUMMARY_TEXT_DELTA]
item_id: str
output_index: int
summary_index: int = 0
delta: str
@ -1490,7 +1491,7 @@ class ReasoningSummaryTextDoneEvent(BaseLiteLLMOpenAIResponseObject):
item_id: str
output_index: int
sequence_number: int
summary_index: int
summary_index: int = 0
text: str
@ -1499,7 +1500,7 @@ class ReasoningSummaryPartDoneEvent(BaseLiteLLMOpenAIResponseObject):
item_id: str
output_index: int
sequence_number: int
summary_index: int
summary_index: int = 0
part: BaseLiteLLMOpenAIResponseObject

66
model_prices_and_context_window.json
View file

@ -33391,6 +33391,72 @@
"source": "https://console.cloud.google.com/vertex-ai/publishers/openai/model-garden/gpt-oss-120b-maas",
"supports_reasoning": true
},
"vertex_ai/xai/grok-4.1-fast-non-reasoning": {
"cache_read_input_token_cost": 5e-08,
"input_cost_per_token": 2e-07,
"litellm_provider": "vertex_ai",
"max_input_tokens": 2000000,
"max_output_tokens": 2000000,
"max_tokens": 2000000,
"mode": "chat",
"output_cost_per_token": 5e-07,
"source": "https://docs.x.ai/docs/models (Vertex AI Model Garden)",
"supports_function_calling": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true,
"supports_web_search": true
},
"vertex_ai/xai/grok-4.1-fast-reasoning": {
"cache_read_input_token_cost": 5e-08,
"input_cost_per_token": 2e-07,
"litellm_provider": "vertex_ai",
"max_input_tokens": 2000000,
"max_output_tokens": 2000000,
"max_tokens": 2000000,
"mode": "chat",
"output_cost_per_token": 5e-07,
"source": "https://docs.x.ai/docs/models (Vertex AI Model Garden)",
"supports_function_calling": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true,
"supports_web_search": true
},
"vertex_ai/xai/grok-4.20-non-reasoning": {
"cache_read_input_token_cost": 2e-07,
"input_cost_per_token": 2e-06,
"litellm_provider": "vertex_ai",
"max_input_tokens": 2000000,
"max_output_tokens": 2000000,
"max_tokens": 2000000,
"mode": "chat",
"output_cost_per_token": 6e-06,
"source": "https://docs.x.ai/docs/models (Vertex AI Model Garden)",
"supports_function_calling": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true,
"supports_web_search": true
},
"vertex_ai/xai/grok-4.20-reasoning": {
"cache_read_input_token_cost": 2e-07,
"input_cost_per_token": 2e-06,
"litellm_provider": "vertex_ai",
"max_input_tokens": 2000000,
"max_output_tokens": 2000000,
"max_tokens": 2000000,
"mode": "chat",
"output_cost_per_token": 6e-06,
"source": "https://docs.x.ai/docs/models (Vertex AI Model Garden)",
"supports_function_calling": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true,
"supports_web_search": true
},
"vertex_ai/qwen/qwen3-235b-a22b-instruct-2507-maas": {
"input_cost_per_token": 2.5e-07,
"litellm_provider": "vertex_ai-qwen_models",

2
pyproject.toml
View file

@ -52,7 +52,7 @@ proxy = [
"azure-identity==1.25.2",
"azure-storage-blob==12.28.0",
"mcp==1.26.0",
"litellm-proxy-extras==0.4.69",
"litellm-proxy-extras==0.4.70",
"litellm-enterprise==0.1.39",
"RestrictedPython==8.1",
"rich==13.9.4",

894
tests/llm_responses_api_testing/test_responses_hooks.py
View file

@ -1,6 +1,9 @@
import asyncio
from contextlib import suppress
from datetime import datetime
import json
from types import SimpleNamespace
from unittest.mock import AsyncMock, MagicMock
import httpx
import pytest
@ -8,8 +11,17 @@ import pytest
import litellm
from litellm.integrations.custom_logger import CustomLogger
from litellm.responses import streaming_iterator as streaming_module
from litellm.responses.streaming_iterator import ResponsesAPIStreamingIterator
from litellm.types.llms.openai import ResponsesAPIStreamEvents
from litellm.responses.streaming_iterator import (
CachedResponsesAPIStreamingIterator,
MockResponsesAPIStreamingIterator,
ResponsesAPIStreamingIterator,
SyncResponsesAPIStreamingIterator,
)
from litellm.types.llms.openai import (
ResponseCompletedEvent,
ResponsesAPIResponse,
ResponsesAPIStreamEvents,
)
from litellm.types.utils import CallTypes
@ -19,15 +31,19 @@ class _FakeLoggingObj:
self.async_success_calls = 0
self.failure_calls = 0
self.async_failure_calls = 0
self.last_success_kwargs = None
self.last_async_success_kwargs = None
self.start_time = datetime.now()
self.model_call_details = {"litellm_params": {}}
# Signature alignment with Logging handlers
def success_handler(self, *args, **kwargs):
self.success_calls += 1
self.last_success_kwargs = kwargs
async def async_success_handler(self, *args, **kwargs):
self.async_success_calls += 1
self.last_async_success_kwargs = kwargs
def failure_handler(self, *args, **kwargs):
self.failure_calls += 1
@ -36,6 +52,115 @@ class _FakeLoggingObj:
self.async_failure_calls += 1
def _make_completed_response(response_id: str = "resp_test") -> ResponseCompletedEvent:
return ResponseCompletedEvent(
type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED,
response=ResponsesAPIResponse(
id=response_id,
created_at=int(datetime.now().timestamp()),
status="completed",
model="test-model",
object="response",
output=[
{
"type": "message",
"id": f"msg_{response_id}",
"status": "completed",
"role": "assistant",
"content": [
{
"type": "output_text",
"text": "cached streamed response",
"annotations": [],
}
],
}
],
),
)
@pytest.mark.asyncio
async def test_log_background_task_failure_logs_task_exceptions(monkeypatch):
error_logger = MagicMock()
monkeypatch.setattr(streaming_module.verbose_logger, "error", error_logger)
async def _boom():
raise RuntimeError("boom")
task = asyncio.create_task(_boom())
with suppress(RuntimeError):
await task
streaming_module._log_background_task_failure(task, task_name="cache write")
error_logger.assert_called_once()
assert error_logger.call_args.args == (
"%s failed: %s",
"cache write",
task.exception(),
)
@pytest.mark.asyncio
async def test_log_background_task_failure_ignores_cancelled_tasks(monkeypatch):
error_logger = MagicMock()
monkeypatch.setattr(streaming_module.verbose_logger, "error", error_logger)
task = asyncio.create_task(asyncio.sleep(1))
task.cancel()
with suppress(asyncio.CancelledError):
await task
streaming_module._log_background_task_failure(task, task_name="cache write")
error_logger.assert_not_called()
def test_content_part_done_event_supports_refusal_and_reasoning_text():
refusal_event = streaming_module._build_content_part_done_event(
item_id="msg_1",
output_index=0,
content_index=0,
part_payload={"type": "refusal", "refusal": "no"},
)
reasoning_event = streaming_module._build_content_part_done_event(
item_id="msg_1",
output_index=0,
content_index=1,
part_payload={"type": "reasoning_text", "reasoning": "because"},
)
unsupported_event = streaming_module._build_content_part_done_event(
item_id="msg_1",
output_index=0,
content_index=2,
part_payload={"type": "image"},
)
assert refusal_event.part.type == "refusal"
assert refusal_event.part.refusal == "no"
assert reasoning_event.part.type == "reasoning_text"
assert reasoning_event.part.reasoning == "because"
assert unsupported_event is None
def test_dump_response_object_handles_model_and_unknown_values():
response = ResponsesAPIResponse(
id="resp_dump",
created_at=int(datetime.now().timestamp()),
status="completed",
model="gpt-4.1-mini",
object="response",
output=[],
)
assert streaming_module._dump_response_object(response)["id"] == "resp_dump"
assert streaming_module._dump_response_object({"type": "message"}) == {
"type": "message"
}
assert streaming_module._dump_response_object(object()) == {}
@pytest.mark.asyncio
async def test_responses_streaming_triggers_hooks(monkeypatch):
"""
@ -167,3 +292,768 @@ async def test_responses_streaming_failure_triggers_failure_handlers():
await asyncio.sleep(0.2)
assert logging_obj.failure_calls >= 1
assert logging_obj.async_failure_calls >= 1
def test_process_chunk_requires_provider_config():
iterator = ResponsesAPIStreamingIterator(
response=httpx.Response(200),
model="test-model",
responses_api_provider_config=None,
logging_obj=_FakeLoggingObj(),
request_data={"foo": "bar"},
call_type=CallTypes.responses.value,
)
with pytest.raises(ValueError, match="responses_api_provider_config is required"):
iterator._process_chunk(json.dumps({"type": "response.completed"}))
def test_process_chunk_wraps_encrypted_content_with_model_id():
openai_types = streaming_module._get_openai_response_types()
class _EncryptedConfig:
def transform_streaming_response(self, **kwargs):
return openai_types.OutputItemAddedEvent(
type=openai_types.ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED,
output_index=0,
item=openai_types.BaseLiteLLMOpenAIResponseObject(
id="rs_123",
type="reasoning",
encrypted_content="ciphertext",
),
)
iterator = ResponsesAPIStreamingIterator(
response=httpx.Response(200),
model="test-model",
responses_api_provider_config=_EncryptedConfig(),
logging_obj=_FakeLoggingObj(),
litellm_metadata={
"encrypted_content_affinity_enabled": True,
"model_info": {"id": "model-123"},
},
request_data={"foo": "bar"},
call_type=CallTypes.responses.value,
)
event = iterator._process_chunk(json.dumps({"type": "response.output_item.added"}))
assert event.item.encrypted_content.startswith("litellm_enc:")
assert event.item.encrypted_content.endswith(";ciphertext")
def test_process_chunk_completed_response_updates_id_and_usage_cost(monkeypatch):
original_include_cost = litellm.include_cost_in_streaming_usage
litellm.include_cost_in_streaming_usage = True
openai_types = streaming_module._get_openai_response_types()
class _CompletedConfig:
def transform_streaming_response(self, **kwargs):
return openai_types.ResponseCompletedEvent(
type=openai_types.ResponsesAPIStreamEvents.RESPONSE_COMPLETED,
response=ResponsesAPIResponse(
id="resp_live",
created_at=int(datetime.now().timestamp()),
status="completed",
model="test-model",
object="response",
output=[],
usage=openai_types.ResponseAPIUsage(
input_tokens=1,
output_tokens=2,
total_tokens=3,
),
),
)
logging_obj = _FakeLoggingObj()
logging_obj._response_cost_calculator = MagicMock(return_value=1.23)
iterator = ResponsesAPIStreamingIterator(
response=httpx.Response(200),
model="test-model",
responses_api_provider_config=_CompletedConfig(),
logging_obj=logging_obj,
litellm_metadata={"model_info": {"id": "model-123"}},
custom_llm_provider="openai",
request_data={"foo": "bar"},
call_type=CallTypes.responses.value,
)
completion_handler = MagicMock()
monkeypatch.setattr(
iterator, "_handle_logging_completed_response", completion_handler
)
try:
# Chunk must include a top-level "response" key so BaseResponsesAPIStreamingIterator
# runs _update_responses_api_response_id_with_model_id (see streaming_iterator.py).
event = iterator._process_chunk(
json.dumps(
{"type": "response.completed", "response": {"id": "resp_live"}}
)
)
finally:
litellm.include_cost_in_streaming_usage = original_include_cost
assert iterator.completed_response is event
assert event.response.id != "resp_live"
assert event.response.id.startswith("resp_")
assert event.response.usage.cost == 1.23
completion_handler.assert_called_once()
def test_process_chunk_failed_response_triggers_failure_logging(monkeypatch):
openai_types = streaming_module._get_openai_response_types()
class _FailedConfig:
def transform_streaming_response(self, **kwargs):
return openai_types.ResponseFailedEvent(
type=openai_types.ResponsesAPIStreamEvents.RESPONSE_FAILED,
response=ResponsesAPIResponse(
id="resp_failed",
created_at=int(datetime.now().timestamp()),
status="failed",
model="test-model",
object="response",
output=[],
error={"message": "provider failed"},
),
)
iterator = ResponsesAPIStreamingIterator(
response=httpx.Response(200),
model="test-model",
responses_api_provider_config=_FailedConfig(),
logging_obj=_FakeLoggingObj(),
request_data={"foo": "bar"},
call_type=CallTypes.responses.value,
)
failure_handler = MagicMock()
monkeypatch.setattr(iterator, "_handle_logging_failed_response", failure_handler)
event = iterator._process_chunk(json.dumps({"type": "response.failed"}))
assert iterator.completed_response is event
failure_handler.assert_called_once()
@pytest.mark.asyncio
async def test_handle_logging_failed_response_uses_response_error_message():
openai_types = streaming_module._get_openai_response_types()
logging_obj = _FakeLoggingObj()
iterator = ResponsesAPIStreamingIterator(
response=httpx.Response(200),
model="test-model",
responses_api_provider_config=SimpleNamespace(),
logging_obj=logging_obj,
request_data={"foo": "bar"},
call_type=CallTypes.responses.value,
)
iterator.completed_response = openai_types.ResponseFailedEvent(
type=openai_types.ResponsesAPIStreamEvents.RESPONSE_FAILED,
response=ResponsesAPIResponse(
id="resp_failed_real",
created_at=int(datetime.now().timestamp()),
status="failed",
model="test-model",
object="response",
output=[],
error={"message": "provider failed"},
),
)
iterator._handle_logging_failed_response()
await asyncio.sleep(0.2)
assert logging_obj.failure_calls == 1
assert logging_obj.async_failure_calls == 1
def test_process_chunk_returns_none_for_invalid_json_and_non_dict_payload():
class _NoopConfig:
def transform_streaming_response(self, **kwargs):
raise AssertionError("should not be called")
iterator = ResponsesAPIStreamingIterator(
response=httpx.Response(200),
model="test-model",
responses_api_provider_config=_NoopConfig(),
logging_obj=_FakeLoggingObj(),
request_data={"foo": "bar"},
call_type=CallTypes.responses.value,
)
assert iterator._process_chunk("not-json") is None
assert iterator._process_chunk(json.dumps(["not", "a", "dict"])) is None
def test_process_chunk_cost_annotation_failure_is_nonfatal(monkeypatch):
original_include_cost = litellm.include_cost_in_streaming_usage
litellm.include_cost_in_streaming_usage = True
openai_types = streaming_module._get_openai_response_types()
class _CompletedConfig:
def transform_streaming_response(self, **kwargs):
return openai_types.ResponseCompletedEvent(
type=openai_types.ResponsesAPIStreamEvents.RESPONSE_COMPLETED,
response=ResponsesAPIResponse(
id="resp_cost_failure",
created_at=int(datetime.now().timestamp()),
status="completed",
model="test-model",
object="response",
output=[],
usage=openai_types.ResponseAPIUsage(
input_tokens=1,
output_tokens=2,
total_tokens=3,
),
),
)
logging_obj = _FakeLoggingObj()
logging_obj._response_cost_calculator = MagicMock(side_effect=RuntimeError("boom"))
iterator = ResponsesAPIStreamingIterator(
response=httpx.Response(200),
model="test-model",
responses_api_provider_config=_CompletedConfig(),
logging_obj=logging_obj,
request_data={"foo": "bar"},
call_type=CallTypes.responses.value,
)
completion_handler = MagicMock()
monkeypatch.setattr(
iterator, "_handle_logging_completed_response", completion_handler
)
try:
event = iterator._process_chunk(json.dumps({"type": "response.completed"}))
finally:
litellm.include_cost_in_streaming_usage = original_include_cost
assert iterator.completed_response is event
assert event.response.usage.cost is None
completion_handler.assert_called_once()
def test_get_completed_response_object_accepts_direct_response():
logging_obj = _FakeLoggingObj()
iterator = SyncResponsesAPIStreamingIterator(
response=httpx.Response(200),
model="test-model",
responses_api_provider_config=SimpleNamespace(),
logging_obj=logging_obj,
request_data={"foo": "bar"},
call_type=CallTypes.responses.value,
)
direct_response = _make_completed_response("resp_direct").response
iterator.completed_response = direct_response
assert iterator._get_completed_response_object() is direct_response
@pytest.mark.asyncio
async def test_responses_streaming_completed_event_persists_async_cache():
logging_obj = _FakeLoggingObj()
original_cache = litellm.cache
litellm.cache = SimpleNamespace(
async_add_cache=AsyncMock(),
add_cache=MagicMock(),
)
caching_handler = SimpleNamespace(
request_kwargs={
"model": "test-model",
"input": "hello",
"stream": True,
"caching": True,
"cache_key": "stale-request-cache-key",
"metadata": None,
"custom_llm_provider": "openai",
},
preset_cache_key="responses-stream-cache-key",
original_function=litellm.aresponses,
async_set_cache=AsyncMock(),
_should_store_result_in_cache=lambda original_function, kwargs: True,
)
logging_obj._llm_caching_handler = caching_handler
iterator = ResponsesAPIStreamingIterator(
response=httpx.Response(200),
model="test-model",
responses_api_provider_config=SimpleNamespace(),
logging_obj=logging_obj,
request_data=caching_handler.request_kwargs,
call_type=CallTypes.aresponses.value,
)
iterator.completed_response = _make_completed_response()
iterator._handle_logging_completed_response()
await asyncio.sleep(0.2)
litellm.cache.async_add_cache.assert_called_once()
assert litellm.cache.async_add_cache.call_args.kwargs["stream"] is True
assert (
litellm.cache.async_add_cache.call_args.kwargs["cache_key"]
== "responses-stream-cache-key"
)
assert "metadata" not in litellm.cache.async_add_cache.call_args.kwargs
assert "custom_llm_provider" not in litellm.cache.async_add_cache.call_args.kwargs
assert (
json.loads(litellm.cache.async_add_cache.call_args.args[0])["id"]
== iterator.completed_response.response.id
)
litellm.cache = original_cache
def test_responses_streaming_completed_event_persists_sync_cache():
logging_obj = _FakeLoggingObj()
original_cache = litellm.cache
litellm.cache = SimpleNamespace(
async_add_cache=AsyncMock(),
add_cache=MagicMock(),
)
caching_handler = SimpleNamespace(
request_kwargs={
"model": "test-model",
"input": "hello",
"stream": True,
"caching": True,
"cache_key": "stale-request-cache-key",
"metadata": None,
"custom_llm_provider": "openai",
},
preset_cache_key="responses-stream-cache-key",
original_function=litellm.responses,
sync_set_cache=MagicMock(),
_should_store_result_in_cache=lambda original_function, kwargs: True,
)
logging_obj._llm_caching_handler = caching_handler
iterator = SyncResponsesAPIStreamingIterator(
response=httpx.Response(200),
model="test-model",
responses_api_provider_config=SimpleNamespace(),
logging_obj=logging_obj,
request_data=caching_handler.request_kwargs,
call_type=CallTypes.responses.value,
)
iterator.completed_response = _make_completed_response("resp_sync")
iterator._handle_logging_completed_response()
litellm.cache.add_cache.assert_called_once()
assert litellm.cache.add_cache.call_args.kwargs["stream"] is True
assert (
litellm.cache.add_cache.call_args.kwargs["cache_key"]
== "responses-stream-cache-key"
)
assert "metadata" not in litellm.cache.add_cache.call_args.kwargs
assert "custom_llm_provider" not in litellm.cache.add_cache.call_args.kwargs
assert (
json.loads(litellm.cache.add_cache.call_args.args[0])["id"]
== iterator.completed_response.response.id
)
litellm.cache = original_cache
def test_log_completed_response_sync_direct_path(monkeypatch):
hook_calls = {"post_call": 0, "metadata": 0}
async def fake_post_call(request_data, response, call_type):
hook_calls["post_call"] += 1
def fake_update_metadata(**kwargs):
hook_calls["metadata"] += 1
monkeypatch.setattr(
streaming_module,
"async_post_call_success_deployment_hook",
fake_post_call,
)
monkeypatch.setattr(
streaming_module,
"update_response_metadata",
fake_update_metadata,
)
logging_obj = _FakeLoggingObj()
iterator = SyncResponsesAPIStreamingIterator(
response=httpx.Response(200),
model="test-model",
responses_api_provider_config=SimpleNamespace(),
logging_obj=logging_obj,
request_data={"foo": "bar"},
call_type=CallTypes.responses.value,
)
iterator._persist_completed_response_before_logging = False
iterator.completed_response = _make_completed_response("resp_log_sync")
iterator._log_completed_response(is_async=False)
asyncio.run(asyncio.sleep(0.2))
assert logging_obj.success_calls == 1
assert logging_obj.async_success_calls == 1
assert hook_calls["post_call"] == 1
assert hook_calls["metadata"] == 1
def test_log_completed_response_falls_back_when_model_validate_fails(monkeypatch):
class _BadSerializableResponse:
@classmethod
def model_validate(cls, value):
raise RuntimeError("nope")
def model_dump(self):
return {"id": "bad"}
logging_obj = _FakeLoggingObj()
iterator = SyncResponsesAPIStreamingIterator(
response=httpx.Response(200),
model="test-model",
responses_api_provider_config=SimpleNamespace(),
logging_obj=logging_obj,
request_data={"foo": "bar"},
call_type=CallTypes.responses.value,
)
iterator._persist_completed_response_before_logging = False
iterator.completed_response = _BadSerializableResponse()
monkeypatch.setattr(iterator, "_run_post_success_hooks", MagicMock())
iterator._log_completed_response(is_async=False)
asyncio.run(asyncio.sleep(0.2))
assert logging_obj.success_calls == 1
assert logging_obj.async_success_calls == 1
@pytest.mark.parametrize(
"scenario",
[
"already_cached",
"not_completed",
"missing_caching_handler",
"not_streaming",
"store_disabled",
"missing_cache_backend",
],
)
def test_persist_completed_response_to_cache_guard_branches(monkeypatch, scenario):
logging_obj = _FakeLoggingObj()
iterator = SyncResponsesAPIStreamingIterator(
response=httpx.Response(200),
model="test-model",
responses_api_provider_config=SimpleNamespace(),
logging_obj=logging_obj,
request_data={"foo": "bar"},
call_type=CallTypes.responses.value,
)
openai_types = streaming_module._get_openai_response_types()
completed_event = _make_completed_response("resp_guard")
iterator.completed_response = completed_event
if scenario == "already_cached":
iterator._completed_response_cached = True
elif scenario == "not_completed":
iterator.completed_response = openai_types.ResponseIncompleteEvent(
type=openai_types.ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE,
response=completed_event.response,
)
elif scenario == "missing_caching_handler":
logging_obj._llm_caching_handler = None
else:
logging_obj._llm_caching_handler = SimpleNamespace(
request_kwargs={
"model": "test-model",
"input": "hello",
"stream": scenario != "not_streaming",
"cache_key": "request-cache-key",
"metadata": None,
"custom_llm_provider": "openai",
},
preset_cache_key=None,
original_function=litellm.responses,
dual_cache=None,
_should_store_result_in_cache=lambda original_function, kwargs: (
scenario != "store_disabled"
),
)
if scenario == "missing_cache_backend":
monkeypatch.setattr(streaming_module.litellm, "cache", None)
else:
monkeypatch.setattr(
streaming_module.litellm,
"cache",
SimpleNamespace(add_cache=MagicMock(), async_add_cache=AsyncMock()),
)
iterator._persist_completed_response_to_cache(is_async=False)
expected_cached_flag = scenario == "already_cached"
assert iterator._completed_response_cached is expected_cached_flag
def test_build_synthetic_response_events_covers_annotations_function_calls_and_refusals():
original_include_cost = litellm.include_cost_in_streaming_usage
litellm.include_cost_in_streaming_usage = True
logging_obj = _FakeLoggingObj()
logging_obj._response_cost_calculator = MagicMock(side_effect=RuntimeError("boom"))
transformed = ResponsesAPIResponse(
id="resp_events",
created_at=int(datetime.now().timestamp()),
status="completed",
model="gpt-4.1-mini",
object="response",
output=[
{
"type": "message",
"id": "msg_events",
"status": "completed",
"role": "assistant",
"content": [
{
"type": "output_text",
"text": "hello world",
"annotations": [{"type": "file_citation", "file_id": "file_1"}],
},
{
"type": "refusal",
"refusal": "no thanks",
},
],
},
{
"type": "function_call",
"id": "fc_events",
"call_id": "call_123",
"name": "lookup",
"arguments": '{"id":1}',
},
],
)
try:
events = streaming_module._build_synthetic_response_events(
transformed=transformed,
logging_obj=logging_obj,
chunk_size=5,
)
finally:
litellm.include_cost_in_streaming_usage = original_include_cost
event_types = [
event.type.value if hasattr(event.type, "value") else str(event.type)
for event in events
]
assert "response.output_text.annotation.added" in event_types
assert "response.refusal.delta" in event_types
assert "response.refusal.done" in event_types
assert "response.function_call_arguments.delta" in event_types
assert "response.function_call_arguments.done" in event_types
assert event_types[-1] == "response.completed"
@pytest.mark.asyncio
async def test_mock_responses_streaming_iterator_async_iteration_logs_completion(
monkeypatch,
):
hook_calls = {"post_call": 0, "metadata": 0}
async def fake_post_call(request_data, response, call_type):
hook_calls["post_call"] += 1
def fake_update_metadata(**kwargs):
hook_calls["metadata"] += 1
monkeypatch.setattr(
streaming_module,
"async_post_call_success_deployment_hook",
fake_post_call,
)
monkeypatch.setattr(
streaming_module,
"update_response_metadata",
fake_update_metadata,
)
class _MockTransformConfig:
def transform_response_api_response(self, **kwargs):
return _make_completed_response("resp_mock").response
logging_obj = _FakeLoggingObj()
iterator = MockResponsesAPIStreamingIterator(
response=httpx.Response(200),
model="test-model",
responses_api_provider_config=_MockTransformConfig(),
logging_obj=logging_obj,
request_data={"model": "test-model", "stream": True},
call_type=CallTypes.responses.value,
)
streamed_events = [event async for event in iterator]
await asyncio.sleep(0.2)
assert streamed_events[0].type == ResponsesAPIStreamEvents.RESPONSE_CREATED
assert streamed_events[-1].type == ResponsesAPIStreamEvents.RESPONSE_COMPLETED
assert logging_obj.success_calls == 1
assert logging_obj.async_success_calls == 1
assert hook_calls["post_call"] == 1
assert hook_calls["metadata"] == 1
def test_mock_responses_streaming_iterator_sync_iteration_logs_completion(monkeypatch):
hook_calls = {"post_call": 0, "metadata": 0}
async def fake_post_call(request_data, response, call_type):
hook_calls["post_call"] += 1
def fake_update_metadata(**kwargs):
hook_calls["metadata"] += 1
monkeypatch.setattr(
streaming_module,
"async_post_call_success_deployment_hook",
fake_post_call,
)
monkeypatch.setattr(
streaming_module,
"update_response_metadata",
fake_update_metadata,
)
class _MockTransformConfig:
def transform_response_api_response(self, **kwargs):
return _make_completed_response("resp_mock_sync").response
logging_obj = _FakeLoggingObj()
iterator = MockResponsesAPIStreamingIterator(
response=httpx.Response(200),
model="test-model",
responses_api_provider_config=_MockTransformConfig(),
logging_obj=logging_obj,
request_data={"model": "test-model", "stream": True},
call_type=CallTypes.responses.value,
)
streamed_events = list(iterator)
asyncio.run(asyncio.sleep(0.2))
assert streamed_events[0].type == ResponsesAPIStreamEvents.RESPONSE_CREATED
assert streamed_events[-1].type == ResponsesAPIStreamEvents.RESPONSE_COMPLETED
assert logging_obj.success_calls == 1
assert logging_obj.async_success_calls == 1
assert hook_calls["post_call"] == 1
assert hook_calls["metadata"] == 1
@pytest.mark.asyncio
async def test_cached_responses_stream_async_hit_triggers_success_callbacks(
monkeypatch,
):
hook_calls = {"post_call": 0, "metadata": 0}
async def fake_post_call(request_data, response, call_type):
hook_calls["post_call"] += 1
def fake_update_metadata(**kwargs):
hook_calls["metadata"] += 1
monkeypatch.setattr(
streaming_module,
"async_post_call_success_deployment_hook",
fake_post_call,
)
monkeypatch.setattr(
streaming_module,
"update_response_metadata",
fake_update_metadata,
)
logging_obj = _FakeLoggingObj()
original_cache = litellm.cache
litellm.cache = SimpleNamespace(
async_add_cache=AsyncMock(),
add_cache=MagicMock(),
)
logging_obj._llm_caching_handler = SimpleNamespace(
request_kwargs={"model": "test-model", "input": "hello", "stream": True},
preset_cache_key="responses-stream-cache-key",
original_function=litellm.aresponses,
_should_store_result_in_cache=lambda original_function, kwargs: True,
)
iterator = CachedResponsesAPIStreamingIterator(
response=_make_completed_response("resp_cached_async").response,
logging_obj=logging_obj,
request_data={"model": "test-model", "input": "hello", "stream": True},
call_type=CallTypes.aresponses.value,
)
streamed_events = [event async for event in iterator]
await asyncio.sleep(0.2)
assert streamed_events[-1].type == ResponsesAPIStreamEvents.RESPONSE_COMPLETED
assert logging_obj.success_calls == 1
assert logging_obj.async_success_calls == 1
assert logging_obj.last_success_kwargs["cache_hit"] is True
assert logging_obj.last_async_success_kwargs["cache_hit"] is True
assert hook_calls["post_call"] == 1
assert hook_calls["metadata"] == 1
litellm.cache.async_add_cache.assert_not_called()
litellm.cache.add_cache.assert_not_called()
litellm.cache = original_cache
def test_cached_responses_stream_sync_hit_triggers_success_callbacks(monkeypatch):
hook_calls = {"post_call": 0, "metadata": 0}
async def fake_post_call(request_data, response, call_type):
hook_calls["post_call"] += 1
def fake_update_metadata(**kwargs):
hook_calls["metadata"] += 1
monkeypatch.setattr(
streaming_module,
"async_post_call_success_deployment_hook",
fake_post_call,
)
monkeypatch.setattr(
streaming_module,
"update_response_metadata",
fake_update_metadata,
)
logging_obj = _FakeLoggingObj()
original_cache = litellm.cache
litellm.cache = SimpleNamespace(
async_add_cache=AsyncMock(),
add_cache=MagicMock(),
)
logging_obj._llm_caching_handler = SimpleNamespace(
request_kwargs={"model": "test-model", "input": "hello", "stream": True},
preset_cache_key="responses-stream-cache-key",
original_function=litellm.responses,
_should_store_result_in_cache=lambda original_function, kwargs: True,
)
iterator = CachedResponsesAPIStreamingIterator(
response=_make_completed_response("resp_cached_sync").response,
logging_obj=logging_obj,
request_data={"model": "test-model", "input": "hello", "stream": True},
call_type=CallTypes.responses.value,
)
streamed_events = list(iterator)
asyncio.run(asyncio.sleep(0.2))
assert streamed_events[-1].type == ResponsesAPIStreamEvents.RESPONSE_COMPLETED
assert logging_obj.success_calls == 1
assert logging_obj.async_success_calls == 1
assert logging_obj.last_success_kwargs["cache_hit"] is True
assert logging_obj.last_async_success_kwargs["cache_hit"] is True
assert hook_calls["post_call"] == 1
assert hook_calls["metadata"] == 1
litellm.cache.async_add_cache.assert_not_called()
litellm.cache.add_cache.assert_not_called()
litellm.cache = original_cache

2
tests/llm_translation/test_anthropic_completion.py
View file

@ -870,7 +870,7 @@ from litellm.constants import RESPONSE_FORMAT_TOOL_NAME
def test_anthropic_json_mode_and_tool_call_response(
json_mode, tool_calls, expect_null_response
):
result = litellm.AnthropicConfig()._transform_response_for_json_mode(
result, _, _ = litellm.AnthropicConfig()._resolve_json_mode_non_streaming(
json_mode=json_mode,
tool_calls=tool_calls,
)

389
tests/local_testing/test_caching_handler.py
View file

@ -19,9 +19,14 @@ import pytest
import litellm
from litellm import aembedding, completion, embedding, aresponses, responses
from litellm.caching.caching import Cache
from litellm.responses.streaming_iterator import CachedResponsesAPIStreamingIterator
from unittest.mock import AsyncMock, patch, MagicMock
from litellm.caching.caching_handler import LLMCachingHandler, CachingHandlerResponse
from litellm.caching.caching_handler import (
LLMCachingHandler,
CachingHandlerResponse,
_should_defer_streaming_cache_hit_callbacks,
)
from litellm.caching.caching import LiteLLMCacheType
from litellm.types.utils import CallTypes
from litellm.types.rerank import RerankResponse
@ -627,6 +632,55 @@ async def test_async_responses_api_caching():
assert cached_response.cached_result._hidden_params["cache_hit"] == True
@pytest.mark.asyncio
async def test_async_get_cache_updates_request_kwargs_for_streaming_responses():
"""
Ensure streamed responses retain the normalized lookup kwargs so a later
cache write can reuse the exact cache key from the read path.
"""
setup_cache()
caching_handler = LLMCachingHandler(
original_function=aresponses,
request_kwargs={"stale": True},
start_time=datetime.now(),
)
logging_obj = LiteLLMLogging(
litellm_call_id=str(datetime.now()),
call_type=CallTypes.aresponses.value,
model="gpt-4o",
messages=[],
function_id=str(uuid.uuid4()),
stream=True,
start_time=datetime.now(),
)
kwargs = {
"model": "gpt-4o",
"input": "hello",
"stream": True,
"caching": True,
}
await caching_handler._async_get_cache(
model="gpt-4o",
original_function=aresponses,
logging_obj=logging_obj,
start_time=datetime.now(),
call_type=CallTypes.aresponses.value,
kwargs=kwargs,
)
assert "stale" not in caching_handler.request_kwargs
assert caching_handler.request_kwargs["model"] == "gpt-4o"
assert caching_handler.request_kwargs["input"] == "hello"
assert caching_handler.request_kwargs["stream"] is True
assert caching_handler.request_kwargs["cache_key"] == litellm.cache.get_cache_key(
**caching_handler.request_kwargs
)
def test_sync_responses_api_caching():
"""
Test that synchronous responses API calls are properly cached and retrieved.
@ -769,6 +823,339 @@ def test_convert_cached_responses_api_result_to_model_response():
assert len(result.output) == 1
def test_sync_get_cache_does_not_eagerly_log_streaming_responses_hits():
litellm.set_verbose = True
setup_cache()
caching_handler = LLMCachingHandler(
original_function=responses, request_kwargs={}, start_time=datetime.now()
)
original_model = "gpt-4o"
responses_api_response = ResponsesAPIResponse(
id="resp_stream_sync_hit",
created_at=int(time.time()),
status="completed",
model=original_model,
object="response",
output=[
{
"type": "message",
"id": "msg_stream_sync_hit",
"status": "completed",
"role": "assistant",
"content": [
{
"type": "output_text",
"text": "Sync streamed cache hit response.",
"annotations": [],
}
],
}
],
)
logging_obj = LiteLLMLogging(
litellm_call_id=str(datetime.now()),
call_type=CallTypes.responses.value,
model=original_model,
messages=[],
function_id=str(uuid.uuid4()),
stream=True,
start_time=datetime.now(),
)
logging_obj.handle_sync_success_callbacks_for_async_calls = MagicMock()
kwargs = {
"model": original_model,
"input": "Tell me a cached story",
"stream": True,
"caching": True,
}
caching_handler.sync_set_cache(result=responses_api_response, kwargs=kwargs)
time.sleep(0.2)
cached_response = caching_handler._sync_get_cache(
model=original_model,
original_function=responses,
logging_obj=logging_obj,
start_time=datetime.now(),
call_type=CallTypes.responses.value,
kwargs=kwargs,
)
assert cached_response.cached_result is not None
assert isinstance(
cached_response.cached_result, CachedResponsesAPIStreamingIterator
)
logging_obj.handle_sync_success_callbacks_for_async_calls.assert_not_called()
def test_sync_get_cache_defers_streaming_completion_hit_callbacks():
litellm.set_verbose = True
setup_cache()
caching_handler = LLMCachingHandler(
original_function=completion, request_kwargs={}, start_time=datetime.now()
)
original_model = "gpt-4o"
logging_obj = LiteLLMLogging(
litellm_call_id=str(datetime.now()),
call_type=CallTypes.completion.value,
model=original_model,
messages=[],
function_id=str(uuid.uuid4()),
stream=True,
start_time=datetime.now(),
)
logging_obj.handle_sync_success_callbacks_for_async_calls = MagicMock()
kwargs = {
"model": original_model,
"messages": [{"role": "user", "content": "Tell me a cached joke"}],
"stream": True,
"caching": True,
}
caching_handler.sync_set_cache(result=chat_completion_response, kwargs=kwargs)
time.sleep(0.2)
cached_response = caching_handler._sync_get_cache(
model=original_model,
original_function=completion,
logging_obj=logging_obj,
start_time=datetime.now(),
call_type=CallTypes.completion.value,
kwargs=kwargs,
)
assert cached_response.cached_result is not None
logging_obj.handle_sync_success_callbacks_for_async_calls.assert_not_called()
def test_should_defer_streaming_cache_hit_callbacks_for_any_streaming_request():
assert (
_should_defer_streaming_cache_hit_callbacks(
kwargs={"stream": True},
)
is True
)
assert (
_should_defer_streaming_cache_hit_callbacks(
kwargs={"stream": False},
)
is False
)
assert (
_should_defer_streaming_cache_hit_callbacks(
kwargs={},
)
is False
)
@pytest.mark.asyncio
async def test_async_get_cache_defers_streaming_completion_hit_callbacks():
litellm.set_verbose = True
setup_cache()
caching_handler = LLMCachingHandler(
original_function=completion, request_kwargs={}, start_time=datetime.now()
)
original_model = "gpt-4o"
kwargs = {
"model": original_model,
"messages": [{"role": "user", "content": "Tell me a cached joke"}],
"stream": True,
"caching": True,
}
await caching_handler.async_set_cache(
result=chat_completion_response,
original_function=litellm.acompletion,
kwargs=kwargs,
)
await asyncio.sleep(0.2)
logging_obj = LiteLLMLogging(
litellm_call_id=str(datetime.now()),
call_type=CallTypes.acompletion.value,
model=original_model,
messages=[],
function_id=str(uuid.uuid4()),
stream=True,
start_time=datetime.now(),
)
caching_handler._async_log_cache_hit_on_callbacks = MagicMock()
cached_response = await caching_handler._async_get_cache(
model=original_model,
original_function=litellm.acompletion,
logging_obj=logging_obj,
start_time=datetime.now(),
call_type=CallTypes.acompletion.value,
kwargs=kwargs,
)
assert cached_response is not None
assert cached_response.cached_result is not None
caching_handler._async_log_cache_hit_on_callbacks.assert_not_called()
def test_convert_cached_streaming_responses_result_to_iterator():
"""
Test that cached streaming Responses results are replayed through a synthetic
streaming iterator instead of being returned as a full response object.
"""
caching_handler = LLMCachingHandler(
original_function=responses, request_kwargs={}, start_time=datetime.now()
)
logging_obj = LiteLLMLogging(
litellm_call_id=str(datetime.now()),
call_type=CallTypes.responses.value,
model="gpt-4o",
messages=[],
function_id=str(uuid.uuid4()),
stream=True,
start_time=datetime.now(),
)
cached_result = {
"id": "resp_stream_cache_test",
"created_at": int(time.time()),
"status": "completed",
"model": "gpt-4o",
"object": "response",
"output": [
{
"type": "message",
"id": "msg_stream_cache_test",
"status": "completed",
"role": "assistant",
"content": [
{
"type": "output_text",
"text": "Streaming cache replay test.",
"annotations": [],
}
],
}
],
}
result = caching_handler._convert_cached_result_to_model_response(
cached_result=cached_result,
call_type=CallTypes.responses.value,
kwargs={"model": "gpt-4o", "input": "test", "stream": True},
logging_obj=logging_obj,
model="gpt-4o",
args=(),
)
assert isinstance(result, CachedResponsesAPIStreamingIterator)
assert result.completed_response is not None
assert result.completed_response.response.id == cached_result["id"]
streamed_events = list(result)
assert streamed_events[0].type == "response.created"
assert streamed_events[1].type == "response.in_progress"
assert streamed_events[2].type == "response.output_item.added"
assert streamed_events[3].type == "response.content_part.added"
assert streamed_events[-4].type == "response.output_text.done"
assert streamed_events[-3].type == "response.content_part.done"
assert streamed_events[-2].type == "response.output_item.done"
assert streamed_events[-1].type == "response.completed"
assert streamed_events[-1].response.id == cached_result["id"]
assert streamed_events[-1].response.output[0].content[0].text == (
"Streaming cache replay test."
)
def test_convert_cached_streaming_reasoning_result_to_iterator():
caching_handler = LLMCachingHandler(
original_function=responses, request_kwargs={}, start_time=datetime.now()
)
logging_obj = LiteLLMLogging(
litellm_call_id=str(datetime.now()),
call_type=CallTypes.responses.value,
model="gpt-4o",
messages=[],
function_id=str(uuid.uuid4()),
stream=True,
start_time=datetime.now(),
)
cached_result = {
"id": "resp_stream_reasoning_cache_test",
"created_at": int(time.time()),
"status": "completed",
"model": "gpt-4o",
"object": "response",
"output": [
{
"type": "reasoning",
"id": "rs_stream_cache_test",
"summary": [
{
"type": "summary_text",
"text": "Cached reasoning summary.",
}
],
}
],
}
result = caching_handler._convert_cached_result_to_model_response(
cached_result=cached_result,
call_type=CallTypes.responses.value,
kwargs={"model": "gpt-4o", "input": "test", "stream": True},
logging_obj=logging_obj,
model="gpt-4o",
args=(),
)
assert isinstance(result, CachedResponsesAPIStreamingIterator)
streamed_events = list(result)
streamed_event_types = [
event.type.value if hasattr(event.type, "value") else str(event.type)
for event in streamed_events
]
assert streamed_event_types[:3] == [
"response.created",
"response.in_progress",
"response.output_item.added",
]
assert streamed_event_types[-4:] == [
"response.reasoning_summary_text.done",
"response.reasoning_summary_part.done",
"response.output_item.done",
"response.completed",
]
assert streamed_event_types.count("response.reasoning_summary_text.delta") >= 1
delta_events = [
event
for event in streamed_events
if (event.type.value if hasattr(event.type, "value") else str(event.type))
== "response.reasoning_summary_text.delta"
]
text_done_event = streamed_events[-4]
part_done_event = streamed_events[-3]
output_item_done_event = streamed_events[-2]
assert all(delta_event.summary_index == 0 for delta_event in delta_events)
assert text_done_event.text == "Cached reasoning summary."
assert text_done_event.summary_index == 0
assert part_done_event.part.type == "summary_text"
assert part_done_event.part.text == "Cached reasoning summary."
assert output_item_done_event.item.type == "reasoning"
assert output_item_done_event.item.summary[0]["text"] == "Cached reasoning summary."
@pytest.mark.asyncio
async def test_responses_api_cache_with_different_inputs():
"""

141
tests/local_testing/test_responses_stream_cache_keys.py Normal file
View file

@ -0,0 +1,141 @@
from datetime import datetime
from unittest.mock import AsyncMock, MagicMock
import pytest
import litellm
from litellm import aresponses
from litellm._uuid import uuid
from litellm.caching.caching_handler import LLMCachingHandler
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
from litellm.types.llms import openai as openai_types
from litellm.types.utils import CallTypes
@pytest.mark.asyncio
async def test_async_get_cache_reuses_preset_cache_key_for_responses():
caching_handler = LLMCachingHandler(
original_function=aresponses,
request_kwargs={},
start_time=datetime.now(),
)
logging_obj = LiteLLMLogging(
litellm_call_id=str(datetime.now()),
call_type=CallTypes.aresponses.value,
model="gpt-4.1-mini",
messages=[],
function_id=str(uuid.uuid4()),
stream=True,
start_time=datetime.now(),
)
original_cache = litellm.cache
mock_cache = MagicMock()
mock_cache.supported_call_types = [CallTypes.aresponses.value]
mock_cache._supports_async.return_value = True
mock_cache.get_cache_key.return_value = "responses-stream-cache-key"
mock_cache.async_get_cache = AsyncMock(return_value=None)
litellm.cache = mock_cache
kwargs = {
"model": "gpt-4.1-mini",
"input": "hello",
"stream": True,
"litellm_params": {},
}
await caching_handler._async_get_cache(
model="gpt-4.1-mini",
original_function=aresponses,
logging_obj=logging_obj,
start_time=datetime.now(),
call_type=CallTypes.aresponses.value,
kwargs=kwargs,
)
assert caching_handler.preset_cache_key == "responses-stream-cache-key"
mock_cache.async_get_cache.assert_awaited_once()
assert (
mock_cache.async_get_cache.call_args.kwargs["cache_key"]
== "responses-stream-cache-key"
)
litellm.cache = original_cache
@pytest.mark.asyncio
async def test_async_get_cache_falls_back_to_sync_cache_for_responses():
caching_handler = LLMCachingHandler(
original_function=aresponses,
request_kwargs={},
start_time=datetime.now(),
)
logging_obj = LiteLLMLogging(
litellm_call_id=str(datetime.now()),
call_type=CallTypes.aresponses.value,
model="gpt-4.1-mini",
messages=[],
function_id=str(uuid.uuid4()),
stream=True,
start_time=datetime.now(),
)
original_cache = litellm.cache
mock_cache = MagicMock()
mock_cache.supported_call_types = [CallTypes.aresponses.value]
mock_cache._supports_async.return_value = False
mock_cache.get_cache_key.return_value = "responses-stream-cache-key"
mock_cache.get_cache.return_value = None
litellm.cache = mock_cache
kwargs = {
"model": "gpt-4.1-mini",
"input": "hello",
"stream": True,
"litellm_params": {},
}
await caching_handler._async_get_cache(
model="gpt-4.1-mini",
original_function=aresponses,
logging_obj=logging_obj,
start_time=datetime.now(),
call_type=CallTypes.aresponses.value,
kwargs=kwargs,
)
assert caching_handler.preset_cache_key == "responses-stream-cache-key"
mock_cache.get_cache.assert_called_once()
assert mock_cache.get_cache.call_args.kwargs["cache_key"] == (
"responses-stream-cache-key"
)
litellm.cache = original_cache
def test_reasoning_summary_events_default_summary_index():
delta_event = openai_types.ReasoningSummaryTextDeltaEvent(
type=openai_types.ResponsesAPIStreamEvents.REASONING_SUMMARY_TEXT_DELTA,
item_id="rs_1",
output_index=0,
delta="abc",
)
text_done_event = openai_types.ReasoningSummaryTextDoneEvent(
type=openai_types.ResponsesAPIStreamEvents.REASONING_SUMMARY_TEXT_DONE,
item_id="rs_1",
output_index=0,
sequence_number=1,
text="abc",
)
part_done_event = openai_types.ReasoningSummaryPartDoneEvent(
type=openai_types.ResponsesAPIStreamEvents.REASONING_SUMMARY_PART_DONE,
item_id="rs_1",
output_index=0,
sequence_number=2,
part=openai_types.BaseLiteLLMOpenAIResponseObject(
type="summary_text",
text="abc",
),
)
assert delta_event.summary_index == 0
assert text_done_event.summary_index == 0
assert part_done_event.summary_index == 0

30
tests/otel_tests/test_e2e_model_access.py
View file

@ -6,13 +6,19 @@ from httpx import AsyncClient
from typing import Any, Optional, List, Literal
# The proxy strips client-supplied `mock_response` unless the calling key or
# team has this admin-metadata flag set. See `_UNTRUSTED_ROOT_CONTROL_FIELDS`
# in litellm/proxy/litellm_pre_call_utils.py.
_ALLOW_CLIENT_MOCK_METADATA = {"allow_client_mock_response": True}
async def generate_key(
session, models: Optional[List[str]] = None, team_id: Optional[str] = None
):
"""Helper function to generate a key with specific model access controls"""
url = "http://0.0.0.0:4000/key/generate"
headers = {"Authorization": "Bearer sk-1234", "Content-Type": "application/json"}
data = {}
data: dict = {"metadata": dict(_ALLOW_CLIENT_MOCK_METADATA)}
if models is not None:
data["models"] = models
if team_id is not None:
@ -25,7 +31,7 @@ async def generate_team(session, models: Optional[List[str]] = None):
"""Helper function to generate a team with specific model access"""
url = "http://0.0.0.0:4000/team/new"
headers = {"Authorization": "Bearer sk-1234", "Content-Type": "application/json"}
data = {}
data: dict = {"metadata": dict(_ALLOW_CLIENT_MOCK_METADATA)}
if models is not None:
data["models"] = models
async with session.post(url, headers=headers, json=data) as response:
@ -111,7 +117,12 @@ async def test_model_access_update():
# Create initial key with restricted access
response = await client.post(
"/key/generate", json={"models": ["openai/gpt-4"]}, headers=headers
"/key/generate",
json={
"models": ["openai/gpt-4"],
"metadata": dict(_ALLOW_CLIENT_MOCK_METADATA),
},
headers=headers,
)
assert response.status_code == 200
key_data = response.json()
@ -214,7 +225,11 @@ async def test_team_model_access_update():
# Create initial team with restricted access
response = await client.post(
"/team/new",
json={"models": ["openai/gpt-4"], "name": "test-team"},
json={
"models": ["openai/gpt-4"],
"name": "test-team",
"metadata": dict(_ALLOW_CLIENT_MOCK_METADATA),
},
headers=headers,
)
assert response.status_code == 200
@ -223,7 +238,12 @@ async def test_team_model_access_update():
# Generate a key for this team
response = await client.post(
"/key/generate", json={"team_id": team_id}, headers=headers
"/key/generate",
json={
"team_id": team_id,
"metadata": dict(_ALLOW_CLIENT_MOCK_METADATA),
},
headers=headers,
)
assert response.status_code == 200
key = response.json()["key"]

11
tests/proxy_unit_tests/test_auth_checks.py
View file

@ -16,6 +16,7 @@ import httpx
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.auth.auth_checks import get_end_user_object
from litellm.caching.caching import DualCache
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.proxy._types import (
LiteLLM_EndUserTable,
LiteLLM_BudgetTable,
@ -48,9 +49,15 @@ async def test_get_end_user_object(customer_spend, customer_budget):
litellm_budget_table=_budget,
blocked=False,
)
_cache = DualCache()
# UserApiKeyCache applies model_type on get/set; plain DualCache returns raw dicts
# and breaks get_end_user_object's typed async_get_cache path.
_cache = UserApiKeyCache()
_key = "end_user_id:{}".format(end_user_id)
_cache.set_cache(key=_key, value=end_user_obj.model_dump())
await _cache.async_set_cache(
key=_key,
value=end_user_obj,
model_type=LiteLLM_EndUserTable,
)
try:
await get_end_user_object(
end_user_id=end_user_id,

7
tests/proxy_unit_tests/test_user_api_key_auth.py
View file

@ -268,7 +268,12 @@ async def test_aaauser_personal_budgets(key_ownership):
test_user_cache = getattr(litellm.proxy.proxy_server, "user_api_key_cache")
assert test_user_cache.get_cache(key=hash_token(user_key)) == valid_token
assert (
test_user_cache.get_cache(
key=hash_token(user_key), model_type=UserAPIKeyAuth
)
== valid_token
)
try:
await user_api_key_auth(request=request, api_key="Bearer " + user_key)

111
tests/router_unit_tests/test_router_endpoints.py
View file

@ -1110,7 +1110,7 @@ def test_initialize_skills_endpoints():
async def test_init_containers_api_endpoints():
"""
Test that _init_containers_api_endpoints calls the original function
directly without model-based routing.
directly when there is no managed container ID (no embedded model_id).
"""
router = Router(model_list=[])
@ -1127,3 +1127,112 @@ async def test_init_containers_api_endpoints():
custom_llm_provider="openai", name="Test Container"
)
assert result == mock_response
@pytest.mark.asyncio
async def test_init_containers_api_endpoints_managed_id_routes_via_generic_fallbacks():
"""
Managed ``cntr_`` IDs embed ``model_id``; router should decode and use
``_ageneric_api_call_with_fallbacks`` so deployment credentials apply.
"""
from litellm.responses.utils import ResponsesAPIRequestUtils
router = Router(
model_list=[
{
"model_name": "azure-router-model",
"litellm_params": {
"model": "azure/gpt-4",
"api_key": "fake-key",
"api_base": "https://westus.api.cognitive.microsoft.com",
},
}
]
)
router._ageneric_api_call_with_fallbacks = AsyncMock()
managed_id = ResponsesAPIRequestUtils._build_container_id(
custom_llm_provider="azure",
model_id="azure-router-model",
container_id="cfile_upstream_abc",
)
await router._init_containers_api_endpoints(
original_function=AsyncMock(),
custom_llm_provider="openai",
container_id=managed_id,
file_id="cfile_xyz",
)
router._ageneric_api_call_with_fallbacks.assert_called_once()
call_kw = router._ageneric_api_call_with_fallbacks.call_args.kwargs
assert call_kw["model"] == "azure-router-model"
assert call_kw["container_id"] == "cfile_upstream_abc"
assert call_kw["file_id"] == "cfile_xyz"
assert call_kw["custom_llm_provider"] == "azure"
@pytest.mark.asyncio
async def test_init_containers_api_endpoints_managed_id_without_model_id_unwraps():
"""
Managed ``cntr_`` IDs may be encoded with an empty ``model_id`` (e.g. when a
streaming response had no router metadata). The router must still unwrap the
managed ID before calling the upstream provider — otherwise the raw
``cntr_...`` token leaks downstream and the provider rejects it.
"""
from litellm.responses.utils import ResponsesAPIRequestUtils
router = Router(model_list=[])
mock_original_function = AsyncMock(return_value={"ok": True})
managed_id = ResponsesAPIRequestUtils._build_container_id(
custom_llm_provider="openai",
model_id=None,
container_id="cfile_upstream_abc",
)
await router._init_containers_api_endpoints(
original_function=mock_original_function,
custom_llm_provider="openai",
container_id=managed_id,
file_id="cfile_xyz",
)
mock_original_function.assert_called_once()
call_kw = mock_original_function.call_args.kwargs
assert call_kw["container_id"] == "cfile_upstream_abc"
assert call_kw["file_id"] == "cfile_xyz"
assert call_kw["custom_llm_provider"] == "openai"
@pytest.mark.asyncio
async def test_init_containers_api_endpoints_managed_id_without_model_id_applies_decoded_provider():
"""
A managed ``cntr_`` ID can encode a non-OpenAI provider (e.g. ``azure``) with
an empty ``model_id`` (streaming events without router ``model_info.id``).
The router must still apply the decoded provider so the request routes to
the correct upstream — not stay on the default ``openai``.
"""
from litellm.responses.utils import ResponsesAPIRequestUtils
router = Router(model_list=[])
mock_original_function = AsyncMock(return_value={"ok": True})
managed_id = ResponsesAPIRequestUtils._build_container_id(
custom_llm_provider="azure",
model_id=None,
container_id="cfile_upstream_abc",
)
await router._init_containers_api_endpoints(
original_function=mock_original_function,
custom_llm_provider="openai",
container_id=managed_id,
file_id="cfile_xyz",
)
mock_original_function.assert_called_once()
call_kw = mock_original_function.call_args.kwargs
assert call_kw["container_id"] == "cfile_upstream_abc"
assert call_kw["file_id"] == "cfile_xyz"
assert call_kw["custom_llm_provider"] == "azure"

70
tests/test_litellm/caching/test_dual_cache.py
View file

@ -1,5 +1,6 @@
import asyncio
import time
import uuid
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
@ -260,3 +261,72 @@ async def test_async_increment_cache_returns_none_when_no_in_memory_cache_and_re
f"Expected None when in_memory_cache is absent and Redis fails, got {result!r}. "
"Returning the delta (1.0) would silently miscalculate rate-limit counters."
)
def test_dual_cache_late_attach_redis_wires_writes_and_ttl_sync():
"""
Typical lazy startup (sync): DualCache runs with in-memory only, then Redis
becomes available and is attached. New writes must reach Redis; keys written
before attach are not backfilled. Optional default_redis_ttl is applied on attach.
"""
in_memory = InMemoryCache()
dual_cache = DualCache(in_memory_cache=in_memory, redis_cache=None)
mock_redis = MagicMock()
mock_redis.set_cache = MagicMock()
mock_redis.async_set_cache = AsyncMock()
key_before = f"before_attach_{uuid.uuid4()}"
val_before = {"phase": "memory_only"}
dual_cache.set_cache(key_before, val_before)
assert in_memory.get_cache(key_before) == val_before
dual_cache.attach_redis_cache(mock_redis, default_redis_ttl=99.0)
assert dual_cache.redis_cache is mock_redis
assert dual_cache.default_redis_ttl == 99.0
mock_redis.set_cache.assert_not_called()
key_after = f"after_attach_{uuid.uuid4()}"
val_after = {"phase": "memory_and_redis"}
dual_cache.set_cache(key_after, val_after)
mock_redis.set_cache.assert_called_once()
assert mock_redis.set_cache.call_args[0][:2] == (key_after, val_after)
assert in_memory.get_cache(key_after) == val_after
@pytest.mark.asyncio
async def test_dual_cache_late_attach_redis_wires_writes_and_ttl_async():
"""
Typical lazy startup (async): DualCache runs with in-memory only, then Redis
becomes available and is attached. New writes must reach Redis; keys written
before attach are not backfilled. Optional default_redis_ttl is applied on attach.
"""
in_memory = InMemoryCache()
dual_cache = DualCache(in_memory_cache=in_memory, redis_cache=None)
mock_redis = MagicMock()
mock_redis.set_cache = MagicMock()
mock_redis.async_set_cache = AsyncMock()
key_before = f"before_attach_{uuid.uuid4()}"
val_before = {"phase": "memory_only"}
await dual_cache.async_set_cache(key_before, val_before)
assert in_memory.get_cache(key_before) == val_before
dual_cache.attach_redis_cache(mock_redis, default_redis_ttl=99.0)
assert dual_cache.redis_cache is mock_redis
assert dual_cache.default_redis_ttl == 99.0
mock_redis.async_set_cache.assert_not_called()
key_after = f"after_attach_{uuid.uuid4()}"
val_after = {"phase": "memory_and_redis"}
await dual_cache.async_set_cache(key_after, val_after)
mock_redis.async_set_cache.assert_called_once()
assert mock_redis.async_set_cache.call_args[0][:2] == (key_after, val_after)
assert in_memory.get_cache(key_after) == val_after

44
tests/test_litellm/caching/test_redis_cache.py
View file

@ -50,6 +50,50 @@ async def test_redis_cache_async_increment(namespace, monkeypatch, redis_no_ping
)
@pytest.mark.asyncio
async def test_redis_cache_async_increment_refresh_ttl_true_bumps_existing_ttl(
monkeypatch, redis_no_ping
):
"""With refresh_ttl=True, every increment should call expire() to bump
the TTL, even when the key already has a TTL (counter-style use)."""
monkeypatch.setenv("REDIS_HOST", "https://my-test-host")
redis_cache = RedisCache()
mock_redis_instance = AsyncMock()
mock_redis_instance.__aenter__.return_value = mock_redis_instance
mock_redis_instance.__aexit__.return_value = None
mock_redis_instance.ttl.return_value = 42 # key already has ~42s left
with patch.object(
redis_cache, "init_async_client", return_value=mock_redis_instance
):
await redis_cache.async_increment(
key="spend:team_member:u:t", value=0.05, refresh_ttl=True
)
mock_redis_instance.expire.assert_awaited_once_with("spend:team_member:u:t", 60)
@pytest.mark.asyncio
async def test_redis_cache_async_increment_default_does_not_bump_existing_ttl(
monkeypatch, redis_no_ping
):
"""Default (refresh_ttl=False) preserves window-style semantics: TTL is
set only on first creation, never refreshed (used by rate-limit windows)."""
monkeypatch.setenv("REDIS_HOST", "https://my-test-host")
redis_cache = RedisCache()
mock_redis_instance = AsyncMock()
mock_redis_instance.__aenter__.return_value = mock_redis_instance
mock_redis_instance.__aexit__.return_value = None
mock_redis_instance.ttl.return_value = 42 # key already has ~42s left
with patch.object(
redis_cache, "init_async_client", return_value=mock_redis_instance
):
await redis_cache.async_increment(key="rate_limit:window", value=1)
mock_redis_instance.expire.assert_not_awaited()
@pytest.mark.asyncio
async def test_redis_client_init_with_socket_timeout(monkeypatch, redis_no_ping):
monkeypatch.setenv("REDIS_HOST", "my-fake-host")

204
tests/test_litellm/containers/test_azure_container_transformation.py
View file

@ -11,6 +11,7 @@ sys.path.insert(0, os.path.abspath("../../../"))
import litellm
from litellm.llms.azure.containers.transformation import AzureContainerConfig
from litellm.llms.base_llm.containers.transformation import BaseContainerConfig
from litellm.responses.utils import ResponsesAPIRequestUtils
from litellm.types.containers.main import (
ContainerFileListResponse,
ContainerListResponse,
@ -518,3 +519,206 @@ class TestAzureContainerKnownFailureRegressions:
c2 = _get_container_provider_config("azure_text")
assert type(c1) is type(c2)
assert isinstance(c1, AzureContainerConfig)
@pytest.mark.asyncio
async def test_proxy_process_request_preserves_managed_container_id(
self, monkeypatch
):
from starlette.requests import Request
from litellm.proxy.container_endpoints import handler_factory
encoded_id = ResponsesAPIRequestUtils._build_container_id(
custom_llm_provider="azure",
model_id="model_abc123",
container_id="cntr_123",
)
captured = {}
async def _mock_base_process_llm_request(
self,
request,
fastapi_response,
user_api_key_dict,
route_type,
**kwargs,
):
captured["data"] = self.data
captured["route_type"] = route_type
return {"id": "cfile_abc"}
from litellm.proxy.common_request_processing import (
ProxyBaseLLMRequestProcessing,
)
monkeypatch.setattr(
ProxyBaseLLMRequestProcessing,
"base_process_llm_request",
_mock_base_process_llm_request,
)
request = Request(
{
"type": "http",
"method": "GET",
"path": "/v1/containers/id/files/id/content",
"headers": [],
"query_string": b"",
}
)
fastapi_response = MagicMock()
await handler_factory._process_request(
request=request,
fastapi_response=fastapi_response,
user_api_key_dict=MagicMock(),
route_type="alist_container_files",
path_params={"container_id": encoded_id},
)
assert captured["route_type"] == "alist_container_files"
assert captured["data"]["container_id"] == encoded_id
assert captured["data"]["custom_llm_provider"] == "openai"
assert "model_id" not in captured["data"]
assert "api_base" not in captured["data"]
@pytest.mark.asyncio
async def test_regression_binary_file_request_routes_through_proxy_processor(
self, monkeypatch
):
from fastapi import Response
from starlette.requests import Request
from litellm.proxy.container_endpoints import handler_factory
encoded_id = ResponsesAPIRequestUtils._build_container_id(
custom_llm_provider="azure",
model_id="model_abc123",
container_id="cntr_123",
)
captured = {}
async def _mock_base_process_llm_request(
self,
request,
fastapi_response,
user_api_key_dict,
route_type,
**kwargs,
):
captured["data"] = self.data
captured["route_type"] = route_type
fastapi_response.headers["x-litellm-call-id"] = "call-123"
return b"csv-bytes"
from litellm.proxy.common_request_processing import (
ProxyBaseLLMRequestProcessing,
)
monkeypatch.setattr(
ProxyBaseLLMRequestProcessing,
"base_process_llm_request",
_mock_base_process_llm_request,
)
request = Request(
{
"type": "http",
"method": "GET",
"path": "/v1/containers/id/files/id/content",
"headers": [],
"query_string": b"",
}
)
fastapi_response = Response()
response = await handler_factory._process_binary_request(
request=request,
fastapi_response=fastapi_response,
container_id=encoded_id,
file_id="cfile_abc",
user_api_key_dict=MagicMock(),
)
assert captured["route_type"] == "aretrieve_container_file_content"
assert captured["data"]["container_id"] == encoded_id
assert captured["data"]["file_id"] == "cfile_abc"
assert captured["data"]["custom_llm_provider"] == "openai"
assert response.status_code == 200
assert response.body == b"csv-bytes"
assert response.headers["x-litellm-call-id"] == "call-123"
@pytest.mark.asyncio
async def test_regression_multipart_upload_request_uses_provider_from_managed_id(
self, monkeypatch
):
from starlette.requests import Request
from litellm.proxy.common_request_processing import (
ProxyBaseLLMRequestProcessing,
)
from litellm.proxy.common_utils import http_parsing_utils
from litellm.proxy.container_endpoints import handler_factory
encoded_id = ResponsesAPIRequestUtils._build_container_id(
custom_llm_provider="azure",
model_id="model_abc123",
container_id="cntr_123",
)
captured = {}
async def _mock_get_form_data(request):
return {"file": "ignored"}
async def _mock_convert_upload_files_to_file_data(form_data):
return {"file": [("data.csv", b"csv-bytes", "text/csv")]}
async def _mock_base_process_llm_request(
self,
request,
fastapi_response,
user_api_key_dict,
route_type,
**kwargs,
):
captured["data"] = self.data
captured["route_type"] = route_type
return {"id": "cfile_abc"}
monkeypatch.setattr(
http_parsing_utils,
"get_form_data",
_mock_get_form_data,
)
monkeypatch.setattr(
http_parsing_utils,
"convert_upload_files_to_file_data",
_mock_convert_upload_files_to_file_data,
)
monkeypatch.setattr(
ProxyBaseLLMRequestProcessing,
"base_process_llm_request",
_mock_base_process_llm_request,
)
request = Request(
{
"type": "http",
"method": "POST",
"path": "/v1/containers/id/files",
"headers": [],
"query_string": b"",
}
)
await handler_factory._process_multipart_upload_request(
request=request,
fastapi_response=MagicMock(),
user_api_key_dict=MagicMock(),
route_type="aupload_container_file",
container_id=encoded_id,
)
assert captured["route_type"] == "aupload_container_file"
assert captured["data"]["container_id"] == encoded_id
assert captured["data"]["custom_llm_provider"] == "openai"

242
tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py
View file

@ -9,8 +9,10 @@ from litellm.litellm_core_utils.prompt_templates.factory import (
BAD_MESSAGE_ERROR_STR,
BedrockConverseMessagesProcessor,
BedrockImageProcessor,
anthropic_messages_pt,
_bedrock_converse_messages_pt,
_convert_to_bedrock_tool_call_invoke,
_convert_to_bedrock_tool_call_result,
anthropic_messages_pt,
convert_to_gemini_tool_call_result,
ollama_pt,
sanitize_messages_for_tool_calling,
@ -2485,10 +2487,6 @@ def test_convert_to_anthropic_tool_result_openai_file_pdf_becomes_document():
inside the tool_result content. Reuses anthropic_process_openai_file_message,
which already handles this for user messages.
"""
from litellm.litellm_core_utils.prompt_templates.factory import (
convert_to_anthropic_tool_result,
)
pdf_b64 = "JVBERi0xLjQKJeLjz9MK"
message = {
"tool_call_id": "toolu_pdf_1",
@ -2505,157 +2503,105 @@ def test_convert_to_anthropic_tool_result_openai_file_pdf_becomes_document():
],
}
result = convert_to_anthropic_tool_result(message)
result = _convert_to_bedrock_tool_call_result(message)
assert result["type"] == "tool_result"
assert result["tool_use_id"] == "toolu_pdf_1"
content = result["content"]
assert isinstance(content, list) and len(content) == 1
block = content[0]
assert block["type"] == "document"
assert block["source"]["type"] == "base64"
assert block["source"]["media_type"] == "application/pdf"
assert block["source"]["data"] == pdf_b64
tool_result = result["toolResult"]
assert len(tool_result["content"]) == 1
assert "document" in tool_result["content"][0]
assert tool_result["content"][0]["document"]["format"] == "pdf"
assert tool_result["content"][0]["document"]["source"]["bytes"] == pdf_b64
def test_convert_to_anthropic_tool_result_image_url_pdf_data_uri_becomes_document():
"""
Regression: a PDF sent as an `image_url` data URI on the tool-result path
must translate to an Anthropic document block (not an image block — Anthropic
rejects image blocks whose media_type is a non-image like application/pdf).
"""
from litellm.litellm_core_utils.prompt_templates.factory import (
convert_to_anthropic_tool_result,
)
def test_bedrock_converse_messages_pt_document_various_formats():
"""Test that various document media types produce the correct format value."""
test_cases = [
("application/pdf", "pdf"),
("text/csv", "csv"),
("text/html", "html"),
("text/plain", "txt"),
("text/markdown", "md"),
(
"application/vnd.openxmlformats-officedocument.wordprocessingml.document",
"docx",
),
]
pdf_b64 = "JVBERi0xLjQKJeLjz9MK"
message = {
"tool_call_id": "toolu_pdf_img_1",
"role": "tool",
"name": "fetch_document",
"content": [
for media_type, expected_format in test_cases:
messages = [
{
"type": "image_url",
"image_url": {
"url": f"data:application/pdf;base64,{pdf_b64}",
"role": "user",
"content": [
{
"type": "document",
"source": {
"type": "base64",
"media_type": media_type,
"data": "dGVzdA==",
},
},
],
}
]
result = _bedrock_converse_messages_pt(
messages, "anthropic.claude-sonnet-4-6", "bedrock"
)
doc_block = result[0]["content"][0]
assert doc_block["document"]["format"] == expected_format, (
f"Expected format '{expected_format}' for media_type '{media_type}', "
f"got '{doc_block['document']['format']}'"
)
def test_bedrock_converse_messages_pt_document_deterministic_name():
"""Test that the same document data always produces the same name."""
messages = [
{
"role": "user",
"content": [
{
"type": "document",
"source": {
"type": "base64",
"media_type": "application/pdf",
"data": "dGVzdA==",
},
},
},
],
}
],
}
]
result = convert_to_anthropic_tool_result(message)
content = result["content"]
assert isinstance(content, list) and len(content) == 1
block = content[0]
assert block["type"] == "document"
assert block["source"]["media_type"] == "application/pdf"
assert block["source"]["data"] == pdf_b64
def test_convert_to_anthropic_tool_result_image_url_unsupported_mime_stays_image_path():
"""
An `image_url` data URI whose mime is neither application/pdf nor text/plain
(e.g. application/json) must NOT be routed through the document path. Anthropic
only accepts application/pdf and text/plain as base64 document media_types —
anything else would produce a document block the API rejects. The old
(pre-fix) behavior was to wrap such data as an image block, which also
fails but stays on the image code path; preserve that failure mode rather
than switching to a document path that is equally broken.
"""
from litellm.litellm_core_utils.prompt_templates.factory import (
convert_to_anthropic_tool_result,
result1 = _bedrock_converse_messages_pt(
messages, "anthropic.claude-sonnet-4-6", "bedrock"
)
result2 = _bedrock_converse_messages_pt(
messages, "anthropic.claude-sonnet-4-6", "bedrock"
)
message = {
"tool_call_id": "toolu_json_1",
"role": "tool",
"name": "fetch_json",
"content": [
{
"type": "image_url",
"image_url": {
"url": "data:application/json;base64,eyJrIjoidiJ9",
name1 = result1[0]["content"][0]["document"]["name"]
name2 = result2[0]["content"][0]["document"]["name"]
assert name1 == name2
def test_bedrock_converse_messages_pt_document_rejects_url_source():
"""Test that a URL-type document source raises a clear error instead of KeyError."""
messages = [
{
"role": "user",
"content": [
{
"type": "document",
"source": {
"type": "url",
"url": "https://example.com/doc.pdf",
},
},
},
],
}
],
}
]
result = convert_to_anthropic_tool_result(message)
content = result["content"]
assert isinstance(content, list) and len(content) == 1
block = content[0]
assert block["type"] == "image", (
f"unsupported mime {block.get('source', {}).get('media_type')!r} "
f"should not be routed to document path; got {block}"
)
def test_convert_to_anthropic_tool_result_image_url_text_plain_data_uri_becomes_document():
"""
text/plain is one of the two mimes Anthropic accepts as a base64 document
media_type. Confirm it routes through the document path so tightening the
gate to {application/pdf, text/plain} (not "application/*") covers both.
"""
from litellm.litellm_core_utils.prompt_templates.factory import (
convert_to_anthropic_tool_result,
)
txt_b64 = "aGVsbG8=" # "hello"
message = {
"tool_call_id": "toolu_txt_1",
"role": "tool",
"name": "fetch_text",
"content": [
{
"type": "image_url",
"image_url": {
"url": f"data:text/plain;base64,{txt_b64}",
},
},
],
}
result = convert_to_anthropic_tool_result(message)
content = result["content"]
assert isinstance(content, list) and len(content) == 1
block = content[0]
assert block["type"] == "document"
assert block["source"]["media_type"] == "text/plain"
assert block["source"]["data"] == txt_b64
def test_convert_to_anthropic_tool_result_image_url_png_still_becomes_image():
"""
Regression: image_url with a real image mime type must continue to translate
to an Anthropic image block. Locks in existing behavior after the
data-URI-mime-type branching for PDFs.
"""
from litellm.litellm_core_utils.prompt_templates.factory import (
convert_to_anthropic_tool_result,
)
png_b64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR4nGNgYGBgAAAABQABXvMqOgAAAABJRU5ErkJggg=="
message = {
"tool_call_id": "toolu_png_1",
"role": "tool",
"name": "fetch_image",
"content": [
{
"type": "image_url",
"image_url": {
"url": f"data:image/png;base64,{png_b64}",
},
},
],
}
result = convert_to_anthropic_tool_result(message)
content = result["content"]
assert isinstance(content, list) and len(content) == 1
block = content[0]
assert block["type"] == "image"
assert block["source"]["media_type"] == "image/png"
with pytest.raises(ValueError, match="only supports base64-encoded"):
_bedrock_converse_messages_pt(
messages, "anthropic.claude-sonnet-4-6", "bedrock"
)

34
tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py
View file

@ -8,6 +8,7 @@ sys.path.insert(
) # Adds the parent directory to the system path
from unittest.mock import MagicMock, patch
from litellm.constants import RESPONSE_FORMAT_TOOL_NAME
from litellm.llms.anthropic.chat.transformation import AnthropicConfig
from litellm.llms.anthropic.experimental_pass_through.messages.transformation import (
AnthropicMessagesConfig,
@ -38,6 +39,39 @@ def test_response_format_transformation_unit_test():
print(result)
def test_anthropic_json_mode_non_streaming_mixed_internal_and_user_tools():
"""Non-streaming + response_format: internal json tool must not require len(tool_calls)==1."""
config = AnthropicConfig()
tool_calls = [
{
"id": "toolu_json",
"type": "function",
"function": {
"name": RESPONSE_FORMAT_TOOL_NAME,
"arguments": '{"values": {"answer": 42}}',
},
"index": 0,
},
{
"id": "toolu_user",
"type": "function",
"function": {
"name": "get_weather",
"arguments": '{"location": "NY"}',
},
"index": 1,
},
]
replacement, filtered, extra = config._resolve_json_mode_non_streaming(
json_mode=True,
tool_calls=tool_calls,
)
assert replacement is None
assert len(filtered) == 1
assert filtered[0]["function"]["name"] == "get_weather"
assert extra == '{"answer": 42}'
def test_calculate_usage():
"""
Do not include cache_creation_input_tokens in the prompt_tokens

74
tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py
View file

@ -1,3 +1,4 @@
import asyncio
import os
import sys
from unittest.mock import AsyncMock, Mock, patch
@ -8,6 +9,8 @@ import pytest
sys.path.insert(
0, os.path.abspath("../../../..")
) # Adds the parent directory to the system path
import litellm
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
from litellm.llms.custom_httpx.llm_http_handler import (
BaseLLMHTTPHandler,
_google_genai_streaming_hidden_params,
@ -103,7 +106,9 @@ def test_fingerprint_agentic_tools_is_deterministic():
tools_a = {"tool_calls": [{"id": "1", "input": {"q": "abc"}, "name": "web_search"}]}
tools_b = {"tool_calls": [{"name": "web_search", "input": {"q": "abc"}, "id": "1"}]}
assert handler._fingerprint_agentic_tools(tools_a) == handler._fingerprint_agentic_tools(tools_b)
assert handler._fingerprint_agentic_tools(
tools_a
) == handler._fingerprint_agentic_tools(tools_b)
@pytest.mark.asyncio
@ -350,3 +355,70 @@ def test_google_genai_streaming_hidden_params_model_info_and_router_fallback():
response_headers=httpx.Headers({}),
)
assert from_router["model_id"] == "router-model-id"
def _build_delete_response_mock(captured: dict):
"""Returns a fake httpx delete that records its kwargs."""
def _response() -> httpx.Response:
return httpx.Response(
status_code=200,
headers={"content-type": "application/json"},
content=b'{"id": "resp_x", "object": "response", "deleted": true}',
request=httpx.Request(method="DELETE", url="https://test.openai.azure.com"),
)
async def fake_async_delete(*args, **kwargs):
captured.update(kwargs)
return _response()
def fake_sync_delete(*args, **kwargs):
captured.update(kwargs)
return _response()
return fake_async_delete, fake_sync_delete
def test_async_delete_responses_omits_body_for_azure():
"""Azure responses DELETE rejects requests with any body. Verify the handler
does not pass `json=` to httpx when the transformer returns an empty dict."""
captured: dict = {}
fake_async_delete, _ = _build_delete_response_mock(captured)
async def run():
with patch.object(AsyncHTTPHandler, "delete", new=fake_async_delete):
await litellm.adelete_responses(
response_id="resp_xyz",
custom_llm_provider="azure",
api_base="https://test.openai.azure.com",
api_key="test-key",
api_version="2025-03-01-preview",
)
asyncio.run(run())
assert "json" not in captured
assert "data" not in captured
assert captured["url"].endswith(
"/openai/responses/resp_xyz?api-version=2025-03-01-preview"
)
def test_sync_delete_responses_omits_body_for_azure():
captured: dict = {}
_, fake_sync_delete = _build_delete_response_mock(captured)
with patch.object(HTTPHandler, "delete", new=fake_sync_delete):
litellm.delete_responses(
response_id="resp_xyz",
custom_llm_provider="azure",
api_base="https://test.openai.azure.com",
api_key="test-key",
api_version="2025-03-01-preview",
)
assert "json" not in captured
assert "data" not in captured
assert captured["url"].endswith(
"/openai/responses/resp_xyz?api-version=2025-03-01-preview"
)

14
tests/test_litellm/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_transformation.py
View file

@ -373,6 +373,20 @@ class TestVertexAIImagenImageGenerationConfig:
assert request["parameters"]["sampleCount"] == 2
assert request["parameters"]["aspectRatio"] == "16:9"
def test_transform_image_generation_request_labels_from_metadata(self):
"""Billing labels from litellm_params.metadata.requester_metadata on predict body."""
request = self.config.transform_image_generation_request(
model="imagegeneration@006",
prompt="A cat",
optional_params={},
litellm_params={
"metadata": {"requester_metadata": {"team": "platform", "env": "prod"}}
},
headers={},
)
assert request["labels"] == {"team": "platform", "env": "prod"}
assert "labels" not in request["parameters"]
def test_transform_image_generation_response(self):
"""Test response transformation"""
mock_response = MagicMock(spec=httpx.Response)

16
tests/test_litellm/llms/vertex_ai/rerank/test_vertex_ai_rerank_transformation.py
View file

@ -216,6 +216,22 @@ class TestVertexAIRerankTransform:
)
assert request_data_default["ignoreRecordDetailsInResponse"] == False
def test_transform_rerank_request_user_labels_from_metadata(self):
"""Discovery Engine Rank API uses userLabels (string map) for billing."""
optional_params = {
"query": "q",
"documents": ["a", "b"],
}
request_data = self.config.transform_rerank_request(
model=self.model,
optional_rerank_params=optional_params,
headers={},
litellm_params={
"metadata": {"requester_metadata": {"app": "litellm", "tier": "1"}}
},
)
assert request_data["userLabels"] == {"app": "litellm", "tier": "1"}
def test_transform_rerank_request_missing_required_params(self):
"""Test that transform_rerank_request handles missing required parameters."""
# Test missing query

182
tests/test_litellm/llms/vertex_ai/rerank/test_vertex_ai_rerank_userlabels_e2e.py Normal file
View file

@ -0,0 +1,182 @@
"""
End-to-end tests for Vertex AI rerank `userLabels` propagation.
These tests go through the full `litellm.rerank()` call path with the HTTP
layer mocked, so they catch plumbing bugs (e.g. `litellm_params` losing
`metadata` between the rerank entrypoint and the Vertex transform) that
unit tests on `VertexAIRerankConfig.transform_rerank_request` miss.
"""
import asyncio
import json
import os
from unittest.mock import AsyncMock, MagicMock, patch
import httpx
import pytest
import litellm
import litellm.llms.vertex_ai.rerank.transformation
def _extract_body(call_kwargs):
"""The rerank handler sends `data=json.dumps(...)`, not `json=...`."""
if "json" in call_kwargs and call_kwargs["json"] is not None:
return call_kwargs["json"]
raw = call_kwargs.get("data")
if isinstance(raw, (bytes, bytearray)):
raw = raw.decode("utf-8")
return json.loads(raw) if raw else None
def _make_mock_rank_response():
mock_response = MagicMock(spec=httpx.Response)
mock_response.status_code = 200
mock_response.headers = {"content-type": "application/json"}
mock_response.json.return_value = {
"records": [
{"id": "0", "score": 0.9, "title": "doc 0", "content": "hello"},
{"id": "1", "score": 0.1, "title": "doc 1", "content": "world"},
]
}
mock_response.text = '{"records": []}'
return mock_response
def _make_async_mock_rank_response():
mock_response = AsyncMock()
mock_response.status_code = 200
mock_response.headers = {"content-type": "application/json"}
mock_response.json = MagicMock(
return_value={
"records": [
{"id": "0", "score": 0.9, "title": "doc 0", "content": "hello"},
]
}
)
mock_response.text = '{"records": []}'
return mock_response
@pytest.fixture
def clean_vertex_env():
saved = {}
for var in (
"GOOGLE_APPLICATION_CREDENTIALS",
"GOOGLE_CLOUD_PROJECT",
"VERTEXAI_PROJECT",
"VERTEXAI_CREDENTIALS",
"VERTEX_AI_CREDENTIALS",
"VERTEX_PROJECT",
"VERTEX_LOCATION",
"VERTEX_AI_PROJECT",
):
if var in os.environ:
saved[var] = os.environ.pop(var)
yield
for var, value in saved.items():
os.environ[var] = value
def _patch_vertex_auth():
return patch.object(
litellm.llms.vertex_ai.rerank.transformation.VertexAIRerankConfig,
"_ensure_access_token",
return_value=("test-access-token", "test-project-2049"),
)
def test_rerank_userlabels_propagates_from_metadata_sync(clean_vertex_env):
"""
`litellm.rerank(metadata={"requester_metadata": {...}})` must end up as
`userLabels` on the Discovery Engine `:rank` request body.
"""
captured = {}
def fake_post(*args, **kwargs):
captured["body"] = _extract_body(kwargs)
return _make_mock_rank_response()
with (
_patch_vertex_auth(),
patch(
"litellm.llms.custom_httpx.http_handler.HTTPHandler.post",
side_effect=fake_post,
),
):
litellm.rerank(
model="vertex_ai/semantic-ranker-default@latest",
query="what is gemini?",
documents=["hello", "world"],
vertex_project="test-project-2049",
vertex_credentials='{"type": "service_account"}',
metadata={"requester_metadata": {"team": "platform", "env": "prod"}},
)
body = captured["body"]
assert body is not None, "expected POST body to be captured"
assert "userLabels" in body, (
"Vertex rerank request body is missing `userLabels` — metadata was "
"lost between litellm.rerank() and transform_rerank_request. "
f"body keys: {sorted(body.keys())}"
)
assert body["userLabels"] == {"team": "platform", "env": "prod"}
def test_rerank_userlabels_propagates_from_metadata_async(clean_vertex_env):
"""Same as the sync test, but through `litellm.arerank`."""
captured = {}
async def fake_post(*args, **kwargs):
captured["body"] = _extract_body(kwargs)
return _make_async_mock_rank_response()
with (
_patch_vertex_auth(),
patch(
"litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post",
side_effect=fake_post,
),
):
asyncio.run(
litellm.arerank(
model="vertex_ai/semantic-ranker-default@latest",
query="what is gemini?",
documents=["hello", "world"],
vertex_project="test-project-2049",
vertex_credentials='{"type": "service_account"}',
metadata={"requester_metadata": {"team": "platform"}},
)
)
body = captured["body"]
assert body is not None
assert body.get("userLabels") == {"team": "platform"}
def test_rerank_userlabels_absent_when_no_metadata(clean_vertex_env):
"""No metadata → no `userLabels` key (don't send empty maps)."""
captured = {}
def fake_post(*args, **kwargs):
captured["body"] = _extract_body(kwargs)
return _make_mock_rank_response()
with (
_patch_vertex_auth(),
patch(
"litellm.llms.custom_httpx.http_handler.HTTPHandler.post",
side_effect=fake_post,
),
):
litellm.rerank(
model="vertex_ai/semantic-ranker-default@latest",
query="what is gemini?",
documents=["hello", "world"],
vertex_project="test-project-2049",
vertex_credentials='{"type": "service_account"}',
)
body = captured["body"]
assert body is not None
assert "userLabels" not in body

64
tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py
View file

@ -15,7 +15,9 @@ from litellm.llms.vertex_ai.common_utils import (
convert_anyof_null_to_nullable,
get_vertex_location_from_url,
get_vertex_project_id_from_url,
pop_vertex_request_labels,
set_schema_property_ordering,
vertex_request_labels_from_litellm_params,
)
@ -1444,3 +1446,65 @@ def test_add_object_type_does_not_add_type_when_anyof_present():
# Verify type was not added (anyOf handles the type)
assert "type" not in input_schema, "type should not be added when anyOf is present"
def test_vertex_request_labels_from_litellm_params_extracts_requester_metadata():
assert vertex_request_labels_from_litellm_params(None) is None
assert vertex_request_labels_from_litellm_params({}) is None
assert vertex_request_labels_from_litellm_params({"metadata": None}) is None
lp = {"metadata": {"requester_metadata": {"team": "analytics", "count": 3}}}
assert vertex_request_labels_from_litellm_params(lp) == {"team": "analytics"}
def test_vertex_request_labels_from_litellm_params_accepts_litellm_metadata():
lp = {
"litellm_metadata": {
"requester_metadata": {"team": "platform", "count": 3}
}
}
assert vertex_request_labels_from_litellm_params(lp) == {"team": "platform"}
def test_vertex_request_labels_prefers_metadata_over_litellm_metadata():
lp = {
"metadata": {"requester_metadata": {"source": "metadata"}},
"litellm_metadata": {"requester_metadata": {"source": "litellm_metadata"}},
}
assert vertex_request_labels_from_litellm_params(lp) == {"source": "metadata"}
def test_pop_vertex_request_labels_prefers_explicit_labels_then_metadata():
optional = {"labels": {"env": "prod"}}
litellm_params = {"metadata": {"requester_metadata": {"team": "x"}}}
assert pop_vertex_request_labels(optional, litellm_params) == {"env": "prod"}
assert "labels" not in optional
optional2: dict = {}
assert pop_vertex_request_labels(optional2, litellm_params) == {"team": "x"}
optional3 = {"labels": {"team": 123}}
assert pop_vertex_request_labels(optional3, litellm_params) == {"team": "x"}
def test_pop_vertex_request_labels_uses_litellm_metadata_when_metadata_absent():
optional: dict = {}
litellm_params = {
"litellm_metadata": {"requester_metadata": {"team": "from_litellm_meta"}}
}
assert pop_vertex_request_labels(optional, litellm_params) == {
"team": "from_litellm_meta"
}
def test_vertex_text_embedding_request_includes_labels_from_metadata():
import litellm
req = litellm.vertexAITextEmbeddingConfig.transform_openai_request_to_vertex_embedding_request(
input="hi",
optional_params={},
model="text-embedding-004",
litellm_params={
"metadata": {"requester_metadata": {"project_id": "cost-center-1"}}
},
)
assert req.get("labels") == {"project_id": "cost-center-1"}

41
tests/test_litellm/llms/vertex_ai/test_vertex_model_garden_openapi.py Normal file
View file

@ -0,0 +1,41 @@
"""Vertex Model Garden: OpenAPI base URL for publisher/model ids vs per-endpoint path."""
import pytest
from litellm.llms.vertex_ai.vertex_model_garden.main import (
_vertex_model_garden_model_id_in_json_body,
create_vertex_url,
)
@pytest.mark.parametrize(
"model,expect_openapi_base",
[
("xai/grok-4.1-fast-reasoning", True),
("openai/foo/bar", True),
("5464397967697903616", False),
("gpt-oss-20b-maas", False),
],
)
def test_create_vertex_url_openapi_vs_deployed_endpoint(
model: str, expect_openapi_base: bool
) -> None:
url = create_vertex_url(
vertex_location="us-central1",
vertex_project="my-project",
stream=False,
model=model,
)
if expect_openapi_base:
assert "/v1/projects/my-project/locations/us-central1/endpoints/openapi" in url
else:
assert (
"/v1beta1/projects/my-project/locations/us-central1/endpoints/"
f"{model}" in url
)
assert "openapi" not in url
def test_model_id_in_json_body_heuristic() -> None:
assert _vertex_model_garden_model_id_in_json_body("xai/grok-4.1-fast-reasoning") is True
assert _vertex_model_garden_model_id_in_json_body("5464397967697903616") is False

39
tests/test_litellm/llms/xai/test_xai_chat_transformation.py Normal file
View file

@ -0,0 +1,39 @@
import os
import sys
sys.path.insert(
0, os.path.abspath("../../../..")
) # Adds the parent directory to the system path
from litellm.llms.xai.chat.transformation import XAIChatConfig
class TestXAIParallelToolCalls:
"""Test suite for XAI parallel tool calls functionality."""
def test_get_supported_openai_params_includes_parallel_tool_calls(self):
"""Test that parallel_tool_calls is in supported parameters."""
config = XAIChatConfig()
supported_params = config.get_supported_openai_params(
"xai/grok-4.20"
)
assert "parallel_tool_calls" in supported_params
def test_transform_request_preserves_parallel_tool_calls(self):
"""Test that transform_request preserves parallel_tool_calls parameter."""
config = XAIChatConfig()
messages = [{"role": "user", "content": "What's the weather like?"}]
optional_params = {"parallel_tool_calls": True}
result = config.transform_request(
model="xai/grok-4.20",
messages=messages,
optional_params=optional_params,
litellm_params={},
headers={},
)
assert result.get("parallel_tool_calls") is True
assert len(result["messages"]) == 1
assert result["messages"][0]["role"] == "user"

100
tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py
View file

@ -673,6 +673,106 @@ class TestHookHeaderMergePriority:
assert headers["X-OAuth"] == "yes"
assert headers["X-Trace-Id"] == "trace-123"
@pytest.mark.asyncio
async def test_m2m_oauth2_does_not_forward_litellm_caller_authorization(self):
"""M2M must not put caller Bearer (LiteLLM API key) into extra_headers (#23652)."""
manager = MCPServerManager()
server = MCPServer(
server_id="test-id",
name="Test Server",
server_name="test_server",
url="https://example.com",
transport=MCPTransport.http,
auth_type=MCPAuth.oauth2,
oauth2_flow="client_credentials",
token_url="https://auth.example.com/token",
)
captured_extra_headers: Dict[str, Any] = {}
async def fake_create_mcp_client(
server, mcp_auth_header=None, extra_headers=None, stdio_env=None
):
captured_extra_headers["value"] = extra_headers
mock_client = MagicMock()
mock_client.call_tool = AsyncMock(return_value=MagicMock())
return mock_client
with patch.object(
manager, "_create_mcp_client", side_effect=fake_create_mcp_client
):
with patch.object(manager, "_build_stdio_env", return_value=None):
try:
await manager._call_regular_mcp_tool(
mcp_server=server,
original_tool_name="test_tool",
arguments={"key": "val"},
tasks=[],
mcp_auth_header=None,
mcp_server_auth_headers=None,
oauth2_headers={"Authorization": "Bearer sk-1234"},
raw_headers={"authorization": "Bearer sk-1234"},
proxy_logging_obj=None,
hook_extra_headers=None,
)
except Exception:
pass
assert captured_extra_headers.get("value") is None
@pytest.mark.asyncio
async def test_m2m_oauth2_skips_authorization_in_configured_extra_headers(self):
"""M2M must not take Authorization from raw_headers even if extra_headers lists it."""
manager = MCPServerManager()
server = MCPServer(
server_id="test-id",
name="Test Server",
server_name="test_server",
url="https://example.com",
transport=MCPTransport.http,
auth_type=MCPAuth.oauth2,
oauth2_flow="client_credentials",
token_url="https://auth.example.com/token",
extra_headers=["Authorization", "X-Custom"],
)
captured_extra_headers: Dict[str, Any] = {}
async def fake_create_mcp_client(
server, mcp_auth_header=None, extra_headers=None, stdio_env=None
):
captured_extra_headers["value"] = extra_headers
mock_client = MagicMock()
mock_client.call_tool = AsyncMock(return_value=MagicMock())
return mock_client
with patch.object(
manager, "_create_mcp_client", side_effect=fake_create_mcp_client
):
with patch.object(manager, "_build_stdio_env", return_value=None):
try:
await manager._call_regular_mcp_tool(
mcp_server=server,
original_tool_name="test_tool",
arguments={"key": "val"},
tasks=[],
mcp_auth_header=None,
mcp_server_auth_headers=None,
oauth2_headers={"Authorization": "Bearer sk-1234"},
raw_headers={
"authorization": "Bearer sk-1234",
"x-custom": "from-client",
},
proxy_logging_obj=None,
hook_extra_headers=None,
)
except Exception:
pass
headers = captured_extra_headers.get("value") or {}
assert "Authorization" not in headers
assert headers.get("X-Custom") == "from-client"
class TestUserAPIKeyAuthJwtClaims:
"""Tests that UserAPIKeyAuth correctly carries jwt_claims."""

394
tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py
View file

@ -17,6 +17,7 @@ from litellm.proxy._types import (
MCPTransport,
UserAPIKeyAuth,
)
from litellm.types.mcp import MCPAuth
from litellm.types.mcp_server.mcp_server_manager import MCPServer
@ -135,6 +136,152 @@ def test_prepare_mcp_server_headers_case_insensitive_extra_headers():
assert extra_headers == {"Authorization": "Bearer token"}
def test_prepare_mcp_server_headers_oauth2_m2m_omits_litellm_caller_authorization():
"""M2M OAuth must not put caller Bearer (LiteLLM API key) into extra_headers (#23652)."""
try:
from litellm.proxy._experimental.mcp_server.server import (
_prepare_mcp_server_headers,
)
except ImportError:
pytest.skip("MCP server not available")
server = MCPServer(
server_id="m2m-server",
name="m2m",
transport=MCPTransport.http,
auth_type=MCPAuth.oauth2,
oauth2_flow="client_credentials",
token_url="https://auth.example.com/token",
)
caller_key = {"Authorization": "Bearer sk-litellm-caller"}
server_auth_header, extra_headers = _prepare_mcp_server_headers(
server=server,
mcp_server_auth_headers=None,
mcp_auth_header=None,
oauth2_headers=caller_key,
raw_headers=None,
)
assert server_auth_header is None
assert extra_headers is None
def test_prepare_mcp_server_headers_oauth2_interactive_copies_oauth2_headers():
"""Interactive OAuth still forwards the user's OAuth token in extra_headers."""
try:
from litellm.proxy._experimental.mcp_server.server import (
_prepare_mcp_server_headers,
)
except ImportError:
pytest.skip("MCP server not available")
user_oauth = {"Authorization": "Bearer upstream-user-token"}
server = MCPServer(
server_id="3lo-server",
name="3lo",
transport=MCPTransport.http,
auth_type=MCPAuth.oauth2,
oauth2_flow=None,
)
server_auth_header, extra_headers = _prepare_mcp_server_headers(
server=server,
mcp_server_auth_headers=None,
mcp_auth_header=None,
oauth2_headers=user_oauth,
raw_headers=None,
)
assert server_auth_header is None
assert extra_headers == user_oauth
def test_prepare_mcp_server_headers_m2m_skips_authorization_from_raw_extra_headers():
"""M2M must not merge caller Authorization from raw_headers when extra_headers lists it."""
try:
from litellm.proxy._experimental.mcp_server.server import (
_prepare_mcp_server_headers,
)
except ImportError:
pytest.skip("MCP server not available")
server = MCPServer(
server_id="m2m-raw",
name="m2m",
transport=MCPTransport.http,
auth_type=MCPAuth.oauth2,
oauth2_flow="client_credentials",
token_url="https://auth.example.com/token",
extra_headers=["Authorization", "X-Custom"],
)
server_auth_header, extra_headers = _prepare_mcp_server_headers(
server=server,
mcp_server_auth_headers=None,
mcp_auth_header=None,
oauth2_headers={"Authorization": "Bearer sk-1234"},
raw_headers={
"authorization": "Bearer sk-1234",
"x-custom": "trace",
},
)
assert server_auth_header is None
assert extra_headers is not None
assert "Authorization" not in extra_headers
assert extra_headers.get("X-Custom") == "trace"
@pytest.mark.asyncio
async def test_call_tool_m2m_skips_authorization_headers():
"""M2M call_tool must not forward caller Authorization in oauth2/raw headers."""
try:
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
MCPServerManager,
)
except ImportError:
pytest.skip("MCP server not available")
manager = MCPServerManager()
server = MCPServer(
server_id="m2m-call-tool",
name="m2m-call-tool",
server_name="m2m-call-tool",
transport=MCPTransport.http,
auth_type=MCPAuth.oauth2,
oauth2_flow="client_credentials",
token_url="https://auth.example.com/token",
client_id="cid",
client_secret="csecret",
extra_headers=["Authorization", "X-Custom"],
)
mock_client = MagicMock()
mock_client.call_tool = AsyncMock(return_value=MagicMock())
with patch.object(
manager, "_create_mcp_client", new=AsyncMock(return_value=mock_client)
) as create_client_mock:
await manager._call_regular_mcp_tool(
mcp_server=server,
original_tool_name="echo",
arguments={"message": "hello"},
tasks=[],
mcp_auth_header=None,
mcp_server_auth_headers=None,
oauth2_headers={"Authorization": "Bearer sk-1234"},
raw_headers={"authorization": "Bearer sk-1234", "x-custom": "trace"},
proxy_logging_obj=None,
)
create_kwargs = create_client_mock.await_args.kwargs
extra_headers = create_kwargs["extra_headers"] or {}
assert "Authorization" not in extra_headers
assert extra_headers.get("X-Custom") == "trace"
@pytest.mark.asyncio
async def test_get_prompts_from_mcp_servers_success():
try:
@ -2288,6 +2435,79 @@ async def test_get_tools_from_mcp_servers_logs_list_tools_to_spendlogs_when_enab
assert spend_meta["per_server_tool_counts"]["server_a"] == 1
@pytest.mark.asyncio
async def test_get_tools_from_mcp_servers_returns_tools_when_success_logging_fails():
"""
Regression test: list_tools should still return fetched tools even if
async_success_handler raises (e.g. serialization errors in logging path).
"""
try:
from litellm.proxy._experimental.mcp_server.server import (
_get_tools_from_mcp_servers,
)
from litellm.proxy._types import UserAPIKeyAuth
except ImportError:
pytest.skip("MCP server not available")
user_auth = UserAPIKeyAuth(api_key="test-key", user_id="test-user")
server_a = MagicMock(name="server_a_obj")
server_a.name = "server_a"
server_a.alias = "server_a"
server_a.server_name = "server_a"
server_a.server_id = "a"
server_a.auth_type = None
server_a.extra_headers = None
tool_1 = MagicMock()
tool_1.name = "server_a-tool_1"
dummy_logging_obj = MagicMock()
dummy_logging_obj.model_call_details = {"metadata": {"spend_logs_metadata": {}}}
dummy_logging_obj.async_success_handler = AsyncMock(
side_effect=TypeError("Object of type Tool is not JSON serializable")
)
with (
patch(
"litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers",
new=AsyncMock(return_value=[server_a]),
),
patch(
"litellm.proxy._experimental.mcp_server.server._prepare_mcp_server_headers",
return_value=(None, None),
),
patch(
"litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager",
) as mock_manager,
patch(
"litellm.proxy._experimental.mcp_server.server.filter_tools_by_allowed_tools",
side_effect=lambda tools, _server: tools,
),
patch(
"litellm.proxy._experimental.mcp_server.server.filter_tools_by_key_team_permissions",
new=AsyncMock(side_effect=lambda tools, **_: tools),
),
patch(
"litellm.proxy._experimental.mcp_server.server.function_setup",
return_value=(dummy_logging_obj, None),
),
):
mock_manager._get_tools_from_server = AsyncMock(return_value=[tool_1])
tools = await _get_tools_from_mcp_servers(
user_api_key_auth=user_auth,
mcp_auth_header=None,
mcp_servers=["server_a"],
mcp_server_auth_headers=None,
log_list_tools_to_spendlogs=True,
list_tools_log_source="mcp_protocol",
)
assert tools == [tool_1]
dummy_logging_obj.async_success_handler.assert_awaited_once()
def test_tool_name_matches_case_insensitive():
"""Test that _tool_name_matches performs case-insensitive comparison.
@ -2719,3 +2939,177 @@ class TestGatewayCreateInitializationOptions:
_mcp_gateway_initialize_instructions.reset(tok)
opts = server.create_initialization_options()
assert getattr(opts, "instructions", None) is None
@pytest.mark.asyncio
async def test_list_tools_with_legacy_db_m2m_server_resolves_oauth2_flow():
"""
P1 Regression: list_tools path must apply _resolve_oauth2_flow to legacy DB
rows where oauth2_flow is NULL but M2M credentials are present.
Without this fix, has_client_credentials returns False and the caller's
Authorization header is forwarded upstream instead of being blocked.
"""
try:
from litellm.proxy._experimental.mcp_server.server import (
_get_tools_from_mcp_servers,
)
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.mcp import MCPAuth
except ImportError:
pytest.skip("MCP server not available")
user_auth = UserAPIKeyAuth(api_key="sk-1234", user_id="test-user")
# Simulate a legacy DB row: OAuth2 with M2M credentials but oauth2_flow=None
legacy_server = MagicMock(name="legacy_m2m_server")
legacy_server.name = "legacy_m2m"
legacy_server.alias = "legacy_m2m"
legacy_server.server_name = "legacy_m2m"
legacy_server.server_id = "legacy-m2m-id"
legacy_server.auth_type = MCPAuth.oauth2
legacy_server.oauth2_flow = None # Legacy: field not set in DB
legacy_server.token_url = "https://oauth.example.com/token"
legacy_server.authorization_url = None
legacy_server.client_id = "client-id"
legacy_server.client_secret = "client-secret"
legacy_server.extra_headers = None
legacy_server.has_client_credentials = False # This is the bug: should be True
legacy_server.model_copy = MagicMock(
side_effect=lambda update: MCPServer(
server_id=legacy_server.server_id,
name=legacy_server.name,
transport=MCPTransport.http,
auth_type=legacy_server.auth_type,
oauth2_flow=update.get("oauth2_flow", legacy_server.oauth2_flow),
token_url=legacy_server.token_url,
authorization_url=legacy_server.authorization_url,
client_id=legacy_server.client_id,
client_secret=legacy_server.client_secret,
)
)
tool_1 = MagicMock()
tool_1.name = "legacy_m2m-tool"
captured_extra_headers = None
async def capture_extra_headers(*args, **kwargs):
nonlocal captured_extra_headers
captured_extra_headers = kwargs.get("extra_headers")
return [tool_1]
with (
patch(
"litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager",
) as mock_manager,
patch(
"litellm.proxy._experimental.mcp_server.server.filter_tools_by_allowed_tools",
side_effect=lambda tools, _server: tools,
),
patch(
"litellm.proxy._experimental.mcp_server.server.filter_tools_by_key_team_permissions",
new=AsyncMock(side_effect=lambda tools, **_: tools),
),
):
mock_manager.get_allowed_mcp_servers = AsyncMock(return_value=["legacy-m2m-id"])
mock_manager.get_mcp_server_by_id = MagicMock(return_value=legacy_server)
mock_manager.filter_server_ids_by_ip_with_info = MagicMock(
return_value=(["legacy-m2m-id"], 0)
)
mock_manager._get_tools_from_server = AsyncMock(
side_effect=capture_extra_headers
)
tools = await _get_tools_from_mcp_servers(
user_api_key_auth=user_auth,
mcp_auth_header=None,
mcp_servers=["legacy_m2m"],
mcp_server_auth_headers=None,
oauth2_headers={"Authorization": "Bearer sk-1234"}, # Caller's token
)
# With P1 fix: _get_allowed_mcp_servers applies _resolve_oauth2_flow,
# so has_client_credentials becomes True and extra_headers should be None
# (caller's Authorization blocked)
assert captured_extra_headers is None, (
"P1 security issue: caller's Authorization header was forwarded to M2M server. "
"Expected None, got: " + str(captured_extra_headers)
)
assert tools == [tool_1]
@pytest.mark.asyncio
async def test_call_tool_empty_extra_headers_returns_none():
"""
P2 Regression: When all configured extra_headers are filtered out (e.g.
Authorization for M2M), the resulting extra_headers should be None, not {}.
Downstream code that checks `if extra_headers is None` will behave
differently if an empty dict is passed instead.
"""
try:
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
MCPServerManager,
)
from litellm.types.mcp import MCPAuth
except ImportError:
pytest.skip("MCP server not available")
manager = MCPServerManager()
# M2M server with only Authorization in extra_headers
m2m_server = MCPServer(
server_id="m2m-srv",
name="m2m_test",
transport=MCPTransport.http,
auth_type=MCPAuth.oauth2,
oauth2_flow="client_credentials",
token_url="https://oauth.example.com/token",
client_id="client-id",
client_secret="client-secret",
extra_headers=["Authorization"], # Will be filtered out for M2M
)
raw_headers = {"Authorization": "Bearer sk-1234", "Content-Type": "application/json"}
captured_extra_headers = None
async def capture_create_mcp_client(*args, **kwargs):
nonlocal captured_extra_headers
captured_extra_headers = kwargs.get("extra_headers")
# Return a mock client
mock_client = AsyncMock()
mock_client.call_tool = AsyncMock(return_value=MagicMock(content=[]))
return mock_client
with (
patch.object(
manager,
"_create_mcp_client",
side_effect=capture_create_mcp_client,
),
patch.object(
manager,
"get_mcp_server_by_id",
return_value=m2m_server,
),
):
try:
await manager._call_regular_mcp_tool(
mcp_server=m2m_server,
original_tool_name="test_tool",
arguments={},
mcp_auth_header=None,
oauth2_headers=None,
raw_headers=raw_headers,
)
except Exception:
pass # We only care about the captured headers
# With P2 fix: extra_headers should be None (not {}) when all headers filtered
assert captured_extra_headers is None, (
"P2 API consistency issue: expected None for empty extra_headers, got: "
+ str(captured_extra_headers)
)

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

@ -922,19 +922,21 @@ async def test_get_tag_objects_batch():
# Simulate 5 tags: 2 cached, 3 uncached
tag_names = ["cached-1", "uncached-1", "cached-2", "uncached-2", "uncached-3"]
# Mock cached tags
cached_tag_1 = {
"tag_name": "cached-1",
"spend": 10.0,
"models": [],
"litellm_budget_table": None,
}
cached_tag_2 = {
"tag_name": "cached-2",
"spend": 20.0,
"models": [],
"litellm_budget_table": None,
}
# Mock cached tags — must be LiteLLM_TagTable instances: the mocked async_get_cache
# bypasses UserApiKeyCache deserialization, so returning plain dicts would flow through
# as dict (production returns models after Codec.deserialize inside the cache).
cached_tag_1 = LiteLLM_TagTable(
tag_name="cached-1",
spend=10.0,
models=[],
litellm_budget_table=None,
)
cached_tag_2 = LiteLLM_TagTable(
tag_name="cached-2",
spend=20.0,
models=[],
litellm_budget_table=None,
)
# Mock DB response for uncached tags
uncached_tag_1 = MagicMock()
@ -980,13 +982,13 @@ async def test_get_tag_objects_batch():
)
# Mock cache behavior - return cached tags, None for uncached
async def mock_get_cache(key):
async def mock_get_cache(*args, **kwargs):
key = kwargs.get("key")
if key == "tag:cached-1":
return cached_tag_1
elif key == "tag:cached-2":
if key == "tag:cached-2":
return cached_tag_2
else:
return None
return None
mock_cache.async_get_cache = AsyncMock(side_effect=mock_get_cache)
mock_cache.async_set_cache = AsyncMock()

10
tests/test_litellm/proxy/auth/test_handle_jwt.py
View file

@ -405,9 +405,9 @@ async def test_sync_user_role_and_teams_cache_invalidation_on_role_change():
mock_cache.async_set_cache.assert_called_once()
call_kwargs = mock_cache.async_set_cache.call_args
assert call_kwargs.kwargs["key"] == "u1"
assert (
call_kwargs.kwargs["value"]["user_role"] == LitellmUserRoles.PROXY_ADMIN.value
)
assert isinstance(call_kwargs.kwargs["value"], LiteLLM_UserTable)
assert call_kwargs.kwargs["value"].user_role == LitellmUserRoles.PROXY_ADMIN.value
assert call_kwargs.kwargs["model_type"] == LiteLLM_UserTable
@pytest.mark.asyncio
@ -452,7 +452,9 @@ async def test_sync_user_role_and_teams_cache_invalidation_on_team_change():
mock_cache.async_set_cache.assert_called_once()
call_kwargs = mock_cache.async_set_cache.call_args
assert call_kwargs.kwargs["key"] == "u1"
assert set(call_kwargs.kwargs["value"]["teams"]) == {"team1", "team2"}
assert isinstance(call_kwargs.kwargs["value"], LiteLLM_UserTable)
assert set(call_kwargs.kwargs["value"].teams) == {"team1", "team2"}
assert call_kwargs.kwargs["model_type"] == LiteLLM_UserTable
@pytest.mark.asyncio

126
tests/test_litellm/proxy/common_utils/test_cache_codec.py Normal file
View file

@ -0,0 +1,126 @@
import logging
from typing import Optional
from unittest.mock import patch
import pytest
from pydantic import BaseModel, ValidationError
from litellm.proxy.common_utils.cache_pydantic_utils import CacheCodec
class _SampleModel(BaseModel):
name: str
count: Optional[int] = None
class _SampleSubModel(_SampleModel):
pass
class TestCacheCodecSerialize:
def test_without_model_type_base_model_dumped_json_safe(self):
m = _SampleModel(name="a", count=1)
out = CacheCodec.serialize(m)
assert out == {"name": "a", "count": 1}
def test_without_model_type_dict_unchanged(self):
d = {"name": "x"}
assert CacheCodec.serialize(d) is d
def test_without_model_type_primitive_unchanged(self):
assert CacheCodec.serialize(42) == 42
def test_with_model_type_dict_validated_and_dumped(self):
out = CacheCodec.serialize({"name": "b", "count": 2}, model_type=_SampleModel)
assert out == {"name": "b", "count": 2}
def test_with_model_type_base_model_validated_and_dumped(self):
m = _SampleModel(name="c", count=None)
out = CacheCodec.serialize(m, model_type=_SampleModel)
assert out == {"name": "c"}
def test_with_model_type_exclude_none_on_dump(self):
out = CacheCodec.serialize({"name": "d"}, model_type=_SampleModel)
assert out == {"name": "d"}
assert "count" not in out
def test_with_model_type_non_dict_non_model_passthrough(self):
assert CacheCodec.serialize("raw", model_type=_SampleModel) == "raw"
def test_with_model_type_invalid_dict_raises(self):
with pytest.raises(ValidationError):
CacheCodec.serialize({"count": 1}, model_type=_SampleModel)
def test_with_model_type_already_correct_instance_skips_revalidation(self):
"""Fast-path: value is already model_type — model_validate must NOT be called."""
m = _SampleModel(name="fast", count=7)
with patch.object(_SampleModel, "model_validate", wraps=_SampleModel.model_validate) as mock_validate:
out = CacheCodec.serialize(m, model_type=_SampleModel)
assert out == {"name": "fast", "count": 7}
mock_validate.assert_not_called()
def test_with_model_type_subclass_instance_skips_revalidation(self):
"""Subclass is isinstance of base → should also take the fast path."""
sub = _SampleSubModel(name="sub", count=2)
with patch.object(_SampleModel, "model_validate", wraps=_SampleModel.model_validate) as mock_validate:
out = CacheCodec.serialize(sub, model_type=_SampleModel)
assert out == {"name": "sub", "count": 2}
mock_validate.assert_not_called()
def test_with_model_type_dict_input_goes_through_model_validate(self):
"""A dict value (not yet an instance) must still go through model_validate."""
raw = {"name": "via-dict", "count": 5}
with patch.object(
_SampleModel, "model_validate", wraps=_SampleModel.model_validate
) as mock_validate:
out = CacheCodec.serialize(raw, model_type=_SampleModel)
assert out == {"name": "via-dict", "count": 5}
mock_validate.assert_called_once()
def test_with_model_type_incompatible_model_raises_validation_error(self):
"""Passing a BaseModel whose fields don't satisfy model_type's required fields raises.
_IncompatibleModel only has `foo: int`, so when Pydantic v2 extracts its
data and validates it against _SampleModel (which requires `name: str`),
a ValidationError is raised.
"""
class _IncompatibleModel(BaseModel):
foo: int # missing required 'name' field of _SampleModel
with pytest.raises(ValidationError):
CacheCodec.serialize(_IncompatibleModel(foo=1), model_type=_SampleModel)
class TestCacheCodecDeserialize:
def test_none_returns_none(self):
assert CacheCodec.deserialize(None, _SampleModel) is None
def test_dict_validates_to_model(self):
m = CacheCodec.deserialize({"name": "e", "count": 3}, _SampleModel)
assert isinstance(m, _SampleModel)
assert m.name == "e"
assert m.count == 3
def test_instance_same_type_returned_as_is(self):
original = _SampleModel(name="f")
m = CacheCodec.deserialize(original, _SampleModel)
assert m is original
def test_subclass_instance_accepted(self):
sub = _SampleSubModel(name="g")
m = CacheCodec.deserialize(sub, _SampleModel)
assert m is sub
def test_wrong_type_returns_none(self):
assert CacheCodec.deserialize("not-a-dict", _SampleModel) is None
def test_invalid_dict_returns_none_and_logs_warning(self, caplog):
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
out = CacheCodec.deserialize({"count": 1}, _SampleModel)
assert out is None
assert any(
"CacheCodec.deserialize" in r.message and "_SampleModel" in r.message
for r in caplog.records
if r.levelno >= logging.WARNING
), f"Expected deserialize validation warning. Records: {[r.message for r in caplog.records]}"

156
tests/test_litellm/proxy/common_utils/test_reset_budget_job.py
View file

@ -1049,3 +1049,159 @@ def test_reset_budget_windows_query_error_does_not_break_team_path(monkeypatch):
asyncio.run(job.reset_budget_windows()) # must not raise
prisma_client.db.litellm_teamtable.update.assert_awaited_once()
# ---------------------------------------------------------------------------
# Counter invalidation on budget reset
# ---------------------------------------------------------------------------
def _make_counter_invalidation_job(monkeypatch):
"""Stub spend_counter_cache so we can observe invalidation calls."""
spend_counter_cache = MagicMock()
spend_counter_cache.in_memory_cache.set_cache = MagicMock()
spend_counter_cache.redis_cache = MagicMock()
spend_counter_cache.redis_cache.async_set_cache = AsyncMock()
fake_module = types.ModuleType("litellm.proxy.proxy_server")
fake_module.spend_counter_cache = spend_counter_cache
monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", fake_module)
return spend_counter_cache
def test_reset_budget_for_team_members_invalidates_redis_counter(monkeypatch):
"""Team-member budget reset clears the Redis spend counter."""
counter_cache = _make_counter_invalidation_job(monkeypatch)
expired_budget = type("B", (), {"budget_id": "budget-1"})
membership = type(
"Membership",
(),
{"user_id": "alice", "team_id": "team-x", "budget_id": "budget-1"},
)
prisma_client = MagicMock()
prisma_client.db.litellm_teammembership.find_many = AsyncMock(
return_value=[membership]
)
prisma_client.db.litellm_teammembership.update_many = AsyncMock(
return_value={"count": 1}
)
job = ResetBudgetJob(proxy_logging_obj=MagicMock(), prisma_client=prisma_client)
asyncio.run(job.reset_budget_for_litellm_team_members([expired_budget]))
counter_cache.in_memory_cache.set_cache.assert_any_call(
key="spend:team_member:alice:team-x", value=0.0, ttl=60
)
counter_cache.redis_cache.async_set_cache.assert_any_await(
key="spend:team_member:alice:team-x", value=0.0, ttl=60
)
def test_reset_budget_for_keys_invalidates_redis_counter(
reset_budget_job, mock_prisma_client, monkeypatch
):
"""Key budget reset must clear the Redis spend counter."""
counter_cache = _make_counter_invalidation_job(monkeypatch)
now = datetime.now(timezone.utc)
mock_prisma_client.data["key"] = [
type(
"Key",
(),
{
"spend": 100.0,
"budget_duration": "30d",
"budget_reset_at": now,
"id": "key-1",
"token": "sk-abc",
},
)
]
asyncio.run(reset_budget_job.reset_budget_for_litellm_keys())
counter_cache.in_memory_cache.set_cache.assert_any_call(
key="spend:key:sk-abc", value=0.0, ttl=60
)
def test_reset_budget_for_users_invalidates_redis_counter(
reset_budget_job, mock_prisma_client, monkeypatch
):
"""User budget reset must clear the Redis spend counter."""
counter_cache = _make_counter_invalidation_job(monkeypatch)
now = datetime.now(timezone.utc)
mock_prisma_client.data["user"] = [
type(
"User",
(),
{
"spend": 50.0,
"budget_duration": "7d",
"budget_reset_at": now,
"id": "user-1",
"user_id": "alice",
},
)
]
asyncio.run(reset_budget_job.reset_budget_for_litellm_users())
counter_cache.in_memory_cache.set_cache.assert_any_call(
key="spend:user:alice", value=0.0, ttl=60
)
def test_reset_budget_for_teams_invalidates_redis_counter(
reset_budget_job, mock_prisma_client, monkeypatch
):
"""Team budget reset must clear the Redis spend counter."""
counter_cache = _make_counter_invalidation_job(monkeypatch)
now = datetime.now(timezone.utc)
mock_prisma_client.data["team"] = [
type(
"Team",
(),
{
"spend": 200.0,
"budget_duration": "1mo",
"budget_reset_at": now,
"id": "team-1",
"team_id": "team-x",
},
)
]
asyncio.run(reset_budget_job.reset_budget_for_litellm_teams())
counter_cache.in_memory_cache.set_cache.assert_any_call(
key="spend:team:team-x", value=0.0, ttl=60
)
def test_reset_budget_for_keys_linked_to_budgets_invalidates_redis_counter(monkeypatch):
"""Resetting keys via budget tier must clear each linked key's counter."""
counter_cache = _make_counter_invalidation_job(monkeypatch)
expired_budget = type("B", (), {"budget_id": "budget-1"})
linked_key = type("Key", (), {"token": "sk-linked"})
prisma_client = MagicMock()
prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(
return_value=[linked_key]
)
prisma_client.db.litellm_verificationtoken.update_many = AsyncMock(
return_value={"count": 1}
)
job = ResetBudgetJob(proxy_logging_obj=MagicMock(), prisma_client=prisma_client)
asyncio.run(job.reset_budget_for_keys_linked_to_budgets([expired_budget]))
counter_cache.in_memory_cache.set_cache.assert_any_call(
key="spend:key:sk-linked", value=0.0, ttl=60
)

219
tests/test_litellm/proxy/common_utils/test_user_api_key_cache.py Normal file
View file

@ -0,0 +1,219 @@
import json
from typing import Any
import pytest
from litellm.caching.in_memory_cache import InMemoryCache
from litellm.caching.redis_cache import RedisCache
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.proxy.proxy_server import UserAPIKeyCacheTTLEnum
class CapturingInMemoryCache(InMemoryCache):
"""Records ``ttl`` passed into ``set_cache`` (what DualCache injects)."""
def __init__(self) -> None:
super().__init__()
self.last_ttl: Any = None
def set_cache(self, key, value, **kwargs): # type: ignore[override]
self.last_ttl = kwargs.get("ttl")
super().set_cache(key, value, **kwargs)
class FakeRedisCache(RedisCache):
"""
In-memory fake that enforces the UserApiKeyCache Redis payload contract.
For user_api_key_cache entries we expect Redis to store a JSON object (dict)
produced by `CacheCodec.serialize(..., model_type=...)`.
This fake:
- raises TypeError if the value is not a dict
- raises TypeError if the dict is not JSON-serializable
Records the ``ttl`` kwarg DualCache forwards on each Redis write for tests.
"""
def __init__(self): # noqa: super().__init__ skipped intentionally
self._store: dict[str, str] = {}
self.last_ttl: Any = None
def set_cache(self, key: str, value: Any, **kwargs): # type: ignore[override]
if not isinstance(value, dict):
raise TypeError("FakeRedisCache only accepts dict payloads")
self.last_ttl = kwargs.get("ttl")
self._store[key] = json.dumps(value)
return True
def get_cache(self, key: str, **kwargs): # type: ignore[override]
raw = self._store.get(key)
if raw is None:
return None
return json.loads(raw)
async def async_set_cache(self, key: str, value: Any, **kwargs): # type: ignore[override]
if not isinstance(value, dict):
raise TypeError("FakeRedisCache only accepts dict payloads")
self.last_ttl = kwargs.get("ttl")
self._store[key] = json.dumps(value)
return True
async def async_get_cache(self, key: str, **kwargs): # type: ignore[override]
raw = self._store.get(key)
if raw is None:
return None
return json.loads(raw)
def delete_cache(self, key: str): # type: ignore[override]
self._store.pop(key, None)
async def async_delete_cache(self, key: str): # type: ignore[override]
self._store.pop(key, None)
def _make_key_obj(token: str = "tok") -> UserAPIKeyAuth:
# Minimal object (UserAPIKeyAuth inherits token from base view).
return UserAPIKeyAuth(token=token)
class TestUserApiKeyCache:
@pytest.mark.asyncio
async def test_async_set_in_memory_gets_enum_default_when_user_api_key_cache_ttl_omitted(
self,
):
"""
If ``general_settings.user_api_key_cache_ttl`` is absent, the proxy never
calls ``update_cache_ttl``; ``user_api_key_cache`` keeps
``default_in_memory_ttl=UserAPIKeyCacheTTLEnum.in_memory_cache_ttl``.
DualCache must forward that as the in-memory ``ttl`` kwarg on each set.
"""
mem = CapturingInMemoryCache()
cache = UserApiKeyCache(
in_memory_cache=mem,
redis_cache=FakeRedisCache(),
default_in_memory_ttl=UserAPIKeyCacheTTLEnum.in_memory_cache_ttl.value,
)
await cache.async_set_cache(
"k",
_make_key_obj("t"),
model_type=UserAPIKeyAuth,
)
expected = UserAPIKeyCacheTTLEnum.in_memory_cache_ttl.value
assert mem.last_ttl == expected
def test_sync_set_in_memory_gets_enum_default_when_user_api_key_cache_ttl_omitted(
self,
):
mem = CapturingInMemoryCache()
cache = UserApiKeyCache(
in_memory_cache=mem,
redis_cache=FakeRedisCache(),
default_in_memory_ttl=UserAPIKeyCacheTTLEnum.in_memory_cache_ttl.value,
)
cache.set_cache("sk", _make_key_obj("s"), model_type=UserAPIKeyAuth)
assert mem.last_ttl == UserAPIKeyCacheTTLEnum.in_memory_cache_ttl.value
@pytest.mark.asyncio
async def test_async_set_forwards_default_in_memory_ttl_to_redis_layer(self):
"""
DualCache injects missing ``ttl`` from ``default_in_memory_ttl`` into kwargs
before calling ``redis_cache.async_set_cache`` — Redis should receive the same
TTL as memory (matches proxy defaults: enum 60s).
"""
fake = FakeRedisCache()
cache = UserApiKeyCache(
redis_cache=fake,
default_in_memory_ttl=60,
)
await cache.async_set_cache(
key="ttl-key",
value=_make_key_obj("ttl-tok"),
model_type=UserAPIKeyAuth,
)
assert fake.last_ttl == 60
@pytest.mark.asyncio
async def test_async_set_explicit_ttl_override_reaches_redis(self):
fake = FakeRedisCache()
cache = UserApiKeyCache(
redis_cache=fake,
default_in_memory_ttl=60,
)
await cache.async_set_cache(
key="k",
value=_make_key_obj("x"),
model_type=UserAPIKeyAuth,
ttl=900,
)
assert fake.last_ttl == 900
def test_sync_set_forwards_default_in_memory_ttl_to_redis_layer(self):
fake = FakeRedisCache()
cache = UserApiKeyCache(
redis_cache=fake,
default_in_memory_ttl=45,
)
cache.set_cache(
"sk",
_make_key_obj("sync"),
model_type=UserAPIKeyAuth,
)
assert fake.last_ttl == 45
@pytest.mark.asyncio
async def test_async_set_typed_stores_serialized_payload_in_memory_and_redis(self):
cache = UserApiKeyCache(redis_cache=FakeRedisCache())
obj = _make_key_obj("abc")
await cache.async_set_cache("k", obj, model_type=UserAPIKeyAuth)
# In-memory hit should still be raw dict (not BaseModel) because wrapper
# stores the serialized payload into both layers.
raw = await cache.in_memory_cache.async_get_cache("k") # type: ignore[union-attr]
assert isinstance(raw, dict)
assert raw["token"] == "abc"
# Redis should also hold the same serialized dict
redis_raw = await cache.redis_cache.async_get_cache("k") # type: ignore[union-attr]
assert redis_raw == raw
@pytest.mark.asyncio
async def test_async_get_typed_returns_model_on_valid_hit(self):
cache = UserApiKeyCache(redis_cache=FakeRedisCache())
await cache.async_set_cache("k", {"token": "abc"}, model_type=UserAPIKeyAuth)
value = await cache.async_get_cache("k", model_type=UserAPIKeyAuth)
assert value is not None
assert isinstance(value, UserAPIKeyAuth)
assert value.token == "abc"
@pytest.mark.asyncio
async def test_async_get_typed_returns_none_on_validation_failure_after_hit(self):
cache = UserApiKeyCache(redis_cache=FakeRedisCache())
# Bypass UserApiKeyCache.serialize: CacheCodec rejects non-dict cached values
# for dict-based models (deserialize returns None).
await cache.in_memory_cache.async_set_cache(
key="k", value="invalid-payload-not-a-dict"
)
value = await cache.async_get_cache("k", model_type=UserAPIKeyAuth)
assert value is None
def test_fake_redis_cache_rejects_non_json_serializable_values(self):
fake = FakeRedisCache()
class NotSerializable:
pass
with pytest.raises(TypeError):
fake.set_cache("k", NotSerializable())
with pytest.raises(TypeError):
fake.set_cache("k2", {"ok": NotSerializable()})

555
tests/test_litellm/proxy/google_endpoints/test_google_api_endpoints.py
View file

@ -4,7 +4,7 @@ Test to verify the Google GenAI proxy API endpoints
"""
import os
import sys
from unittest.mock import AsyncMock, patch
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
@ -13,520 +13,171 @@ sys.path.insert(
) # Adds the parent directory to the system path
def test_google_generate_content_endpoint():
"""Test that the google_generate_content endpoint correctly routes requests"""
# Skip this test if we can't import the required modules due to missing dependencies
try:
from fastapi import FastAPI
from fastapi.testclient import TestClient
def _build_test_client():
from fastapi import FastAPI
from fastapi.testclient import TestClient
from litellm.proxy.google_endpoints.endpoints import router as google_router
from litellm.proxy.google_endpoints.endpoints import router as google_router
app = FastAPI()
app.include_router(google_router)
return TestClient(app)
def _patch_base_process(return_value=None):
"""Patch ProxyBaseLLMRequestProcessing.base_process_llm_request so endpoint
tests don't run the full pipeline. Returns the AsyncMock so callers can
inspect call args."""
if return_value is None:
return_value = {"test": "response"}
return patch(
"litellm.proxy.google_endpoints.endpoints.ProxyBaseLLMRequestProcessing.base_process_llm_request",
new_callable=AsyncMock,
return_value=return_value,
)
def test_google_generate_content_endpoint():
"""generateContent routes through ProxyBaseLLMRequestProcessing with the
agenerate_content route_type — that pipeline runs pre_call_hook +
during_call_hook + post_call_success_hook for every guardrail callback."""
try:
client = _build_test_client()
except ImportError as e:
pytest.skip(f"Skipping test due to missing dependency: {e}")
# Create a FastAPI app and include the router (required for FastAPI 0.120+)
app = FastAPI()
app.include_router(google_router)
# Create a test client
client = TestClient(app)
# Mock the router's agenerate_content method
with patch("litellm.proxy.proxy_server.llm_router") as mock_router:
mock_router.agenerate_content = AsyncMock(return_value={"test": "response"})
# Send a request to the endpoint
with _patch_base_process() as mock_base:
response = client.post(
"/v1beta/models/test-model:generateContent",
json={"contents": [{"role": "user", "parts": [{"text": "Hello"}]}]},
)
# Verify the response
assert response.status_code == 200
assert response.json() == {"test": "response"}
# Verify that agenerate_content was called
mock_router.agenerate_content.assert_called_once()
mock_base.assert_called_once()
kwargs = mock_base.call_args.kwargs
assert kwargs["route_type"] == "agenerate_content"
assert kwargs["model"] == "test-model"
def test_google_stream_generate_content_endpoint():
"""Test that the google_stream_generate_content endpoint correctly routes streaming requests"""
# Skip this test if we can't import the required modules due to missing dependencies
"""streamGenerateContent must route through the same processor with the
streaming route_type so the guardrail pipeline runs."""
try:
from fastapi import FastAPI
from fastapi.testclient import TestClient
from litellm.proxy.google_endpoints.endpoints import router as google_router
client = _build_test_client()
except ImportError as e:
pytest.skip(f"Skipping test due to missing dependency: {e}")
# Create a FastAPI app and include the router (required for FastAPI 0.120+)
app = FastAPI()
app.include_router(google_router)
# Create a test client
client = TestClient(app)
# Mock the router's agenerate_content_stream method to return a stream
async def mock_stream_generator():
yield 'data: {"test": "stream_chunk_1"}\n\n'
yield 'data: {"test": "stream_chunk_2"}\n\n'
yield "data: [DONE]\n\n"
with patch("litellm.proxy.proxy_server.llm_router") as mock_router:
mock_router.agenerate_content_stream = AsyncMock(
return_value=mock_stream_generator()
)
# Send a request to the endpoint
with (
_patch_base_process() as mock_base,
patch(
"litellm.proxy.google_endpoints.endpoints.ProxyBaseLLMRequestProcessing.__init__",
return_value=None,
) as mock_init,
):
response = client.post(
"/v1beta/models/test-model:streamGenerateContent",
json={"contents": [{"role": "user", "parts": [{"text": "Hello"}]}]},
)
# Verify the response
assert response.status_code == 200
mock_base.assert_called_once()
kwargs = mock_base.call_args.kwargs
assert kwargs["route_type"] == "agenerate_content_stream"
assert kwargs["model"] == "test-model"
# Verify that agenerate_content_stream was called with correct parameters
mock_router.agenerate_content_stream.assert_called_once()
call_args = mock_router.agenerate_content_stream.call_args
assert call_args[1]["stream"] is True
assert call_args[1]["model"] == "test-model"
assert call_args[1]["contents"] == [
# stream=True must be forced into the data the processor receives.
init_kwargs = mock_init.call_args.kwargs
assert init_kwargs["data"]["stream"] is True
assert init_kwargs["data"]["model"] == "test-model"
assert init_kwargs["data"]["contents"] == [
{"role": "user", "parts": [{"text": "Hello"}]}
]
def test_google_generate_content_with_cost_tracking_metadata():
"""Test that the google_generate_content endpoint includes user metadata for cost tracking"""
def test_google_generate_content_data_flows_through_processor():
"""The body the client sends must reach ProxyBaseLLMRequestProcessing
intact so the pipeline can apply guardrails to it."""
try:
from fastapi import FastAPI
from fastapi.testclient import TestClient
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.google_endpoints.endpoints import router as google_router
client = _build_test_client()
except ImportError as e:
pytest.skip(f"Skipping test due to missing dependency: {e}")
# Create a FastAPI app and include the router (required for FastAPI 0.120+)
app = FastAPI()
app.include_router(google_router)
# Create a test client
client = TestClient(app)
# Mock all required proxy server dependencies
with (
patch("litellm.proxy.proxy_server.llm_router") as mock_router,
patch("litellm.proxy.proxy_server.general_settings", {}),
patch("litellm.proxy.proxy_server.proxy_config") as mock_proxy_config,
patch("litellm.proxy.proxy_server.version", "1.0.0"),
_patch_base_process(),
patch(
"litellm.proxy.litellm_pre_call_utils.add_litellm_data_to_request"
) as mock_add_data,
"litellm.proxy.google_endpoints.endpoints.ProxyBaseLLMRequestProcessing.__init__",
return_value=None,
) as mock_init,
):
mock_router.agenerate_content = AsyncMock(return_value={"test": "response"})
# Mock add_litellm_data_to_request to return data with metadata
async def mock_add_litellm_data(
data, request, user_api_key_dict, proxy_config, general_settings, version
):
# Simulate adding user metadata
data["litellm_metadata"] = {
"user_api_key_user_id": "test-user-id",
"user_api_key_team_id": "test-team-id",
"user_api_key": "hashed-key",
}
return data
mock_add_data.side_effect = mock_add_litellm_data
# Send a request to the endpoint
response = client.post(
client.post(
"/v1beta/models/test-model:generateContent",
json={"contents": [{"role": "user", "parts": [{"text": "Hello"}]}]},
headers={"Authorization": "Bearer sk-test-key"},
)
# Verify the response
assert response.status_code == 200
# Verify that add_litellm_data_to_request was called
mock_add_data.assert_called_once()
# Verify that agenerate_content was called with metadata
mock_router.agenerate_content.assert_called_once()
call_args = mock_router.agenerate_content.call_args
called_data = call_args[1]
# Verify that litellm_metadata exists and contains user information
assert "litellm_metadata" in called_data
assert called_data["litellm_metadata"]["user_api_key_user_id"] == "test-user-id"
assert called_data["litellm_metadata"]["user_api_key_team_id"] == "test-team-id"
def test_google_stream_generate_content_with_cost_tracking_metadata():
"""Test that the google_stream_generate_content endpoint includes user metadata for cost tracking"""
try:
from fastapi import FastAPI
from fastapi.testclient import TestClient
from litellm.proxy.google_endpoints.endpoints import router as google_router
except ImportError as e:
pytest.skip(f"Skipping test due to missing dependency: {e}")
# Create a FastAPI app and include the router (required for FastAPI 0.120+)
app = FastAPI()
app.include_router(google_router)
# Create a test client
client = TestClient(app)
# Mock the router's agenerate_content_stream method to return a stream
mock_stream = AsyncMock()
mock_stream.__aiter__ = lambda self: mock_stream
mock_stream.__anext__.side_effect = StopAsyncIteration
# Mock all required proxy server dependencies
with (
patch("litellm.proxy.proxy_server.llm_router") as mock_router,
patch("litellm.proxy.proxy_server.general_settings", {}),
patch("litellm.proxy.proxy_server.proxy_config") as mock_proxy_config,
patch("litellm.proxy.proxy_server.version", "1.0.0"),
patch(
"litellm.proxy.litellm_pre_call_utils.add_litellm_data_to_request"
) as mock_add_data,
):
mock_router.agenerate_content_stream = AsyncMock(return_value=mock_stream)
# Mock add_litellm_data_to_request to return data with metadata
async def mock_add_litellm_data(
data, request, user_api_key_dict, proxy_config, general_settings, version
):
# Simulate adding user metadata
data["litellm_metadata"] = {
"user_api_key_user_id": "test-user-id",
"user_api_key_team_id": "test-team-id",
"user_api_key": "hashed-key",
}
return data
mock_add_data.side_effect = mock_add_litellm_data
# Send a request to the endpoint
response = client.post(
"/v1beta/models/test-model:streamGenerateContent",
json={"contents": [{"role": "user", "parts": [{"text": "Hello"}]}]},
headers={"Authorization": "Bearer sk-test-key"},
)
# Verify the response
assert response.status_code == 200
# Verify that add_litellm_data_to_request was called
mock_add_data.assert_called_once()
# Verify that agenerate_content_stream was called with metadata
mock_router.agenerate_content_stream.assert_called_once()
call_args = mock_router.agenerate_content_stream.call_args
called_data = call_args[1]
# Verify that litellm_metadata exists and contains user information
assert "litellm_metadata" in called_data
assert called_data["litellm_metadata"]["user_api_key_user_id"] == "test-user-id"
assert called_data["litellm_metadata"]["user_api_key_team_id"] == "test-team-id"
# Verify stream is set to True
assert called_data["stream"] is True
def test_google_generate_content_with_system_instruction():
"""
Test that systemInstruction is correctly passed through from the endpoint to the router.
This test verifies the fix for systemInstruction being dropped when forwarding
requests to Vertex AI through the Google GenAI endpoint.
"""
try:
from fastapi import FastAPI
from fastapi.testclient import TestClient
from litellm.proxy.google_endpoints.endpoints import router as google_router
except ImportError as e:
pytest.skip(f"Skipping test due to missing dependency: {e}")
# Create a FastAPI app and include the router
app = FastAPI()
app.include_router(google_router)
# Create a test client
client = TestClient(app)
# Mock all required proxy server dependencies
with (
patch("litellm.proxy.proxy_server.llm_router") as mock_router,
patch("litellm.proxy.proxy_server.general_settings", {}),
patch("litellm.proxy.proxy_server.proxy_config") as mock_proxy_config,
patch("litellm.proxy.proxy_server.version", "1.0.0"),
patch(
"litellm.proxy.litellm_pre_call_utils.add_litellm_data_to_request"
) as mock_add_data,
):
mock_router.agenerate_content = AsyncMock(return_value={"test": "response"})
# Mock add_litellm_data_to_request to pass through data unchanged
async def mock_add_litellm_data(
data, request, user_api_key_dict, proxy_config, general_settings, version
):
return data
mock_add_data.side_effect = mock_add_litellm_data
# Define the systemInstruction to test
system_instruction = {"parts": [{"text": "Your name is Doodle."}]}
# Send a request with systemInstruction
response = client.post(
"/v1beta/models/gemini-2.5-pro:generateContent",
json={
"systemInstruction": system_instruction,
"contents": [
{"parts": [{"text": "What is your name?"}], "role": "user"}
],
},
headers={"Authorization": "Bearer sk-test-key"},
)
# Verify the response
assert response.status_code == 200
# Verify that agenerate_content was called
mock_router.agenerate_content.assert_called_once()
call_args = mock_router.agenerate_content.call_args
called_data = call_args[1]
# Verify that systemInstruction is present in the call arguments
assert "systemInstruction" in called_data
assert called_data["systemInstruction"] == system_instruction
assert (
called_data["systemInstruction"]["parts"][0]["text"]
== "Your name is Doodle."
)
# Verify contents are also present
assert "contents" in called_data
assert len(called_data["contents"]) == 1
assert called_data["contents"][0]["role"] == "user"
def test_google_generate_content_with_image_config():
"""
Test that imageConfig is correctly passed through from generationConfig to the router.
This test verifies that imageConfig parameters (aspectRatio, imageSize) are preserved
when forwarding requests to Google GenAI through the endpoint.
"""
try:
from fastapi import FastAPI
from fastapi.testclient import TestClient
from litellm.proxy.google_endpoints.endpoints import router as google_router
except ImportError as e:
pytest.skip(f"Skipping test due to missing dependency: {e}")
# Create a FastAPI app and include the router
app = FastAPI()
app.include_router(google_router)
# Create a test client
client = TestClient(app)
# Mock all required proxy server dependencies
with (
patch("litellm.proxy.proxy_server.llm_router") as mock_router,
patch("litellm.proxy.proxy_server.general_settings", {}),
patch("litellm.proxy.proxy_server.proxy_config") as mock_proxy_config,
patch("litellm.proxy.proxy_server.version", "1.0.0"),
patch(
"litellm.proxy.litellm_pre_call_utils.add_litellm_data_to_request"
) as mock_add_data,
):
mock_router.agenerate_content = AsyncMock(return_value={"test": "response"})
# Mock add_litellm_data_to_request to pass through data unchanged
async def mock_add_litellm_data(
data, request, user_api_key_dict, proxy_config, general_settings, version
):
return data
mock_add_data.side_effect = mock_add_litellm_data
# Send a request with generationConfig containing imageConfig
response = client.post(
"/v1beta/models/gemini-3-pro-image-preview:generateContent",
json={
"contents": [
{
"role": "user",
"parts": [
{
"text": "Create a vibrant infographic about photosynthesis"
}
],
}
],
"contents": [{"role": "user", "parts": [{"text": "Hello"}]}],
"systemInstruction": {"parts": [{"text": "Your name is Doodle."}]},
"generationConfig": {
"responseModalities": ["TEXT", "IMAGE"],
"imageConfig": {"aspectRatio": "9:16", "imageSize": "4K"},
},
},
headers={"Authorization": "Bearer sk-test-key"},
)
# Verify the response
assert response.status_code == 200
# Verify that agenerate_content was called
mock_router.agenerate_content.assert_called_once()
call_args = mock_router.agenerate_content.call_args
called_data = call_args[1]
# Verify that config is present in the call arguments
assert "config" in called_data
# Verify that imageConfig is preserved in the config
assert "imageConfig" in called_data["config"]
assert called_data["config"]["imageConfig"]["aspectRatio"] == "9:16"
assert called_data["config"]["imageConfig"]["imageSize"] == "4K"
# Verify that responseModalities is also preserved
assert "responseModalities" in called_data["config"]
assert called_data["config"]["responseModalities"] == ["TEXT", "IMAGE"]
# Verify contents are also present
assert "contents" in called_data
assert len(called_data["contents"]) == 1
assert called_data["contents"][0]["role"] == "user"
data = mock_init.call_args.kwargs["data"]
assert data["model"] == "test-model"
assert data["contents"] == [{"role": "user", "parts": [{"text": "Hello"}]}]
assert data["systemInstruction"] == {
"parts": [{"text": "Your name is Doodle."}]
}
# generationConfig arrives intact here; the rename to `config` is
# done downstream in route_request (see test_route_llm_request).
assert data["generationConfig"]["responseModalities"] == ["TEXT", "IMAGE"]
assert data["generationConfig"]["imageConfig"]["aspectRatio"] == "9:16"
def test_google_generate_content_metadata_and_trace_id_callbacks():
"""Test that google_generate_content sets litellm_call_id and logging_obj for callbacks (e.g. S3, Langfuse)"""
def test_google_generate_content_forwards_call_id_header():
"""The endpoint must forward the x-litellm-call-id header to the processor
so the helper can stamp it on the logging object. Trace continuity from
client → callbacks (S3, Langfuse, etc.) depends on this header surviving
the hop through these endpoints."""
try:
from fastapi import FastAPI
from fastapi.testclient import TestClient
from litellm.proxy.google_endpoints.endpoints import router as google_router
client = _build_test_client()
except ImportError as e:
pytest.skip(f"Skipping test due to missing dependency: {e}")
# Create a FastAPI app and include the router
app = FastAPI()
app.include_router(google_router)
# Create a test client
client = TestClient(app)
# Mock all required proxy server dependencies
with (
patch("litellm.proxy.proxy_server.llm_router") as mock_router,
patch("litellm.proxy.proxy_server.general_settings", {}),
patch("litellm.proxy.proxy_server.proxy_config") as mock_proxy_config,
patch("litellm.proxy.proxy_server.version", "1.0.0"),
patch(
"litellm.proxy.litellm_pre_call_utils.add_litellm_data_to_request"
) as mock_add_data,
):
mock_router.agenerate_content = AsyncMock(return_value={"test": "response"})
# Mock add_litellm_data_to_request to return data with metadata
async def mock_add_litellm_data(
data, request, user_api_key_dict, proxy_config, general_settings, version
):
# Simulate adding user metadata
data["litellm_metadata"] = {
"user_api_key_user_id": "test-user-id",
}
return data
mock_add_data.side_effect = mock_add_litellm_data
# Send a request to the endpoint with x-litellm-call-id header
test_call_id = "test-custom-call-id"
response = client.post(
with _patch_base_process() as mock_base:
client.post(
"/v1beta/models/test-model:generateContent",
json={"contents": [{"role": "user", "parts": [{"text": "Hello"}]}]},
headers={
"Authorization": "Bearer sk-test-key",
"x-litellm-call-id": test_call_id,
},
headers={"x-litellm-call-id": "trace-abc-123"},
)
assert response.status_code == 200
mock_router.agenerate_content.assert_called_once()
call_args = mock_router.agenerate_content.call_args
called_data = call_args[1]
# Verify that the litellm_logging_obj got assigned in the final called_data to router
assert "litellm_logging_obj" in called_data
assert "litellm_call_id" in called_data
assert called_data["litellm_call_id"] == test_call_id
forwarded_request = mock_base.call_args.kwargs["request"]
assert forwarded_request.headers.get("x-litellm-call-id") == "trace-abc-123"
def test_google_stream_generate_content_metadata_and_trace_id_callbacks():
"""Test that google_stream_generate_content sets litellm_call_id and logging_obj for callbacks"""
def test_google_count_tokens_unchanged():
"""countTokens has its own path and isn't affected by the pipeline change."""
try:
from fastapi import FastAPI
from fastapi.testclient import TestClient
from litellm.proxy.google_endpoints.endpoints import router as google_router
client = _build_test_client()
except ImportError as e:
pytest.skip(f"Skipping test due to missing dependency: {e}")
app = FastAPI()
app.include_router(google_router)
client = TestClient(app)
fake_response = MagicMock()
fake_response.original_response = {
"totalTokens": 7,
"promptTokensDetails": [{"modality": "TEXT", "tokenCount": 7}],
}
fake_response.total_tokens = 7
mock_stream = AsyncMock()
mock_stream.__aiter__ = lambda self: mock_stream
mock_stream.__anext__.side_effect = StopAsyncIteration
with (
patch("litellm.proxy.proxy_server.llm_router") as mock_router,
patch("litellm.proxy.proxy_server.general_settings", {}),
patch("litellm.proxy.proxy_server.proxy_config") as mock_proxy_config,
patch("litellm.proxy.proxy_server.version", "1.0.0"),
patch(
"litellm.proxy.litellm_pre_call_utils.add_litellm_data_to_request"
) as mock_add_data,
with patch(
"litellm.proxy.proxy_server.token_counter",
new_callable=AsyncMock,
return_value=fake_response,
):
mock_router.agenerate_content_stream = AsyncMock(return_value=mock_stream)
async def mock_add_litellm_data(
data, request, user_api_key_dict, proxy_config, general_settings, version
):
data["litellm_metadata"] = {
"user_api_key_user_id": "test-user-id",
}
return data
mock_add_data.side_effect = mock_add_litellm_data
test_call_id = "test-custom-stream-call-id"
response = client.post(
"/v1beta/models/test-model:streamGenerateContent",
json={"contents": [{"role": "user", "parts": [{"text": "Hello stream"}]}]},
headers={
"Authorization": "Bearer sk-test-key",
"x-litellm-call-id": test_call_id,
},
"/v1beta/models/test-model:countTokens",
json={"contents": [{"role": "user", "parts": [{"text": "Hello"}]}]},
)
assert response.status_code == 200
mock_router.agenerate_content_stream.assert_called_once()
call_args = mock_router.agenerate_content_stream.call_args
called_data = call_args[1]
assert "litellm_logging_obj" in called_data
assert "litellm_call_id" in called_data
assert called_data["litellm_call_id"] == test_call_id
body = response.json()
assert body["totalTokens"] == 7

57
tests/test_litellm/proxy/management_endpoints/test_access_group_endpoints.py
View file

@ -738,15 +738,32 @@ def test_delete_access_group_patches_cached_team_and_key(
return_value=None
)
# Build cached key object (returned from user_api_key_cache)
if key_cache_group_ids is not None:
cached_key = UserAPIKeyAuth(
token="hashed-key-1",
access_group_ids=list(key_cache_group_ids),
# user_api_key_cache is queried both for teams (fallback after dual_cache) and
# hashed keys — return the right stub per ``key``. A single AsyncMock(return_value=key)
# would wrongly serve the key blob for ``team_id:team-1`` and trigger team patching.
# Use a synchronous side_effect (not async def): AsyncMock awaits coroutine side_effects
# inconsistently across Python/unittest versions; sync returns are awaited as immediate results.
def user_cache_get_side_effect(*args, **kwargs):
cache_key = (
kwargs.get("key") if "key" in kwargs else (args[0] if args else None)
)
mock_cache.async_get_cache = AsyncMock(return_value=cached_key)
else:
mock_cache.async_get_cache = AsyncMock(return_value=None)
if cache_key == "team_id:team-1":
if team_cache_group_ids is None:
return None
return LiteLLM_TeamTableCachedObj(
team_id="team-1",
access_group_ids=list(team_cache_group_ids),
)
if cache_key == "hashed-key-1":
if key_cache_group_ids is None:
return None
return UserAPIKeyAuth(
token="hashed-key-1",
access_group_ids=list(key_cache_group_ids),
)
return None
mock_cache.async_get_cache = AsyncMock(side_effect=user_cache_get_side_effect)
resp = client.delete("/v1/access_group/ag-to-delete")
assert resp.status_code == 204
@ -803,7 +820,7 @@ def test_delete_access_group_patches_cached_team_and_key(
def test_delete_access_group_patches_key_cached_as_dict(client_and_mocks):
"""Delete correctly patches a key cached as a raw dict (not UserAPIKeyAuth)."""
"""Delete patches key cache — mock returns UserAPIKeyAuth (what UserApiKeyCache emits after deserialize)."""
client, mock_prisma, mock_access_group_table, mock_cache, mock_proxy_logging = (
client_and_mocks
)
@ -826,12 +843,24 @@ def test_delete_access_group_patches_key_cached_as_dict(client_and_mocks):
return_value=None
)
# Key cached as a plain dict (as can happen with Redis serialization)
# Serialized shape from Redis dict; UserApiKeyCache.async_get_cache(model_type=...) yields a model — simulate that.
cached_key_payload = {
"token": "hashed-key-dict",
"access_group_ids": ["ag-to-delete", "ag-other"],
}
def user_cache_get_dict_when_key_matches(*args, **kwargs):
cache_key = (
kwargs.get("key") if "key" in kwargs else (args[0] if args else None)
)
if cache_key == "team_id:team-1":
return None
if cache_key == "hashed-key-dict":
return UserAPIKeyAuth.model_validate(cached_key_payload)
return None
mock_cache.async_get_cache = AsyncMock(
return_value={
"token": "hashed-key-dict",
"access_group_ids": ["ag-to-delete", "ag-other"],
}
side_effect=user_cache_get_dict_when_key_matches
)
resp = client.delete("/v1/access_group/ag-to-delete")

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

@ -5512,6 +5512,9 @@ async def test_update_team_guardrails_with_org_id():
return_value=mock_updated_team
)
mock_prisma.jsonify_team_object = MagicMock(side_effect=lambda db_data: db_data)
# async_get_cache must be an AsyncMock so `await` in get_org_object works
mock_cache.async_get_cache = AsyncMock(return_value=None)
mock_cache.async_set_cache = AsyncMock()
# Mock llm_router
mock_router = MagicMock()

236
tests/test_litellm/proxy/test_filter_models_by_team_access_group.py Normal file
View file

@ -0,0 +1,236 @@
"""
Tests for _filter_models_by_team_id resolving access group names.
Verifies that when a team's `models` field contains an access group name
(e.g., "Group-A"), the filter resolves it to the member model names before
looking up deployments — matching the behavior of the auth path in
auth_checks.py:model_in_access_group().
"""
import os
import sys
from unittest.mock import AsyncMock, MagicMock
import pytest
sys.path.insert(0, os.path.abspath("../../.."))
from litellm.proxy.proxy_server import _filter_models_by_team_id
def _make_model(model_name: str, model_id: str, access_groups: list[str] = None):
"""Helper to build a model dict matching the router's format."""
return {
"model_name": model_name,
"litellm_params": {"model": model_name},
"model_info": {
"id": model_id,
"access_groups": access_groups or [],
},
}
def _make_team(models: list[str], team_id: str = "team_alpha"):
"""Helper to build a mock team DB object."""
mock = MagicMock()
mock.model_dump.return_value = {
"team_id": team_id,
"team_alias": "Team Alpha",
"models": models,
"max_budget": None,
"spend": 0.0,
"blocked": False,
"members_with_roles": [],
"metadata": {},
}
return mock
@pytest.mark.asyncio
async def test_filter_resolves_access_group_names():
"""
When team.models contains an access group name, _filter_models_by_team_id
should resolve it to the member models and return only those deployments.
"""
# Models on the proxy
gpt4o = _make_model("gpt-4o", "id-1", ["Group-A"])
gpt5 = _make_model("gpt-5", "id-2", ["Group-A"])
claude = _make_model("claude-3", "id-3", ["Group-B"])
all_models = [gpt4o, gpt5, claude]
# Router mock
mock_router = MagicMock()
# get_model_access_groups returns {group_name: [model_names]}
mock_router.get_model_access_groups.return_value = {
"Group-A": ["gpt-4o", "gpt-5"],
"Group-B": ["claude-3"],
}
# get_model_list returns deployments matching a model_name
def fake_get_model_list(model_name=None, team_id=None):
return [m for m in all_models if m["model_name"] == model_name]
mock_router.get_model_list = MagicMock(side_effect=fake_get_model_list)
# Team has models: ["Group-A"] — an access group name, not a literal model
team_db = _make_team(models=["Group-A"])
# Prisma mock
mock_prisma = MagicMock()
mock_prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=team_db)
mock_prisma.db.litellm_proxymodeltable.find_many = AsyncMock(return_value=[])
result = await _filter_models_by_team_id(
all_models=all_models,
team_id="team_alpha",
prisma_client=mock_prisma,
llm_router=mock_router,
)
result_ids = {m["model_info"]["id"] for m in result}
# Should include gpt-4o and gpt-5 (Group-A), but NOT claude-3 (Group-B)
assert result_ids == {
"id-1",
"id-2",
}, f"Expected Group-A models only, got {result_ids}"
# Verify DB fallback query received resolved model names, not access group name
call_kwargs = mock_prisma.db.litellm_proxymodeltable.find_many.call_args[1]
assert set(call_kwargs["where"]["model_name"]["in"]) == {
"gpt-4o",
"gpt-5",
}, "find_many should receive resolved model names, not the access group name"
@pytest.mark.asyncio
async def test_filter_resolves_mix_of_access_groups_and_literal_names():
"""
When team.models contains both an access group name and a literal model name,
both should be resolved correctly.
"""
gpt4o = _make_model("gpt-4o", "id-1", ["Group-A"])
gpt5 = _make_model("gpt-5", "id-2", ["Group-A"])
claude = _make_model("claude-3", "id-3", ["Group-B"])
mistral = _make_model("mistral-large", "id-4", []) # no access group
all_models = [gpt4o, gpt5, claude, mistral]
mock_router = MagicMock()
mock_router.get_model_access_groups.return_value = {
"Group-A": ["gpt-4o", "gpt-5"],
"Group-B": ["claude-3"],
}
def fake_get_model_list(model_name=None, team_id=None):
return [m for m in all_models if m["model_name"] == model_name]
mock_router.get_model_list = MagicMock(side_effect=fake_get_model_list)
# Team has access to Group-A (access group) + mistral-large (literal name)
team_db = _make_team(models=["Group-A", "mistral-large"])
mock_prisma = MagicMock()
mock_prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=team_db)
mock_prisma.db.litellm_proxymodeltable.find_many = AsyncMock(return_value=[])
result = await _filter_models_by_team_id(
all_models=all_models,
team_id="team_alpha",
prisma_client=mock_prisma,
llm_router=mock_router,
)
result_ids = {m["model_info"]["id"] for m in result}
# Group-A models + mistral-large, but NOT claude-3
assert result_ids == {
"id-1",
"id-2",
"id-4",
}, f"Expected Group-A + mistral-large, got {result_ids}"
@pytest.mark.asyncio
async def test_filter_excludes_models_from_other_access_group():
"""
Models belonging only to a different access group must not appear in results.
"""
gpt4o = _make_model("gpt-4o", "id-1", ["Group-A"])
claude = _make_model("claude-3", "id-3", ["Group-B"])
llama = _make_model("llama-4", "id-4", ["Group-B"])
all_models = [gpt4o, claude, llama]
mock_router = MagicMock()
mock_router.get_model_access_groups.return_value = {
"Group-A": ["gpt-4o"],
"Group-B": ["claude-3", "llama-4"],
}
def fake_get_model_list(model_name=None, team_id=None):
return [m for m in all_models if m["model_name"] == model_name]
mock_router.get_model_list = MagicMock(side_effect=fake_get_model_list)
team_db = _make_team(models=["Group-A"])
mock_prisma = MagicMock()
mock_prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=team_db)
mock_prisma.db.litellm_proxymodeltable.find_many = AsyncMock(return_value=[])
result = await _filter_models_by_team_id(
all_models=all_models,
team_id="team_alpha",
prisma_client=mock_prisma,
llm_router=mock_router,
)
result_names = {m["model_name"] for m in result}
assert "claude-3" not in result_names, "Group-B model should not be accessible"
assert "llama-4" not in result_names, "Group-B model should not be accessible"
assert "gpt-4o" in result_names, "Group-A model should be accessible"
@pytest.mark.asyncio
async def test_filter_db_fallback_receives_resolved_model_names():
"""
When get_model_list returns no results (forcing the DB fallback path),
the DB query should receive resolved model names, not the raw access group name.
"""
gpt4o = _make_model("gpt-4o", "id-1", ["Group-A"])
all_models = [gpt4o]
mock_router = MagicMock()
mock_router.get_model_access_groups.return_value = {
"Group-A": ["gpt-4o", "gpt-5"],
}
# get_model_list returns nothing — forces reliance on the DB fallback
mock_router.get_model_list = MagicMock(return_value=[])
team_db = _make_team(models=["Group-A"])
# DB returns a model that the router didn't find
mock_db_model = MagicMock()
mock_db_model.model_id = "id-db-1"
mock_prisma = MagicMock()
mock_prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=team_db)
mock_prisma.db.litellm_proxymodeltable.find_many = AsyncMock(
return_value=[mock_db_model]
)
result = await _filter_models_by_team_id(
all_models=all_models,
team_id="team_alpha",
prisma_client=mock_prisma,
llm_router=mock_router,
)
# Verify DB query received resolved names, not "Group-A"
call_kwargs = mock_prisma.db.litellm_proxymodeltable.find_many.call_args[1]
queried_names = set(call_kwargs["where"]["model_name"]["in"])
assert queried_names == {
"gpt-4o",
"gpt-5",
}, f"DB query should receive resolved model names, got {queried_names}"
assert "Group-A" not in queried_names, "Raw access group name should not be in DB query"

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

@ -5750,3 +5750,90 @@ class TestLazyFeatureMiddleware:
assert attempts == [
"called"
], f"failing register_fn should be invoked once, not on every request; got {attempts}"
@pytest.mark.asyncio
async def test_get_current_spend_redis_clean_miss_skips_stale_in_memory():
"""When Redis is reachable and cleanly returns None (TTL expired,
counter genuinely absent), the read must reseed from DB - NOT fall
through to per-pod in-memory which only contains this pod's writes.
Pre-fix in multi-pod deployments, in-memory contained a stale local
subset (e.g. $30) while DB had the true cross-pod total ($500). The
fall-through returned $30, enforcement passed, bypass.
"""
from litellm.caching.dual_cache import DualCache
from litellm.proxy.proxy_server import get_current_spend
counter_cache = DualCache()
counter_key = "spend:team_member:user-1:team-1"
# Per-pod stale in-memory: only this pod's writes, not cross-pod truth.
counter_cache.in_memory_cache.set_cache(key=counter_key, value=30.0)
# Redis cleanly returns None (key expired or never written on this pod).
fake_redis = AsyncMock()
fake_redis.async_get_cache = AsyncMock(return_value=None)
fake_redis.async_increment = AsyncMock(return_value=500.0)
counter_cache.redis_cache = fake_redis
# DB has the authoritative cross-pod spend.
db_row = MagicMock()
db_row.spend = 500.0
fake_prisma = MagicMock()
fake_prisma.db.litellm_teammembership.find_unique = AsyncMock(return_value=db_row)
import litellm.proxy.proxy_server as ps
orig_counter, orig_prisma = ps.spend_counter_cache, ps.prisma_client
ps.spend_counter_cache = counter_cache
ps.prisma_client = fake_prisma
try:
spend = await get_current_spend(counter_key=counter_key, fallback_spend=0.0)
assert spend == 500.0, (
f"expected DB-authoritative 500.0 on clean Redis miss, got {spend} "
f"(stale per-pod in-memory $30 would have caused multi-pod bypass)"
)
finally:
ps.spend_counter_cache = orig_counter
ps.prisma_client = orig_prisma
@pytest.mark.asyncio
async def test_get_current_spend_redis_error_falls_back_to_in_memory():
"""When Redis raises, the read should still degrade to in-memory rather
than going straight to DB - in-memory is at least same-pod-fresh and
cheaper than a DB query during a Redis outage."""
from litellm.caching.dual_cache import DualCache
from litellm.proxy.proxy_server import get_current_spend
counter_cache = DualCache()
counter_key = "spend:team_member:user-1:team-1"
counter_cache.in_memory_cache.set_cache(key=counter_key, value=42.0)
fake_redis = AsyncMock()
fake_redis.async_get_cache = AsyncMock(side_effect=ConnectionError("redis down"))
counter_cache.redis_cache = fake_redis
fake_prisma = MagicMock()
fake_prisma.db.litellm_teammembership.find_unique = AsyncMock(
return_value=MagicMock(spend=999.0)
)
import litellm.proxy.proxy_server as ps
orig_counter, orig_prisma = ps.spend_counter_cache, ps.prisma_client
ps.spend_counter_cache = counter_cache
ps.prisma_client = fake_prisma
try:
spend = await get_current_spend(counter_key=counter_key, fallback_spend=0.0)
assert spend == 42.0, (
f"expected in-memory fallback 42.0 on Redis error, got {spend} "
f"(should not have hit DB when Redis errored)"
)
# DB query should NOT have fired - in-memory short-circuits.
fake_prisma.db.litellm_teammembership.find_unique.assert_not_awaited()
finally:
ps.spend_counter_cache = orig_counter
ps.prisma_client = orig_prisma

145
tests/test_litellm/proxy/test_redis_auth_cache_flag.py Normal file
View file

@ -0,0 +1,145 @@
"""
Tests for the enable_redis_auth_cache litellm_settings flag.
Verifies that _init_cache attaches Redis to user_api_key_cache only when
the flag is explicitly set to True, and leaves it in-memory-only otherwise.
"""
from contextlib import contextmanager
import json
from unittest.mock import MagicMock, patch
import pytest
import litellm
import litellm.proxy.proxy_server as ps
from litellm.caching.caching import RedisCache
from litellm.caching.dual_cache import DualCache
# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
class _FakeRedisCache(RedisCache):
"""
Minimal RedisCache subclass that passes isinstance checks without
requiring a real Redis connection. __init__ is bypassed so no
network calls are made.
"""
def __init__(self): # noqa: super().__init__ skipped intentionally
self._store = {}
def set_cache(self, key, value, **kwargs): # type: ignore[override]
# Enforce Redis JSON-serializable payload contract.
self._store[key] = json.dumps(value)
return True
def get_cache(self, key, **kwargs): # type: ignore[override]
raw = self._store.get(key)
if raw is None:
return None
return json.loads(raw)
@contextmanager
def _patched_init_cache(litellm_settings: dict, cache_params: dict):
"""
Context manager that:
1. Replaces the module-level globals with fresh DualCache instances.
2. Patches ``litellm.Cache`` (locally imported inside _init_cache) so
it returns a fake cache whose ``.cache`` attribute is a
_FakeRedisCache (passes the isinstance guard in _init_cache).
3. Extracts enable_redis_auth_cache from litellm_settings and passes it
as the second argument to _init_cache (matching production behaviour).
4. Yields (user_api_key_cache, spend_counter_cache) after calling
_init_cache, then restores everything.
"""
fake_redis = _FakeRedisCache()
mock_litellm_cache = MagicMock()
mock_litellm_cache.cache = fake_redis
fresh_user_cache = DualCache()
fresh_spend_cache = DualCache()
enable_redis_auth_cache = litellm_settings.get("enable_redis_auth_cache", False)
with (
patch.object(ps, "user_api_key_cache", fresh_user_cache),
patch.object(ps, "spend_counter_cache", fresh_spend_cache),
patch.object(ps, "llm_router", None),
# Cache is locally imported inside _init_cache: patch it at source.
patch("litellm.Cache", return_value=mock_litellm_cache),
):
litellm.cache = None
ps.ProxyConfig()._init_cache(cache_params, enable_redis_auth_cache)
yield fresh_user_cache, fresh_spend_cache
# ---------------------------------------------------------------------------
# Tests
# ---------------------------------------------------------------------------
class TestRedisAuthCacheFlag:
def test_flag_true_attaches_redis_to_user_api_key_cache(self):
"""When enable_redis_auth_cache=True, user_api_key_cache.redis_cache must be set."""
with _patched_init_cache(
litellm_settings={"enable_redis_auth_cache": True},
cache_params={"type": "redis", "host": "localhost", "port": 6379},
) as (user_cache, _):
assert user_cache.redis_cache is not None, (
"Redis should be attached to user_api_key_cache when "
"enable_redis_auth_cache=True"
)
def test_flag_false_leaves_user_api_key_cache_in_memory_only(self):
"""When enable_redis_auth_cache=False, user_api_key_cache must stay in-memory."""
with _patched_init_cache(
litellm_settings={"enable_redis_auth_cache": False},
cache_params={"type": "redis", "host": "localhost", "port": 6379},
) as (user_cache, _):
assert user_cache.redis_cache is None, (
"user_api_key_cache must remain in-memory-only when "
"enable_redis_auth_cache=False"
)
def test_flag_absent_leaves_user_api_key_cache_in_memory_only(self):
"""When enable_redis_auth_cache is not set at all, default is in-memory-only."""
with _patched_init_cache(
litellm_settings={},
cache_params={"type": "redis", "host": "localhost", "port": 6379},
) as (user_cache, _):
assert user_cache.redis_cache is None, (
"user_api_key_cache must remain in-memory-only when "
"enable_redis_auth_cache is absent from litellm_settings"
)
def test_spend_counter_cache_always_gets_redis_regardless_of_flag(self):
"""spend_counter_cache must receive Redis regardless of the auth-cache flag."""
for flag_value in (True, False, None):
ls = (
{"enable_redis_auth_cache": flag_value}
if flag_value is not None
else {}
)
with _patched_init_cache(
litellm_settings=ls,
cache_params={"type": "redis", "host": "localhost", "port": 6379},
) as (_, spend_cache):
assert spend_cache.redis_cache is not None, (
f"spend_counter_cache must always get Redis "
f"(enable_redis_auth_cache={flag_value!r})"
)
def test_flag_false_spend_gets_redis_but_user_cache_does_not(self):
"""Explicit False: spend cache wired, auth cache left in-memory."""
with _patched_init_cache(
litellm_settings={"enable_redis_auth_cache": False},
cache_params={"type": "redis", "host": "localhost", "port": 6379},
) as (user_cache, spend_cache):
assert spend_cache.redis_cache is not None
assert user_cache.redis_cache is None

52
tests/test_litellm/proxy/test_route_llm_request.py
View file

@ -239,3 +239,55 @@ async def test_route_request_with_router_settings_override_preserves_existing():
assert call_kwargs["num_retries"] == 10
# Key/team timeout should be applied since not in request
assert call_kwargs["timeout"] == 30
@pytest.mark.parametrize(
"route_type", ["agenerate_content", "agenerate_content_stream"]
)
@pytest.mark.asyncio
async def test_route_request_maps_generation_config_for_google_routes(route_type):
"""For Google generate_content routes, route_request must rename
`generationConfig` (Google's wire format) to `config` (the kwarg the
router method expects). Without this mapping the request reaches the
LLM with the field under the wrong name and the config is dropped."""
data = {
"model": "gemini-2.5-flash",
"contents": [{"role": "user", "parts": [{"text": "Hello"}]}],
"generationConfig": {
"responseModalities": ["TEXT", "IMAGE"],
"imageConfig": {"aspectRatio": "9:16", "imageSize": "4K"},
},
}
llm_router = MagicMock()
getattr(llm_router, route_type).return_value = "ok"
await route_request(data, llm_router, None, route_type)
call_kwargs = getattr(llm_router, route_type).call_args[1]
assert "generationConfig" not in call_kwargs
assert "config" in call_kwargs
assert call_kwargs["config"]["responseModalities"] == ["TEXT", "IMAGE"]
assert call_kwargs["config"]["imageConfig"]["aspectRatio"] == "9:16"
assert call_kwargs["config"]["imageConfig"]["imageSize"] == "4K"
@pytest.mark.parametrize(
"route_type", ["agenerate_content", "agenerate_content_stream"]
)
@pytest.mark.asyncio
async def test_route_request_preserves_existing_config_for_google_routes(route_type):
"""If the caller already supplies `config`, route_request must not
overwrite it with `generationConfig`."""
data = {
"model": "gemini-2.5-flash",
"contents": [{"role": "user", "parts": [{"text": "Hello"}]}],
"config": {"existing": True},
"generationConfig": {"shouldNotWin": True},
}
llm_router = MagicMock()
getattr(llm_router, route_type).return_value = "ok"
await route_request(data, llm_router, None, route_type)
call_kwargs = getattr(llm_router, route_type).call_args[1]
assert call_kwargs["config"] == {"existing": True}

28
tests/test_litellm/router_strategy/test_router_tag_routing.py
View file

@ -346,6 +346,34 @@ def test_tag_routing_with_list_of_tags_match_all():
assert not is_valid_deployment_tag(["default"], ["teamA"], match_any=False)
def test_strict_tag_routing_without_request_tags_blocks_header_regex_fallback():
"""
When tag_filtering_match_any=False, deployments with plain tags must require
those request tags before header regex can match. A spoofed User-Agent must
not route to a tagged deployment when the request has no tags.
"""
from litellm.router_strategy.tag_based_routing import _match_deployment
deployment = {
"model_name": "restricted-model",
"litellm_params": {
"model": "gpt-4o",
"tags": ["internal"],
"tag_regex": ["^User-Agent: internal-tool"],
},
}
assert (
_match_deployment(
deployment=deployment,
request_tags=None,
header_strings=["User-Agent: internal-tool"],
match_any=False,
)
is None
)
@pytest.mark.asyncio()
async def test_router_free_paid_tier_with_responses_api():
"""

110
tests/test_litellm/router_utils/test_router_utils_common_utils.py
View file

@ -6,6 +6,7 @@ import pytest
from litellm import Router
from litellm.router_utils.common_utils import (
_deployment_supports_web_search,
add_model_file_id_mappings,
filter_team_based_models,
filter_web_search_deployments,
)
@ -362,3 +363,112 @@ def test_invalidate_model_group_info_cache():
# Invalidate and verify cache is cleared
router._invalidate_model_group_info_cache()
assert router._cached_get_model_group_info.cache_info().currsize == 0
class TestAddModelFileIdMappings:
"""Test cases for add_model_file_id_mappings.
The router may pass either a list of deployment dicts (multiple matched
deployments) or a single deployment dict (when a specific deployment was
resolved, e.g. because the requested model matched a `model_info.id`).
Both shapes must produce a `{model_id: file_id}` mapping by extracting
`model_info.id` from each deployment.
"""
@staticmethod
def _make_response(file_id: str):
response = Mock()
response.id = file_id
return response
def test_should_map_each_deployment_id_when_given_list(self):
deployments = [
{
"model_name": "gpt-4",
"litellm_params": {"model": "gpt-4"},
"model_info": {"id": "deployment-1"},
},
{
"model_name": "gpt-4",
"litellm_params": {"model": "gpt-4"},
"model_info": {"id": "deployment-2"},
},
]
responses = [self._make_response("file-1"), self._make_response("file-2")]
result = add_model_file_id_mappings(deployments, responses)
assert result == {"deployment-1": "file-1", "deployment-2": "file-2"}
def test_should_extract_model_info_id_when_given_single_deployment_dict(self):
"""Regression test: when `_common_checks_available_deployment` resolves
a specific deployment (returned as a dict, not a list), the function
must still extract `model_info.id` rather than iterate over the
deployment's own keys (`model_name`, `litellm_params`, `model_info`).
"""
deployment = {
"model_name": "gpt-4",
"litellm_params": {"model": "gpt-4", "api_key": "sk-test"},
"model_info": {"id": "deployment-1", "mode": "chat"},
}
responses = [self._make_response("file-1")]
result = add_model_file_id_mappings(deployment, responses)
assert result == {"deployment-1": "file-1"}
assert all(isinstance(v, str) for v in result.values())
def test_should_handle_batch_model_when_id_matches_model_name(self):
"""Regression test for the batch-model case: when `model_info.id` is
intentionally set equal to `model_name`, the router resolves a single
deployment via `has_model_id` and returns it as a dict. The mapping
must contain only `{id: file_id}` with string values so the resulting
`LiteLLM_ManagedFileTable` Pydantic validation passes.
"""
deployment = {
"model_name": "openai/openai/gpt-5.5-batch",
"litellm_params": {
"model": "openai/gpt-5.5",
"api_key": "sk-test",
"tpm": 40000000,
"rpm": 15000,
},
"model_info": {
"id": "openai/openai/gpt-5.5-batch",
"mode": "batch",
"base_model": "gpt-5.5",
"access_groups": ["default-models"],
},
}
responses = [self._make_response("file-batch-1")]
result = add_model_file_id_mappings(deployment, responses)
# Bug case would have produced keys ["model_name", "litellm_params",
# "model_info"] with non-string values.
assert result == {"openai/openai/gpt-5.5-batch": "file-batch-1"}
assert "litellm_params" not in result
assert "model_info" not in result
def test_should_skip_deployment_when_model_info_id_missing(self):
deployments = [
{
"model_name": "gpt-4",
"litellm_params": {"model": "gpt-4"},
"model_info": {},
},
{
"model_name": "gpt-4",
"litellm_params": {"model": "gpt-4"},
"model_info": {"id": "deployment-2"},
},
]
responses = [self._make_response("file-1"), self._make_response("file-2")]
result = add_model_file_id_mappings(deployments, responses)
assert result == {"deployment-2": "file-2"}
def test_should_return_empty_mapping_when_given_empty_list(self):
result = add_model_file_id_mappings([], [])
assert result == {}

8
ui/litellm-dashboard/e2e_tests/tests/modelsPage/addModel.spec.ts
View file

@ -141,7 +141,9 @@ test.describe("Add Model", () => {
await page.waitForTimeout(1000);
// Verify the model appears in the results count (not "Showing 0 results")
await expect(page.getByText(/Showing \d+ - \d+ of \d+ results/)).toBeVisible({ timeout: 15_000 });
await expect(page.getByTestId("models-results-count")).toHaveText(/Showing \d+ - \d+ of \d+ results/, {
timeout: 15_000,
});
// Verify the model name appears in the table body
const tableBody = page.locator("table tbody");
@ -181,7 +183,9 @@ test.describe("Add Model", () => {
await page.waitForTimeout(1000);
// Verify the model appears in the results count (not "Showing 0 results")
await expect(page.getByText(/Showing \d+ - \d+ of \d+ results/)).toBeVisible({ timeout: 15_000 });
await expect(page.getByTestId("models-results-count")).toHaveText(/Showing \d+ - \d+ of \d+ results/, {
timeout: 15_000,
});
// Verify the wildcard model appears in the table body (wildcard models show as "cohere/*")
const tableBody = page.locator("table tbody");

41
ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/ModelsAndEndpointsView.tsx
View file

@ -46,6 +46,8 @@ interface GlobalRetryPolicyObject {
[retryPolicyKey: string]: number;
}
const HEALTH_PAGE_SIZE = 50;
const ModelsAndEndpointsView: React.FC<ModelDashboardProps> = ({ premiumUser, teams }) => {
const { accessToken, token, userRole, userId: userID } = useAuthorized();
const [addModelForm] = Form.useForm();
@ -62,6 +64,7 @@ const ModelsAndEndpointsView: React.FC<ModelDashboardProps> = ({ premiumUser, te
const [selectedModelId, setSelectedModelId] = useState<string | null>(null);
const [selectedTeamId, setSelectedTeamId] = useState<string | null>(null);
const [selectedTabIndex, setSelectedTabIndex] = useState(0);
const [healthCurrentPage, setHealthCurrentPage] = useState(1);
const [showMissingProviderBanner, setShowMissingProviderBanner] = useState(() => {
if (typeof window !== "undefined") {
return localStorage.getItem("hideMissingProviderBanner") !== "true";
@ -71,6 +74,10 @@ const ModelsAndEndpointsView: React.FC<ModelDashboardProps> = ({ premiumUser, te
const queryClient = useQueryClient();
const { data: modelDataResponse, isLoading: isLoadingModels, refetch: refetchModels } = useModelsInfo();
const { data: healthModelDataResponse, isLoading: isLoadingHealthModels } = useModelsInfo(
healthCurrentPage,
HEALTH_PAGE_SIZE,
);
const { data: modelCostMapData, isLoading: isLoadingModelCostMap } = useModelCostMap();
const { data: credentialsResponse, isLoading: isLoadingCredentials } = useCredentials();
const credentialsList = credentialsResponse?.credentials || [];
@ -104,12 +111,12 @@ const ModelsAndEndpointsView: React.FC<ModelDashboardProps> = ({ premiumUser, te
return modelDataResponse.data.map((model: any) => model.model_name);
}, [modelDataResponse?.data]);
const allModelIdsOnProxy = useMemo<string[]>(() => {
if (!modelDataResponse?.data) return [];
return modelDataResponse.data
const healthModelIdsOnProxy = useMemo<string[]>(() => {
if (!healthModelDataResponse?.data) return [];
return healthModelDataResponse.data
.map((model: any) => model.model_info?.id)
.filter((id: string | undefined): id is string => Boolean(id));
}, [modelDataResponse?.data]);
}, [healthModelDataResponse?.data]);
const getProviderFromModel = (model: string) => {
if (modelCostMapData !== null && modelCostMapData !== undefined) {
@ -125,6 +132,20 @@ const ModelsAndEndpointsView: React.FC<ModelDashboardProps> = ({ premiumUser, te
return transformModelData(modelDataResponse, getProviderFromModel);
}, [modelDataResponse?.data, getProviderFromModel]);
const processedHealthModelData = useMemo(() => {
if (!healthModelDataResponse?.data) return { data: [] };
return transformModelData(healthModelDataResponse, getProviderFromModel);
}, [healthModelDataResponse?.data, getProviderFromModel]);
const healthPaginationMeta = useMemo(() => {
return {
total_count: healthModelDataResponse?.total_count ?? 0,
current_page: healthModelDataResponse?.current_page ?? healthCurrentPage,
total_pages: healthModelDataResponse?.total_pages ?? 1,
size: healthModelDataResponse?.size ?? HEALTH_PAGE_SIZE,
};
}, [healthModelDataResponse, healthCurrentPage]);
const isProxyAdmin = userRole && isProxyAdminRole(userRole);
const isInternalUser = userRole && internalUserRoles.includes(userRole);
const isUserTeamAdmin = userID && isUserTeamAdminForAnyTeam(teams, userID);
@ -166,7 +187,8 @@ const ModelsAndEndpointsView: React.FC<ModelDashboardProps> = ({ premiumUser, te
const handleRefreshClick = () => {
const currentDate = new Date();
setLastRefreshed(currentDate.toLocaleTimeString([], { hour: '2-digit', minute: '2-digit' }));
setLastRefreshed(currentDate.toLocaleTimeString([], { hour: "2-digit", minute: "2-digit" }));
setHealthCurrentPage(1);
queryClient.invalidateQueries({ queryKey: ["models", "list"] });
refetchModels();
};
@ -441,11 +463,16 @@ const ModelsAndEndpointsView: React.FC<ModelDashboardProps> = ({ premiumUser, te
<TabPanel>
<HealthCheckComponent
accessToken={accessToken}
modelData={processedModelData}
all_models_on_proxy={allModelIdsOnProxy}
modelData={processedHealthModelData}
all_models_on_proxy={healthModelIdsOnProxy}
getDisplayModelName={getDisplayModelName}
setSelectedModelId={setSelectedModelId}
teams={teams}
isLoading={isLoadingHealthModels}
paginationMeta={healthPaginationMeta}
currentPage={healthCurrentPage}
pageSize={HEALTH_PAGE_SIZE}
onPageChange={setHealthCurrentPage}
/>
</TabPanel>
<ModelRetrySettingsTab

58
ui/litellm-dashboard/src/components/model_dashboard/HealthCheckComponent.test.tsx
View file

@ -92,18 +92,62 @@ describe("HealthCheckComponent", () => {
expect(mockIndividualModelHealthCheckCall).not.toHaveBeenCalledWith("token-123", "gpt-4");
});
it("should show pagination controls and request the next page", async () => {
const onPageChange = vi.fn();
const modelData = {
data: [
{
model_name: "gpt-4",
model_info: { id: "deployment-1" },
litellm_model_name: "gpt-4",
},
],
};
await act(async () => {
render(
<HealthCheckComponent
accessToken="token"
modelData={modelData}
all_models_on_proxy={["deployment-1"]}
getDisplayModelName={getDisplayModelName}
paginationMeta={{
total_count: 75,
current_page: 1,
total_pages: 2,
size: 50,
}}
currentPage={1}
pageSize={50}
onPageChange={onPageChange}
/>,
);
});
await act(async () => {
await new Promise((r) => setTimeout(r, 0));
});
expect(screen.getByTestId("health-results-count")).toHaveTextContent("Showing 1 - 50 of 75 results");
await act(async () => {
screen.getByRole("button", { name: "Next" }).click();
});
expect(onPageChange).toHaveBeenCalledWith(2);
});
describe("latest_health_checks keyed by model id", () => {
it("should show status from latest_health_checks when keys match model ids", async () => {
const modelData = {
data: [
{
model_name: "gpt-4",
model_info: { id: "id-alpha" },
litellm_model_name: "gpt-4",
{
model_name: "gpt-4",
model_info: { id: "id-alpha" },
litellm_model_name: "gpt-4",
},
{
model_name: "gpt-4",
model_info: { id: "id-beta" },
{
model_name: "gpt-4",
model_info: { id: "id-beta" },
litellm_model_name: "gpt-4",
},
],

Some files were not shown because too many files have changed in this diff Show more