scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-07 02:59:05 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 103

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

Browse source 
# Conflicts:
#	litellm/proxy/auth/auth_checks.py
This commit is contained in:
user 2026-05-01 11:37:43 -07:00
commit 2ed4b5968d
94 changed files with 6068 additions and 941 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")

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

10
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

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

2
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

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"

300
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
@ -846,7 +847,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]:
"""
@ -869,9 +870,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:
@ -885,14 +889,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)}")
@ -903,7 +909,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"].
@ -960,7 +966,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:
"""
@ -1042,7 +1048,7 @@ async 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,
@ -1073,10 +1079,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,
@ -1111,9 +1119,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
@ -1131,7 +1141,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]:
@ -1164,12 +1174,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)
@ -1185,11 +1195,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}")
@ -1200,7 +1212,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]:
@ -1239,7 +1251,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"]:
@ -1259,9 +1271,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:
@ -1273,10 +1288,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:
@ -1444,7 +1461,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,
@ -1463,12 +1480,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")
@ -1530,7 +1547,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,
)
@ -1551,13 +1569,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,
)
@ -1565,7 +1591,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)
@ -1578,13 +1604,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
@ -1600,12 +1627,13 @@ async def _cache_key_object(
value=cached_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
@ -1653,7 +1681,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],
@ -1714,38 +1742,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,
@ -1811,20 +1839,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)
@ -1842,7 +1871,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:
"""
@ -1864,13 +1893,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:
@ -1916,7 +1944,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:
@ -1998,6 +2026,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
@ -2005,6 +2034,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,
)
@ -2026,7 +2056,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]:
@ -2053,12 +2083,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:
@ -2088,13 +2118,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,
)
@ -2297,7 +2329,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,
@ -2315,17 +2347,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 _copy_user_api_key_auth_for_cache(
user_api_key_obj=UserAPIKeyAuth(**cached_key_obj)
)
elif isinstance(cached_key_obj, UserAPIKeyAuth):
return _copy_user_api_key_auth_for_cache(user_api_key_obj=cached_key_obj)
if user_api_key_auth is not None:
return _copy_user_api_key_auth_for_cache(user_api_key_obj=user_api_key_auth)
if check_cache_only:
raise Exception(
@ -2392,7 +2421,7 @@ def _copy_user_api_key_auth_for_cache(
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]:
@ -2408,12 +2437,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:
@ -2424,14 +2453,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
@ -2440,7 +2470,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]:
@ -2460,14 +2490,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)
@ -2493,7 +2521,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)
@ -2505,7 +2534,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,
@ -2536,12 +2565,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}}
@ -2555,16 +2584,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."
@ -2577,7 +2606,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]:
"""
@ -2635,7 +2664,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]:
"""
@ -2654,7 +2683,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]:
"""
@ -2673,7 +2702,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]:
"""
@ -3397,7 +3426,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."""
@ -3465,7 +3494,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:
"""
@ -3772,7 +3801,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]:
"""
@ -3787,12 +3816,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(
@ -3811,6 +3840,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
@ -3820,7 +3850,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,
):
"""
@ -3914,7 +3944,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,

60
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
@ -1902,7 +1914,7 @@ async def _reserve_budget_after_common_checks(
team_object: Optional[LiteLLM_TeamTableCachedObj],
user_object: Optional[LiteLLM_UserTable],
prisma_client: Optional[PrismaClient],
user_api_key_cache: DualCache,
user_api_key_cache: UserApiKeyCache,
proxy_logging_obj: ProxyLogging,
skip_budget_checks: bool,
end_user_id: Optional[str] = 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

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)

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

27
litellm/proxy/middleware/prometheus_auth_middleware.py
View file

@ -3,6 +3,7 @@ Prometheus Auth Middleware - Pure ASGI implementation
"""
import json
from typing import Any, List, MutableMapping
from fastapi import Request
from starlette.types import ASGIApp, Receive, Scope, Send
@ -40,8 +41,17 @@ class PrometheusAuthMiddleware:
# Only run auth if configured to do so
if litellm.require_auth_for_metrics_endpoint is True:
# Construct Request only when auth is actually needed
request = Request(scope, receive)
# user_api_key_auth reads the request body, which consumes ASGI `receive`.
# Buffer those messages and replay them for the inner app; otherwise a
# successful auth would forward an exhausted receive and /metrics hangs.
buffered_messages: List[MutableMapping[str, Any]] = []
async def receive_for_auth() -> MutableMapping[str, Any]:
message = await receive()
buffered_messages.append(message)
return message
request = Request(scope, receive_for_auth)
api_key = request.headers.get(_AUTHORIZATION_HEADER) or ""
try:
@ -70,5 +80,18 @@ class PrometheusAuthMiddleware:
)
return
replay_idx = 0
async def receive_replay() -> MutableMapping[str, Any]:
nonlocal replay_idx
if replay_idx < len(buffered_messages):
msg = buffered_messages[replay_idx]
replay_idx += 1
return msg
return await receive()
await self.app(scope, receive_replay, send)
return
# Pass through to the inner application
await self.app(scope, receive, send)

375
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(
@ -2290,14 +2295,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
@ -2355,41 +2362,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)
@ -2421,31 +2435,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. "
@ -2464,36 +2480,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. "
@ -2520,32 +2532,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. "
@ -3213,8 +3225,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:
@ -3230,7 +3243,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.
@ -3556,7 +3591,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}"
@ -3827,21 +3868,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
@ -6570,7 +6613,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:
@ -6619,7 +6662,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."""
@ -7259,7 +7302,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
@ -10561,6 +10604,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,
@ -10583,78 +10721,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

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

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

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

@ -927,19 +927,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()
@ -985,13 +987,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]}"

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

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

31
tests/test_litellm/proxy/middleware/test_prometheus_auth_middleware.py
View file

@ -26,6 +26,15 @@ async def fake_valid_auth(request, api_key):
return
async def fake_valid_auth_reads_body(request, api_key, **kwargs):
"""
Like real user_api_key_auth, consumes the ASGI body stream. Regression test
for successful auth passing a drained receive to the inner app (hang).
"""
await request.body()
return
async def fake_invalid_auth(request, api_key):
print("running fake invalid auth", request, api_key)
# Simulate invalid auth by raising an exception.
@ -62,6 +71,28 @@ def app_with_middleware():
return app
def test_valid_auth_metrics_after_body_consumed(app_with_middleware, monkeypatch):
"""
Auth that reads the request body must not cause /metrics to hang on success.
"""
litellm.require_auth_for_metrics_endpoint = True
monkeypatch.setattr(
"litellm.proxy.middleware.prometheus_auth_middleware.user_api_key_auth",
fake_valid_auth_reads_body,
)
client = TestClient(app_with_middleware)
headers = {SpecialHeaders.openai_authorization.value: "valid"}
response = client.get("/metrics", headers=headers)
assert response.status_code == 200, response.text
assert response.json() == {"msg": "metrics OK"}
response = client.get("/metrics/", headers=headers)
assert response.status_code == 200, response.text
assert response.json() == {"msg": "metrics OK"}
def test_valid_auth_metrics(app_with_middleware, monkeypatch):
"""
Test that a request to /metrics (and /metrics/) with valid auth headers passes.

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"

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

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

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

@ -26,6 +26,16 @@ interface HealthCheckComponentProps {
getDisplayModelName: (model: any) => string;
setSelectedModelId?: (modelId: string) => void;
teams?: Team[] | null;
isLoading?: boolean;
paginationMeta?: {
total_count: number;
current_page: number;
total_pages: number;
size: number;
};
currentPage?: number;
pageSize?: number;
onPageChange?: (page: number) => void;
}
const HealthCheckComponent: React.FC<HealthCheckComponentProps> = ({
@ -35,6 +45,11 @@ const HealthCheckComponent: React.FC<HealthCheckComponentProps> = ({
getDisplayModelName,
setSelectedModelId,
teams,
isLoading = false,
paginationMeta,
currentPage = 1,
pageSize = 50,
onPageChange,
}) => {
const [modelHealthStatuses, setModelHealthStatuses] = useState<{ [key: string]: HealthStatus }>({});
const [selectedModelsForHealth, setSelectedModelsForHealth] = useState<string[]>([]);
@ -95,19 +110,19 @@ const HealthCheckComponent: React.FC<HealthCheckComponentProps> = ({
const fullError = checkData.error_message || undefined;
healthStatusMap[modelId] = {
status: checkData.status || "unknown",
lastCheck: checkData.checked_at ? new Date(checkData.checked_at).toLocaleString() : "None",
lastSuccess:
checkData.status === "healthy"
? checkData.checked_at
? new Date(checkData.checked_at).toLocaleString()
: "None"
: "None",
loading: false,
error: fullError ? extractMeaningfulError(fullError) : undefined,
fullError: fullError,
successResponse: checkData.status === "healthy" ? checkData : undefined,
};
status: checkData.status || "unknown",
lastCheck: checkData.checked_at ? new Date(checkData.checked_at).toLocaleString() : "None",
lastSuccess:
checkData.status === "healthy"
? checkData.checked_at
? new Date(checkData.checked_at).toLocaleString()
: "None"
: "None",
loading: false,
error: fullError ? extractMeaningfulError(fullError) : undefined,
fullError: fullError,
successResponse: checkData.status === "healthy" ? checkData : undefined,
};
});
}
} catch (healthError) {
@ -448,6 +463,13 @@ const HealthCheckComponent: React.FC<HealthCheckComponentProps> = ({
}
};
const handlePageChange = (page: number) => {
setSelectedModelsForHealth([]);
setAllModelsSelected(false);
setModelHealthStatuses({});
onPageChange?.(page);
};
const getStatusBadge = (status: string) => {
switch (status) {
case "healthy":
@ -490,6 +512,36 @@ const HealthCheckComponent: React.FC<HealthCheckComponentProps> = ({
setSelectedSuccessDetails(null);
};
const healthTableData = (modelData?.data ?? []).map((model: any) => {
const modelId = model.model_info?.id;
const healthStatus = modelId ? modelHealthStatuses[modelId] : null;
const status = healthStatus || {
status: "none",
lastCheck: "None",
loading: false,
};
return {
model_name: model.model_name,
model_info: model.model_info,
provider: model.provider,
litellm_model_name: model.litellm_model_name,
health_status: status.status,
last_check: status.lastCheck,
last_success: status.lastSuccess || "None",
health_loading: status.loading,
health_error: status.error,
health_full_error: status.fullError,
};
});
const shouldShowPagination = Boolean(paginationMeta && onPageChange);
const totalCount = paginationMeta?.total_count ?? 0;
const totalPages = paginationMeta?.total_pages ?? 1;
const pageForDisplay = paginationMeta?.current_page ?? currentPage;
const pageSizeForDisplay = paginationMeta?.size ?? pageSize;
const resultsStart = shouldShowPagination && totalCount > 0 ? (pageForDisplay - 1) * pageSizeForDisplay + 1 : 0;
const resultsEnd = shouldShowPagination ? Math.min(pageForDisplay * pageSizeForDisplay, totalCount) : 0;
return (
<div>
<div className="mb-6">
@ -522,6 +574,38 @@ const HealthCheckComponent: React.FC<HealthCheckComponentProps> = ({
</div>
<div>
{shouldShowPagination && (
<div className="flex justify-between items-center mb-3">
<span data-testid="health-results-count" className="text-sm text-gray-700">
{totalCount > 0
? `Showing ${resultsStart} - ${resultsEnd} of ${totalCount} results`
: "Showing 0 results"}
</span>
<div className="flex items-center space-x-2">
<button
onClick={() => handlePageChange(currentPage - 1)}
disabled={isLoading || currentPage === 1}
className={`px-3 py-1 text-sm border rounded-md ${
isLoading || currentPage === 1 ? "bg-gray-100 text-gray-400 cursor-not-allowed" : "hover:bg-gray-50"
}`}
>
Previous
</button>
<button
onClick={() => handlePageChange(currentPage + 1)}
disabled={isLoading || currentPage >= totalPages}
className={`px-3 py-1 text-sm border rounded-md ${
isLoading || currentPage >= totalPages
? "bg-gray-100 text-gray-400 cursor-not-allowed"
: "hover:bg-gray-50"
}`}
>
Next
</button>
</div>
</div>
)}
<ModelDataTable
columns={healthCheckColumns(
modelHealthStatuses,
@ -537,28 +621,8 @@ const HealthCheckComponent: React.FC<HealthCheckComponentProps> = ({
setSelectedModelId,
teams,
)}
data={modelData.data.map((model: any) => {
const modelId = model.model_info?.id;
const healthStatus = modelId ? modelHealthStatuses[modelId] : null;
const status = healthStatus || {
status: "none",
lastCheck: "None",
loading: false,
};
return {
model_name: model.model_name,
model_info: model.model_info,
provider: model.provider,
litellm_model_name: model.litellm_model_name,
health_status: status.status,
last_check: status.lastCheck,
last_success: status.lastSuccess || "None",
health_loading: status.loading,
health_error: status.error,
health_full_error: status.fullError,
};
})}
isLoading={false}
data={healthTableData}
isLoading={isLoading}
/>
</div>

2
uv.lock generated
View file

@ -3418,7 +3418,7 @@ source = { editable = "enterprise" }
[[package]]
name = "litellm-proxy-extras"
version = "0.4.69"
version = "0.4.70"
source = { editable = "litellm-proxy-extras" }
[[package]]