Merge remote-tracking branch 'upstream/litellm_internal_staging' into add-gemma-4-26b-vertex-ai

# Conflicts:
#	litellm/llms/vertex_ai/common_utils.py
This commit is contained in:
Emerson Gomes 2026-05-01 14:18:38 -05:00
commit c91e2ded8e
No known key found for this signature in database
GPG key ID: D3DF28AB5D1B5E17
86 changed files with 6708 additions and 1022 deletions

View file

@ -1,75 +0,0 @@
name: Check Lazy OpenAPI Snapshot
on:
pull_request:
branches:
- main
- litellm_internal_staging
- "litellm_**"
permissions:
contents: read
checks: write
concurrency:
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
cancel-in-progress: true
jobs:
verify:
runs-on: ubuntu-latest
timeout-minutes: 10
steps:
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
persist-credentials: false
- name: Set up Python
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
with:
python-version: "3.12"
- name: Set up uv
uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7
with:
version: "0.10.9"
- name: Cache uv dependencies
uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0
with:
path: |
~/.cache/uv
.venv
key: ${{ runner.os }}-uv-${{ hashFiles('uv.lock') }}
restore-keys: |
${{ runner.os }}-uv-
- name: Install dependencies
run: uv sync --frozen --all-groups --all-extras
- name: Regenerate snapshot to /tmp
id: regen
run: |
cp litellm/proxy/_lazy_openapi_snapshot.json /tmp/snapshot.committed.json
uv run --no-sync python -m litellm.proxy._lazy_openapi_snapshot
mv litellm/proxy/_lazy_openapi_snapshot.json /tmp/snapshot.fresh.json
mv /tmp/snapshot.committed.json litellm/proxy/_lazy_openapi_snapshot.json
- name: Compare
id: diff
continue-on-error: true
run: |
diff -q /tmp/snapshot.fresh.json litellm/proxy/_lazy_openapi_snapshot.json
- name: Mark neutral if drift
if: steps.diff.outcome == 'failure'
uses: LouisBrunner/checks-action@6b626ffbad7cc56fd58627f774b9067e6118af23 # v2.0.0
with:
token: ${{ secrets.GITHUB_TOKEN }}
name: lazy-openapi-snapshot
conclusion: neutral
output: |
{
"title": "Lazy openapi snapshot is stale",
"summary": "Run `python -m litellm.proxy._lazy_openapi_snapshot` and commit the regenerated `litellm/proxy/_lazy_openapi_snapshot.json`. Not blocking — the snapshot will regenerate at release if not committed."
}

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

View file

@ -68,7 +68,7 @@ Managing LLM calls across providers gets complicated fast — different SDKs, au
<td><img height="60" alt="Stripe" src="https://github.com/user-attachments/assets/f7296d4f-9fbd-460d-9d05-e4df31697c4b" /></td>
<td><img height="60" alt="image" src="https://github.com/user-attachments/assets/436fca71-988b-40bb-b5fe-8450c80fdbd0" /></td>
<td><img height="60" alt="Google ADK" src="https://github.com/user-attachments/assets/caf270a2-5aee-45c4-8222-41a2070c4f19" /></td>
<td><img height="60" alt="Greptile" src="https://github.com/user-attachments/assets/0be4bd8a-7cfa-48d3-9090-f415fe948280" /></td>
<td><img height="60" alt="Greptile" src="https://github.com/user-attachments/assets/3db0ae72-0843-4005-a56d-bba1dde2193d" /></td>
<td><img height="60" alt="OpenHands" src="https://github.com/user-attachments/assets/a6150c4c-149e-4cae-888b-8b92be6e003f" /></td>
<td><h2>Netflix</h2></td>
<td><img height="60" alt="OpenAI Agents SDK" src="https://github.com/user-attachments/assets/c02f7be0-8c2e-4d27-aea7-7c024bfaebc0" /></td>

View file

@ -857,10 +857,16 @@ async def project_info(
where={"team_id": project.team_id}
)
if team:
is_team_member = (
user_api_key_dict.user_id in team.admins
or user_api_key_dict.user_id in team.members
)
caller_user_id = user_api_key_dict.user_id
for m in team.members_with_roles or []:
m_user_id = (
m.get("user_id")
if isinstance(m, dict)
else getattr(m, "user_id", None)
)
if m_user_id == caller_user_id:
is_team_member = True
break
if not (is_admin or is_team_member):
raise HTTPException(
@ -911,20 +917,20 @@ async def list_projects(
include={"litellm_budget_table": True, "object_permission": True}
)
else:
# Get projects for teams the user belongs to
user_teams = await prisma_client.db.litellm_teamtable.find_many(
where={
"OR": [
{"members": {"has": user_api_key_dict.user_id}},
{"admins": {"has": user_api_key_dict.user_id}},
]
}
# Look up the user's team memberships via the reverse-index on
# LiteLLM_UserTable.teams (maintained by team_member_add alongside
# members_with_roles). This avoids a full scan of all team rows.
user_record = await prisma_client.db.litellm_usertable.find_unique(
where={"user_id": user_api_key_dict.user_id},
)
user_team_ids = (
user_record.teams
if user_record is not None and user_record.teams
else []
)
team_ids = [team.team_id for team in user_teams]
projects = await prisma_client.db.litellm_projecttable.find_many(
where={"team_id": {"in": team_ids}},
where={"team_id": {"in": user_team_ids}},
include={"litellm_budget_table": True, "object_permission": True},
)

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:

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

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:

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

View file

@ -2,11 +2,23 @@
Arize Phoenix API client for fetching prompt versions from Arize Phoenix.
"""
import urllib.parse
from typing import Any, Dict, Optional
from litellm.llms.custom_httpx.http_handler import HTTPHandler
def _sanitize_id(identifier: str) -> str:
"""Reject path traversal characters and URL-encode the identifier."""
if any(c in identifier for c in ("/", "\\", "#", "?")):
raise ValueError(
f"Invalid identifier {identifier!r}: contains disallowed characters"
)
if ".." in identifier:
raise ValueError(f"Invalid identifier {identifier!r}: path traversal detected")
return urllib.parse.quote(identifier, safe="")
class ArizePhoenixClient:
"""
Client for interacting with Arize Phoenix API to fetch prompt versions.
@ -53,7 +65,8 @@ class ArizePhoenixClient:
Returns:
Dictionary containing prompt version data, or None if not found
"""
url = f"{self.api_base}/v1/prompt_versions/{prompt_version_id}"
safe_id = _sanitize_id(prompt_version_id)
url = f"{self.api_base}/v1/prompt_versions/{safe_id}"
try:
# Use the underlying httpx client directly to avoid query param extraction

View file

@ -3,11 +3,27 @@ BitBucket API client for fetching .prompt files from BitBucket repositories.
"""
import base64
import urllib.parse
from typing import Any, Dict, List, Optional
from litellm.llms.custom_httpx.http_handler import HTTPHandler
def _sanitize_file_path(file_path: str) -> str:
"""Reject path traversal and URL-encode each path segment."""
if "#" in file_path or "?" in file_path:
raise ValueError(
f"Invalid file path {file_path!r}: contains URL special characters"
)
parts = file_path.split("/")
for part in parts:
if part == "..":
raise ValueError(
f"Invalid file path {file_path!r}: path traversal detected"
)
return "/".join(urllib.parse.quote(part, safe="") for part in parts)
class BitBucketClient:
"""
Client for interacting with BitBucket API to fetch .prompt files.
@ -72,7 +88,8 @@ class BitBucketClient:
Returns:
File content as string, or None if file not found
"""
url = f"{self.base_url}/repositories/{self.workspace}/{self.repository}/src/{self.branch}/{file_path}"
safe_path = _sanitize_file_path(file_path)
url = f"{self.base_url}/repositories/{self.workspace}/{self.repository}/src/{self.branch}/{safe_path}"
try:
response = self.http_handler.get(url, headers=self.headers)
@ -119,7 +136,8 @@ class BitBucketClient:
Returns:
List of file paths
"""
url = f"{self.base_url}/repositories/{self.workspace}/{self.repository}/src/{self.branch}/{directory_path}"
safe_dir = _sanitize_file_path(directory_path) if directory_path else ""
url = f"{self.base_url}/repositories/{self.workspace}/{self.repository}/src/{self.branch}/{safe_dir}"
try:
response = self.http_handler.get(url, headers=self.headers)
@ -211,7 +229,8 @@ class BitBucketClient:
Returns:
Dictionary containing file metadata, or None if file not found
"""
url = f"{self.base_url}/repositories/{self.workspace}/{self.repository}/src/{self.branch}/{file_path}"
safe_path = _sanitize_file_path(file_path)
url = f"{self.base_url}/repositories/{self.workspace}/{self.repository}/src/{self.branch}/{safe_path}"
try:
# Use GET with Range header to get just the headers (HEAD equivalent)

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

View file

@ -31,15 +31,23 @@ def load_cli_token() -> Optional[dict]:
return None
def get_litellm_gateway_api_key() -> Optional[str]:
def get_litellm_gateway_api_key(
expected_base_url: Optional[str] = None,
) -> Optional[str]:
"""
Get the stored CLI API key for use with LiteLLM SDK.
This function reads the token file created by `litellm-proxy login`
and returns the API key for use in Python scripts.
Args:
expected_base_url: When provided, the key is only returned if it was
originally issued for this URL. Pass the target server URL to
prevent credential leakage when the client is pointed at a
different (possibly malicious) server.
Returns:
str: The API key if found, None otherwise
str: The API key if found (and origin matches), None otherwise
Example:
>>> import litellm
@ -53,6 +61,10 @@ def get_litellm_gateway_api_key() -> Optional[str]:
>>> )
"""
token_data = load_cli_token()
if token_data and "key" in token_data:
return token_data["key"]
return None
if not token_data or "key" not in token_data:
return None
if expected_base_url is not None:
stored_url = token_data.get("base_url")
if stored_url != expected_base_url.rstrip("/"):
return None
return token_data["key"]

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(

View file

@ -97,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", "google/gemma-4-26b-a4b-it-maas", "openai/gpt-oss-120b")
model: The model name (e.g., "llama3-405b", "gemini-pro", "gemma/gemma-3-12b-it", "google/gemma-4-26b-a4b-it-maas", "openai/gpt-oss-120b", "xai/grok-4.1-fast-non-reasoning")
litellm_params: Optional litellm parameters dict that may contain base_model for routing
Returns:
@ -119,6 +119,9 @@ def get_vertex_ai_model_route(
>>> get_vertex_ai_model_route("openai/gpt-oss-120b")
VertexAIModelRoute.MODEL_GARDEN
>>> 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"})
VertexAIModelRoute.GEMINI # Numeric endpoints with api_base use HTTP path
"""
@ -152,8 +155,11 @@ def get_vertex_ai_model_route(
if "gemma/" in model or model.startswith("google/gemma-"):
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
@ -259,8 +265,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"

View file

@ -3,7 +3,7 @@ Google AI Studio /batchEmbedContents Embeddings Endpoint
"""
import json
from typing import Any, Dict, Literal, Optional, Union
from typing import Any, Dict, List, Literal, Optional, Tuple, Union
import httpx
@ -13,8 +13,8 @@ from litellm.llms.custom_httpx.http_handler import (
HTTPHandler,
get_async_httpx_client,
)
from litellm.types.llms.openai import EmbeddingInput
from litellm.types.llms.vertex_ai import (
GeminiEmbeddingInput,
VertexAIBatchEmbeddingsRequestBody,
VertexAIBatchEmbeddingsResponseObject,
)
@ -23,7 +23,6 @@ from litellm.types.utils import EmbeddingResponse
from ..gemini.vertex_and_google_ai_studio_gemini import VertexLLM
from .batch_embed_content_transformation import (
_is_file_reference,
_is_multimodal_input,
process_embed_content_response,
process_response,
transform_openai_input_gemini_content,
@ -32,9 +31,24 @@ from .batch_embed_content_transformation import (
class GoogleBatchEmbeddings(VertexLLM):
@staticmethod
def _flatten_and_detect_file_refs(
input: GeminiEmbeddingInput,
) -> Tuple[List[str], bool]:
"""Flatten nested input lists and detect file references."""
input_list = [input] if isinstance(input, str) else input
flat_elements = [
e
for item in input_list
for e in (item if isinstance(item, list) else [item])
if isinstance(e, str)
]
has_file_refs = any(_is_file_reference(e) for e in flat_elements)
return flat_elements, has_file_refs
def _resolve_file_references(
self,
input: EmbeddingInput,
input: GeminiEmbeddingInput,
api_key: str,
sync_handler: HTTPHandler,
) -> Dict[str, Dict[str, str]]:
@ -42,7 +56,7 @@ class GoogleBatchEmbeddings(VertexLLM):
Resolve Gemini file references (files/...) to get mime_type and uri.
Args:
input: EmbeddingInput that may contain file references
input: GeminiEmbeddingInput that may contain file references
api_key: Gemini API key
sync_handler: HTTP client
@ -73,7 +87,7 @@ class GoogleBatchEmbeddings(VertexLLM):
async def _async_resolve_file_references(
self,
input: EmbeddingInput,
input: GeminiEmbeddingInput,
api_key: str,
async_handler: AsyncHTTPHandler,
) -> Dict[str, Dict[str, str]]:
@ -81,7 +95,7 @@ class GoogleBatchEmbeddings(VertexLLM):
Async version of _resolve_file_references.
Args:
input: EmbeddingInput that may contain file references
input: GeminiEmbeddingInput that may contain file references
api_key: Gemini API key
async_handler: Async HTTP client
@ -110,10 +124,10 @@ class GoogleBatchEmbeddings(VertexLLM):
return resolved_files
def batch_embeddings(
def batch_embeddings( # noqa: PLR0915
self,
model: str,
input: EmbeddingInput,
input: GeminiEmbeddingInput,
print_verbose,
model_response: EmbeddingResponse,
custom_llm_provider: Literal["gemini", "vertex_ai"],
@ -151,8 +165,7 @@ class GoogleBatchEmbeddings(VertexLLM):
optional_params = optional_params or {}
is_multimodal = _is_multimodal_input(input)
use_embed_content = is_multimodal or (custom_llm_provider == "vertex_ai")
use_embed_content = custom_llm_provider == "vertex_ai"
mode: Literal["embedding", "batch_embedding"]
if use_embed_content:
mode = "embedding"
@ -215,8 +228,22 @@ class GoogleBatchEmbeddings(VertexLLM):
resolved_files=resolved_files,
)
else:
flat_elements, has_file_refs = self._flatten_and_detect_file_refs(input)
if has_file_refs and not api_key:
raise ValueError(
"An API key is required to resolve Gemini file references (files/...). "
"Pass api_key= or set GEMINI_API_KEY."
)
resolved_files = {}
if api_key and has_file_refs:
resolved_files = self._resolve_file_references(
input=flat_elements, api_key=api_key, sync_handler=sync_handler
)
request_data = transform_openai_input_gemini_content(
input=input, model=model, optional_params=optional_params
input=input,
model=model,
optional_params=optional_params,
resolved_files=resolved_files,
)
## LOGGING
@ -264,7 +291,7 @@ class GoogleBatchEmbeddings(VertexLLM):
url: str,
data: Optional[Union[VertexAIBatchEmbeddingsRequestBody, dict]],
model_response: EmbeddingResponse,
input: EmbeddingInput,
input: GeminiEmbeddingInput,
timeout: Optional[Union[float, httpx.Timeout]],
headers={},
client: Optional[AsyncHTTPHandler] = None,
@ -303,8 +330,22 @@ class GoogleBatchEmbeddings(VertexLLM):
resolved_files=resolved_files,
)
else:
flat_elements, has_file_refs = self._flatten_and_detect_file_refs(input)
if has_file_refs and not api_key:
raise ValueError(
"An API key is required to resolve Gemini file references (files/...). "
"Pass api_key= or set GEMINI_API_KEY."
)
resolved_files = {}
if api_key and has_file_refs:
resolved_files = await self._async_resolve_file_references(
input=flat_elements, api_key=api_key, async_handler=async_handler
)
data = transform_openai_input_gemini_content(
input=input, model=model, optional_params=optional_params or {}
input=input,
model=model,
optional_params=optional_params or {},
resolved_files=resolved_files,
)
## LOGGING

View file

@ -6,12 +6,12 @@ Why separate file? Make it easy to see how transformation works
from typing import Dict, List, Optional, Tuple
from litellm.types.llms.openai import EmbeddingInput
from litellm.types.llms.vertex_ai import (
BlobType,
ContentType,
EmbedContentRequest,
FileDataType,
GeminiEmbeddingInput,
PartType,
VertexAIBatchEmbeddingsRequestBody,
VertexAIBatchEmbeddingsResponseObject,
@ -114,33 +114,77 @@ def _parse_data_url(data_url: str) -> Tuple[str, str]:
return media_type, base64_data
def _is_multimodal_input(input: EmbeddingInput) -> bool:
def _is_multimodal_input(input: GeminiEmbeddingInput) -> bool:
"""
Check if the input contains multimodal data (data URIs, file references, or GCS URLs).
Check if the input contains multimodal data (data URIs, file references,
GCS URLs, or nested lists for combined embeddings).
Args:
input: EmbeddingInput (str or List[str])
input: GeminiEmbeddingInput str, List[str], or List[List[str]] for combined embeddings
Returns:
bool: True if any element is a data URI, file reference, or GCS URL
bool: True if any element is multimodal or a nested list
"""
if isinstance(input, str):
input_list = [input]
else:
input_list = input
return _is_multimodal_element(input)
for element in input_list:
if isinstance(element, str):
if element.startswith("data:") and ";base64," in element:
return True
if _is_file_reference(element):
return True
if _is_gcs_url(element):
for element in input:
if isinstance(element, list):
if any(
_is_multimodal_element(sub) for sub in element if isinstance(sub, str)
):
return True
elif isinstance(element, str) and _is_multimodal_element(element):
return True
return False
def _is_multimodal_element(element: str) -> bool:
"""Check if a single string element is multimodal."""
if element.startswith("data:") and ";base64," in element:
return True
if _is_file_reference(element):
return True
if _is_gcs_url(element):
return True
return False
def _build_part_for_input(
element: str,
resolved_files: Optional[Dict[str, Dict[str, str]]] = None,
) -> PartType:
"""
Build a single PartType for an input element, handling text, data URIs,
file references, and GCS URLs.
"""
resolved_files = resolved_files or {}
if element.startswith("data:") and ";base64," in element:
mime_type, base64_data = _parse_data_url(element)
blob: BlobType = {"mime_type": mime_type, "data": base64_data}
return PartType(inline_data=blob)
elif _is_gcs_url(element):
mime_type = _infer_mime_type_from_gcs_url(element)
file_data: FileDataType = {
"mime_type": mime_type,
"file_uri": element,
}
return PartType(file_data=file_data)
elif _is_file_reference(element):
if element not in resolved_files:
raise ValueError(f"File reference {element} not resolved")
file_info = resolved_files[element]
file_data_ref: FileDataType = {
"mime_type": file_info["mime_type"],
"file_uri": file_info["uri"],
}
return PartType(file_data=file_data_ref)
else:
return PartType(text=element)
_SUPPORTED_EMBED_PARAMS = {"outputDimensionality", "taskType", "title"}
@ -155,37 +199,60 @@ def _filter_embed_params(optional_params: dict) -> dict:
def transform_openai_input_gemini_content(
input: EmbeddingInput, model: str, optional_params: dict
input: GeminiEmbeddingInput,
model: str,
optional_params: dict,
resolved_files: Optional[Dict[str, Dict[str, str]]] = None,
) -> VertexAIBatchEmbeddingsRequestBody:
"""
The content to embed. Only the parts.text fields will be counted.
Transform OpenAI embedding input to Gemini batchEmbedContents format.
Each input element becomes a separate EmbedContentRequest, supporting
text, data URIs, file references, and GCS URLs.
If an element is a list (nested input), all sub-elements are combined
into a single content with multiple parts, producing one combined
embedding for the group.
Examples:
input=["text", "image"] 2 separate embeddings
input=[["text", "image"]] 1 combined embedding
input=[["text", "image"], "x"] 2 embeddings (1 combined + 1 separate)
"""
gemini_model_name = "models/{}".format(model)
gemini_params = _filter_embed_params(optional_params)
input_list = [input] if isinstance(input, str) else input
requests: List[EmbedContentRequest] = []
if isinstance(input, str):
for element in input_list:
if isinstance(element, list):
if not element:
raise ValueError("Nested input list must not be empty")
for sub in element:
if not isinstance(sub, str):
raise ValueError(
f"Elements inside a nested input list must be strings, got {type(sub)}"
)
parts = [
_build_part_for_input(sub, resolved_files=resolved_files)
for sub in element
]
else:
parts = [_build_part_for_input(element, resolved_files=resolved_files)]
request = EmbedContentRequest(
model=gemini_model_name,
content=ContentType(parts=[PartType(text=input)]),
content=ContentType(parts=parts),
**gemini_params,
)
requests.append(request)
else:
for i in input:
request = EmbedContentRequest(
model=gemini_model_name,
content=ContentType(parts=[PartType(text=i)]),
**gemini_params,
)
requests.append(request)
return VertexAIBatchEmbeddingsRequestBody(requests=requests)
def transform_openai_input_gemini_embed_content(
input: EmbeddingInput,
input: GeminiEmbeddingInput,
model: str,
optional_params: dict,
resolved_files: Optional[Dict[str, Dict[str, str]]] = None,
@ -194,7 +261,7 @@ def transform_openai_input_gemini_embed_content(
Transform OpenAI embedding input to Gemini embedContent format (multimodal).
Args:
input: EmbeddingInput (str or List[str]) with text, data URIs, or file references
input: GeminiEmbeddingInput with text, data URIs, or file references
model: Model name
optional_params: Additional parameters (taskType, outputDimensionality, etc.)
resolved_files: Dict mapping file names (files/abc) to {mime_type, uri}
@ -210,31 +277,14 @@ def transform_openai_input_gemini_embed_content(
parts: List[PartType] = []
for element in input_list:
if isinstance(element, list):
raise ValueError(
"Nested (combined) embeddings are not supported on the embedContent path. "
"Use the batchEmbedContents path or pass a flat list instead."
)
if not isinstance(element, str):
raise ValueError(f"Unsupported input type: {type(element)}")
if element.startswith("data:") and ";base64," in element:
mime_type, base64_data = _parse_data_url(element)
blob: BlobType = {"mime_type": mime_type, "data": base64_data}
parts.append(PartType(inline_data=blob))
elif _is_gcs_url(element):
mime_type = _infer_mime_type_from_gcs_url(element)
file_data: FileDataType = {
"mime_type": mime_type,
"file_uri": element,
}
parts.append(PartType(file_data=file_data))
elif _is_file_reference(element):
if element not in resolved_files:
raise ValueError(f"File reference {element} not resolved")
file_info = resolved_files[element]
file_data_ref: FileDataType = {
"mime_type": file_info["mime_type"],
"file_uri": file_info["uri"],
}
parts.append(PartType(file_data=file_data_ref))
else:
parts.append(PartType(text=element))
parts.append(_build_part_for_input(element, resolved_files=resolved_files))
request_body: dict = {
"content": ContentType(parts=parts),
@ -245,7 +295,7 @@ def transform_openai_input_gemini_embed_content(
def process_embed_content_response(
input: EmbeddingInput,
input: GeminiEmbeddingInput,
model_response: EmbeddingResponse,
model: str,
response_json: dict,
@ -291,7 +341,7 @@ def process_embed_content_response(
def process_response(
input: EmbeddingInput,
input: GeminiEmbeddingInput,
model_response: EmbeddingResponse,
model: str,
_predictions: VertexAIBatchEmbeddingsResponseObject,
@ -308,8 +358,29 @@ def process_response(
model_response.data = openai_embeddings
model_response.model = model
input_text = get_formatted_prompt(data={"input": input}, call_type="embedding")
prompt_tokens = token_counter(model=model, text=input_text)
has_nested = isinstance(input, list) and any(isinstance(e, list) for e in input)
if _is_multimodal_input(input) or has_nested:
input_list = input if isinstance(input, list) else [input]
text_elements: List[str] = []
for e in input_list:
if isinstance(e, list):
text_elements.extend(
sub
for sub in e
if isinstance(sub, str) and not _is_multimodal_element(sub)
)
elif isinstance(e, str) and not _is_multimodal_element(e):
text_elements.append(e)
if text_elements:
input_text = get_formatted_prompt(
data={"input": text_elements}, call_type="embedding"
)
prompt_tokens = token_counter(model=model, text=input_text)
else:
prompt_tokens = 0
else:
input_text = get_formatted_prompt(data={"input": input}, call_type="embedding")
prompt_tokens = token_counter(model=model, text=input_text)
model_response.usage = Usage(
prompt_tokens=prompt_tokens, total_tokens=prompt_tokens
)

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,

View file

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

View file

@ -33351,6 +33351,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",

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(

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"

View file

@ -668,6 +668,8 @@ class LiteLLMRoutes(enum.Enum):
"/models/{model_id}",
"/guardrails/list",
"/v2/guardrails/list",
"/project/list",
"/project/info",
]
+ spend_tracking_routes
+ key_management_routes
@ -692,6 +694,9 @@ class LiteLLMRoutes(enum.Enum):
"/model/{model_id}/update",
"/prompt/list",
"/prompt/info",
# Project read routes - endpoint scopes results to caller's teams (non-admin)
"/project/list",
"/project/info",
# Invitation routes - org/team admins checked in endpoint via _user_has_admin_privileges
"/invitation/new",
"/invitation/delete",

View file

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

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,

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

View file

@ -53,12 +53,16 @@ def clear_token() -> None:
os.remove(token_file)
def get_stored_api_key() -> Optional[str]:
"""Get the stored API key from token file"""
# Use the SDK-level utility
def get_stored_api_key(expected_base_url: Optional[str] = None) -> Optional[str]:
"""Get the stored API key from token file.
If expected_base_url is provided, the key is only returned when it was
originally issued for that URL. This prevents credential leakage when the
CLI is pointed at a different (possibly malicious) server.
"""
from litellm.litellm_core_utils.cli_token_utils import get_litellm_gateway_api_key
return get_litellm_gateway_api_key()
return get_litellm_gateway_api_key(expected_base_url=expected_base_url)
# Team selection utilities
@ -572,9 +576,11 @@ def login(ctx: click.Context):
api_key = auth_result["api_key"]
user_id = auth_result["user_id"]
# Save token data (simplified for CLI - we just need the key)
# Save token data. base_url is stored so we can verify origin
# before reusing the key on a subsequent CLI invocation.
save_token(
{
"base_url": base_url.rstrip("/"),
"key": api_key,
"user_id": user_id or "cli-user",
"user_email": "unknown",

View file

@ -74,9 +74,10 @@ def cli(ctx: click.Context, base_url: str, api_key: Optional[str]) -> None:
"""LiteLLM Proxy CLI - Manage your LiteLLM proxy server"""
ctx.ensure_object(dict)
# If no API key provided via flag or environment variable, try to load from saved token
# If no API key provided via flag or environment variable, try to load from saved token.
# Pass base_url so we only use the stored key when it was issued for this server.
if api_key is None:
api_key = get_stored_api_key()
api_key = get_stored_api_key(expected_base_url=base_url)
ctx.obj["base_url"] = base_url
ctx.obj["api_key"] = api_key

View file

@ -28,12 +28,17 @@ class Client:
api_key (Optional[str]): API key for authentication. If provided, it will be sent as a Bearer token.
timeout: Request timeout in seconds (default: 30)
"""
self._base_url = base_url.rstrip("/") # Remove trailing slash if present
self._api_key = get_litellm_gateway_api_key() or api_key
self._base_url = base_url.rstrip("/")
# Only use the stored CLI key when it was issued for this server.
self._api_key = api_key or get_litellm_gateway_api_key(
expected_base_url=self._base_url
)
# Initialize resource clients
self.http = HTTPClient(base_url=base_url, api_key=api_key, timeout=timeout)
self.http = HTTPClient(
base_url=base_url, api_key=self._api_key, timeout=timeout
)
self.models = ModelsManagementClient(
base_url=self._base_url, api_key=self._api_key
)

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

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

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

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

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
)

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)

View file

@ -29,6 +29,10 @@ ILLEGAL_DISPLAY_PARAMS = [
"exception", # internal; not JSON-serializable, never for display
"litellm_metadata", # internal tracking metadata with auth objects; not for display
]
# Provider routing fields. Allowed for proxy admins so they can see which
# region/version a deployment is checking; gated at the endpoint layer for
# non-admin callers (see _strip_admin_only_fields_from_health_result).
ADMIN_ONLY_HEALTH_DISPLAY_PARAMS = ("api_base", "api_version")
MINIMAL_DISPLAY_PARAMS = ["model", "mode_error"]

View file

@ -20,6 +20,7 @@ from litellm.proxy._types import (
CallInfo,
EnterpriseLicenseData,
Litellm_EntityType,
LitellmUserRoles,
ProxyErrorTypes,
ProxyException,
UserAPIKeyAuth,
@ -28,6 +29,7 @@ from litellm.proxy._types import (
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler
from litellm.proxy.health_check import (
ADMIN_ONLY_HEALTH_DISPLAY_PARAMS,
_clean_endpoint_data,
_update_litellm_params_for_health_check,
perform_health_check,
@ -723,6 +725,90 @@ async def _save_background_health_checks_to_db(
# Continue execution - don't let database save failure break health checks
_PROXY_ADMIN_ROLES = frozenset(
{
LitellmUserRoles.PROXY_ADMIN.value,
# View-only admins are operators (oncall, support); they need the
# routing fields (api_base, api_version) to diagnose health and tell
# which provider region a check is hitting. They cannot mutate config
# so granting them the read-only view is safe.
LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY.value,
}
)
def _is_proxy_admin(user_api_key_dict: UserAPIKeyAuth) -> bool:
"""
Return True if the caller has a proxy-admin role (full or view-only).
user_role on UserAPIKeyAuth can be either a LitellmUserRoles enum or its
string value depending on how the auth path constructed the object, so we
compare against the raw value rather than the enum identity.
"""
role = user_api_key_dict.user_role
if role is None:
return False
role_value = role.value if hasattr(role, "value") else role
return role_value in _PROXY_ADMIN_ROLES
def _strip_admin_only_fields_from_health_result(result: dict) -> dict:
"""
Return a copy of the /health response with provider routing fields
(``api_base``, ``api_version``) removed from each healthy/unhealthy
endpoint entry. Used to hide those fields from non-admin callers while
still showing them which deployments they own and whether each one is
healthy. Proxy admins receive the unmodified result.
"""
out = dict(result)
drop = set(ADMIN_ONLY_HEALTH_DISPLAY_PARAMS)
for key in ("healthy_endpoints", "unhealthy_endpoints"):
eps = out.get(key)
if isinstance(eps, list):
out[key] = [
(
{k: v for k, v in ep.items() if k not in drop}
if isinstance(ep, dict)
else ep
)
for ep in eps
]
return out
def _filter_health_check_results_by_model_ids(
results: dict, allowed_model_ids: set
) -> dict:
"""
Restrict a cached background health-check result dict to endpoints whose
model_id is in ``allowed_model_ids``.
Endpoints without a model_id (e.g. CLI-model entries that predate the
model_id wiring) are dropped conservatively we cannot prove they belong
to the caller, so they are excluded rather than leaked.
Each retained endpoint is shallow-copied before being returned, so any
downstream transform (e.g. _strip_admin_only_fields_from_health_result)
cannot accidentally mutate the shared ``health_check_results`` cache.
"""
healthy = [
dict(ep)
for ep in (results.get("healthy_endpoints") or [])
if ep.get("model_id") in allowed_model_ids
]
unhealthy = [
dict(ep)
for ep in (results.get("unhealthy_endpoints") or [])
if ep.get("model_id") in allowed_model_ids
]
return {
"healthy_endpoints": healthy,
"unhealthy_endpoints": unhealthy,
"healthy_count": len(healthy),
"unhealthy_count": len(unhealthy),
}
async def _perform_health_check_and_save(
model_list,
target_model,
@ -771,6 +857,7 @@ async def _perform_health_check_and_save(
@router.get("/health", tags=["health"], dependencies=[Depends(user_api_key_auth)])
async def health_endpoint(
response: Response,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
model: Optional[str] = fastapi.Query(
None, description="Specify the model name (optional)"
@ -838,11 +925,26 @@ async def health_endpoint(
detail={"error": f"Model with ID {model_id} not found"},
)
is_admin = _is_proxy_admin(user_api_key_dict)
def _post_process(result: dict) -> dict:
# api_base / api_version reveal which provider/region/internal host the
# deployment talks to; only proxy admins receive them. Non-admin keys
# still see model/model_id and the healthy/unhealthy status. We also
# set a header so non-admin clients that previously parsed those
# fields can detect the change programmatically.
if is_admin:
return result
response.headers["Litellm-Health-Field-Notice"] = (
"api_base and api_version are admin-only on this endpoint"
)
return _strip_admin_only_fields_from_health_result(result)
try:
if llm_model_list is None:
# if no router set, check if user set a model using litellm --model ollama/llama2
if user_model is not None:
return await _perform_health_check_and_save(
cli_result = await _perform_health_check_and_save(
model_list=[],
target_model=None,
cli_model=user_model,
@ -853,20 +955,59 @@ async def health_endpoint(
model_id=None, # CLI model doesn't have model_id
max_concurrency=health_check_concurrency,
)
return _post_process(cli_result)
raise HTTPException(
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
detail={"error": "Model list not initialized"},
)
_llm_model_list = copy.deepcopy(llm_model_list)
### FILTER MODELS FOR ONLY THOSE USER HAS ACCESS TO ###
# Live path: scope by model_name (every deployment has one).
# Cache path: scope by model_id (the cache is keyed on model_id).
# Consequence: a deployment whose model_name the caller can access
# but which lacks model_info.id will appear in the live /health
# response but NOT in the background-cache /health response. This is
# surfaced via the "warnings" field below so operators can fix the
# missing model_info.id rather than guess at the discrepancy.
if len(user_api_key_dict.models) > 0:
pass
else:
pass #
allowed_models = set(user_api_key_dict.models)
_llm_model_list = [
m for m in _llm_model_list if m.get("model_name") in allowed_models
]
if use_background_health_checks:
return health_check_results
if len(user_api_key_dict.models) > 0:
allowed_model_ids = {
(m.get("model_info") or {}).get("id")
for m in _llm_model_list
if (m.get("model_info") or {}).get("id")
}
filtered = _filter_health_check_results_by_model_ids(
health_check_results, allowed_model_ids
)
if not allowed_model_ids:
# Caller has accessible model_names but none of the
# matching deployments expose a model_info.id, so the
# cache filter (which keys on model_id) drops every
# entry. Surface this both as a warning log and a
# structured "warnings" field on the response so the
# caller can distinguish "no deployments found" from
# "deployments excluded due to missing model_info.id".
verbose_proxy_logger.warning(
"health_endpoint: scoped key %s has accessible models %s "
"but none of the matching deployments carry a model_info.id; "
"background health-check cache will return an empty result.",
user_api_key_dict.user_id,
list(user_api_key_dict.models),
)
filtered["warnings"] = [
"Some accessible deployments are missing model_info.id "
"and were excluded from this response. Ask a proxy admin "
"to populate model_info.id for these models."
]
return _post_process(filtered)
return _post_process(health_check_results)
else:
return await _perform_health_check_and_save(
router_result = await _perform_health_check_and_save(
model_list=_llm_model_list,
target_model=target_model,
cli_model=None,
@ -877,6 +1018,7 @@ async def health_endpoint(
model_id=model_id,
max_concurrency=health_check_concurrency,
)
return _post_process(router_result)
except Exception as e:
verbose_proxy_logger.error(
"litellm.proxy.proxy_server.py::health_endpoint(): Exception occured - {}".format(

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
]

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

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

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

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)

View file

@ -1,6 +1,7 @@
import asyncio
import json
import time
import urllib.parse
from datetime import datetime
from typing import Literal, Optional
from urllib.parse import urlparse
@ -203,8 +204,16 @@ class AssemblyAIPassthroughLoggingHandler:
)
if _api_key is None:
raise ValueError("AssemblyAI API key not found")
if (
any(c in transcript_id for c in ("/", "\\", "#", "?"))
or ".." in transcript_id
):
raise ValueError(
f"Invalid transcript_id {transcript_id!r}: contains disallowed characters"
)
safe_transcript_id = urllib.parse.quote(transcript_id, safe="")
try:
url = f"{_base_url}/v2/transcript/{transcript_id}"
url = f"{_base_url}/v2/transcript/{safe_transcript_id}"
headers = {
"Authorization": f"Bearer {_api_key}",
"Content-Type": "application/json",

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(
@ -2014,14 +2019,16 @@ async def update_cache( # noqa: PLR0915
else:
hashed_token = token
verbose_proxy_logger.debug("_update_key_cache: hashed_token=%s", hashed_token)
existing_spend_obj: LiteLLM_VerificationTokenView = await user_api_key_cache.async_get_cache(key=hashed_token) # type: ignore
existing_spend_obj = await user_api_key_cache.async_get_cache(
key=hashed_token, model_type=UserAPIKeyAuth
)
verbose_proxy_logger.debug(
f"_update_key_cache: existing_spend_obj={existing_spend_obj}"
)
if existing_spend_obj is None:
return
else:
existing_spend = existing_spend_obj.spend
existing_spend = existing_spend_obj.spend or 0.0
# Calculate the new cost by adding the existing cost and response_cost
new_spend = existing_spend + response_cost
@ -2079,41 +2086,48 @@ async def update_cache( # noqa: PLR0915
existing_team_member_spend + response_cost
)
# Update the cost column for the given token
# Existing spend_obj is mutated; UserApiKeyCache.async_set_cache_pipeline turns
# BaseModel values into dicts for Redis (same Codec path as async_set_cache).
existing_spend_obj.spend = new_spend
values_to_update_in_cache.append((hashed_token, existing_spend_obj))
### UPDATE USER SPEND ###
async def _update_user_cache():
## UPDATE CACHE FOR USER ID + GLOBAL PROXY
if response_cost is None:
return
user_ids = [user_id]
try:
for _id in user_ids:
# Fetch the existing cost for the given user
if _id is None:
continue
existing_spend_obj = await user_api_key_cache.async_get_cache(key=_id)
if existing_spend_obj is None:
cached_user = await user_api_key_cache.async_get_cache(key=_id)
if cached_user is None:
# do nothing if there is no cache value
return
existing_spend_obj = CacheCodec.deserialize(
cached_user, LiteLLM_UserTable
)
if existing_spend_obj is None:
return
verbose_proxy_logger.debug(
f"_update_user_db: existing spend: {existing_spend_obj}; response_cost: {response_cost}"
)
if isinstance(existing_spend_obj, dict):
existing_spend = existing_spend_obj["spend"]
else:
existing_spend = existing_spend_obj.spend
existing_spend = existing_spend_obj.spend or 0.0
# Calculate the new cost by adding the existing cost and response_cost
new_spend = existing_spend + response_cost
# Update the cost column for the given user
if isinstance(existing_spend_obj, dict):
existing_spend_obj["spend"] = new_spend
values_to_update_in_cache.append((_id, existing_spend_obj))
else:
existing_spend_obj.spend = new_spend
values_to_update_in_cache.append((_id, existing_spend_obj.json()))
existing_spend_obj.spend = new_spend
values_to_update_in_cache.append(
(
_id,
CacheCodec.serialize(
existing_spend_obj, model_type=LiteLLM_UserTable
),
)
)
## UPDATE GLOBAL PROXY ##
global_proxy_spend = await user_api_key_cache.async_get_cache(
key="{}:spend".format(litellm_proxy_admin_name)
@ -2145,31 +2159,33 @@ async def update_cache( # noqa: PLR0915
_id = "end_user_id:{}".format(end_user_id)
try:
# Fetch the existing cost for the given user
existing_spend_obj = await user_api_key_cache.async_get_cache(key=_id)
if existing_spend_obj is None:
cached_end_user = await user_api_key_cache.async_get_cache(key=_id)
if cached_end_user is None:
# if user does not exist in LiteLLM_UserTable, create a new user
# do nothing if end-user not in api key cache
return
existing_spend_obj = CacheCodec.deserialize(
cached_end_user, LiteLLM_EndUserTable
)
if existing_spend_obj is None:
return
verbose_proxy_logger.debug(
f"_update_end_user_db: existing spend: {existing_spend_obj}; response_cost: {response_cost}"
)
if existing_spend_obj is None:
existing_spend = 0
else:
if isinstance(existing_spend_obj, dict):
existing_spend = existing_spend_obj["spend"]
else:
existing_spend = existing_spend_obj.spend
existing_spend = existing_spend_obj.spend or 0.0
# Calculate the new cost by adding the existing cost and response_cost
new_spend = existing_spend + response_cost
# Update the cost column for the given user
if isinstance(existing_spend_obj, dict):
existing_spend_obj["spend"] = new_spend
values_to_update_in_cache.append((_id, existing_spend_obj))
else:
existing_spend_obj.spend = new_spend
values_to_update_in_cache.append((_id, existing_spend_obj.json()))
existing_spend_obj.spend = new_spend
values_to_update_in_cache.append(
(
_id,
CacheCodec.serialize(
existing_spend_obj, model_type=LiteLLM_EndUserTable
),
)
)
except Exception as e:
verbose_proxy_logger.warning(
"Spend tracking - failed to update end user spend in cache. "
@ -2188,36 +2204,32 @@ async def update_cache( # noqa: PLR0915
_id = "team_id:{}".format(team_id)
try:
# Fetch the existing cost for the given user
existing_spend_obj: Optional[LiteLLM_TeamTable] = (
await user_api_key_cache.async_get_cache(key=_id)
cached_team = await user_api_key_cache.async_get_cache(key=_id)
if cached_team is None:
# do nothing if team not in api key cache
return
existing_spend_obj: Optional[LiteLLM_TeamTableCachedObj] = (
CacheCodec.deserialize(cached_team, LiteLLM_TeamTableCachedObj)
)
if existing_spend_obj is None:
# do nothing if team not in api key cache
return
verbose_proxy_logger.debug(
f"_update_team_db: existing spend: {existing_spend_obj}; response_cost: {response_cost}"
)
if existing_spend_obj is None:
existing_spend: Optional[float] = 0.0
else:
if isinstance(existing_spend_obj, dict):
existing_spend = existing_spend_obj["spend"]
else:
existing_spend = existing_spend_obj.spend
if existing_spend is None:
existing_spend = 0.0
existing_spend: float = existing_spend_obj.spend or 0.0
# Calculate the new cost by adding the existing cost and response_cost
new_spend = existing_spend + response_cost
# Update the cost column for the given user
if isinstance(existing_spend_obj, dict):
existing_spend_obj["spend"] = new_spend
values_to_update_in_cache.append((_id, existing_spend_obj))
else:
existing_spend_obj.spend = new_spend
values_to_update_in_cache.append((_id, existing_spend_obj))
existing_spend_obj.spend = new_spend
values_to_update_in_cache.append(
(
_id,
CacheCodec.serialize(
existing_spend_obj, model_type=LiteLLM_TeamTableCachedObj
),
)
)
except Exception as e:
verbose_proxy_logger.warning(
"Spend tracking - failed to update team spend in cache. "
@ -2244,32 +2256,32 @@ async def update_cache( # noqa: PLR0915
cache_key = f"tag:{tag_name}"
# Fetch the existing tag object from cache
existing_tag_obj = await user_api_key_cache.async_get_cache(
key=cache_key
)
if existing_tag_obj is None:
cached_tag = await user_api_key_cache.async_get_cache(key=cache_key)
if cached_tag is None:
# do nothing if tag not in api key cache
continue
existing_tag_obj = CacheCodec.deserialize(cached_tag, LiteLLM_TagTable)
if existing_tag_obj is None:
continue
verbose_proxy_logger.debug(
f"_update_tag_cache: existing spend for tag={tag_name}: {existing_tag_obj}; response_cost: {response_cost}"
)
if isinstance(existing_tag_obj, dict):
existing_spend = existing_tag_obj.get("spend", 0) or 0
else:
existing_spend = getattr(existing_tag_obj, "spend", 0) or 0
existing_spend = existing_tag_obj.spend or 0.0
# Calculate the new cost by adding the existing cost and response_cost
new_spend = existing_spend + response_cost
# Update the spend column for the given tag
if isinstance(existing_tag_obj, dict):
existing_tag_obj["spend"] = new_spend
values_to_update_in_cache.append((cache_key, existing_tag_obj))
else:
existing_tag_obj.spend = new_spend
values_to_update_in_cache.append((cache_key, existing_tag_obj))
existing_tag_obj.spend = new_spend
values_to_update_in_cache.append(
(
cache_key,
CacheCodec.serialize(
existing_tag_obj, model_type=LiteLLM_TagTable
),
)
)
except Exception as e:
verbose_proxy_logger.warning(
"Spend tracking - failed to update tag spend in cache. "
@ -2937,8 +2949,9 @@ class ProxyConfig:
def _init_cache(
self,
cache_params: dict,
enable_redis_auth_cache: bool = False,
):
global redis_usage_cache, llm_router
global redis_usage_cache, llm_router, general_settings
from litellm import Cache
if "default_in_memory_ttl" in cache_params:
@ -2954,7 +2967,29 @@ class ProxyConfig:
):
## INIT PROXY REDIS USAGE CLIENT ##
redis_usage_cache = litellm.cache.cache
spend_counter_cache.redis_cache = redis_usage_cache
spend_counter_cache.attach_redis_cache(
redis_usage_cache,
default_redis_ttl=litellm.default_redis_ttl,
)
# Note: PKCE verifier storage uses redis_usage_cache directly (not
# user_api_key_cache) to avoid routing all API-key lookups through Redis.
if enable_redis_auth_cache is True:
user_api_key_cache.attach_redis_cache(
redis_usage_cache,
default_redis_ttl=litellm.default_redis_ttl,
)
verbose_proxy_logger.info(
"enable_redis_auth_cache=True: attached Redis to "
"user_api_key_cache — virtual-key lookups are now "
"shared across all proxy workers."
)
else:
verbose_proxy_logger.info(
"enable_redis_auth_cache is not set: user_api_key_cache "
"remains in-memory only (per-worker). Set "
"litellm_settings.enable_redis_auth_cache: true to share "
"the auth cache across workers and reduce DB load."
)
litellm_config_cache.redis_cache = redis_usage_cache
# Note: PKCE verifier storage uses redis_usage_cache directly (not
# user_api_key_cache) to avoid routing all API-key lookups through Redis.
@ -3280,7 +3315,13 @@ class ProxyConfig:
cache_params[key] = get_secret(value)
## to pass a complete url, or set ssl=True, etc. just set it as `os.environ[REDIS_URL] = <your-redis-url>`, _redis.py checks for REDIS specific environment variables
self._init_cache(cache_params=cache_params)
self._init_cache(
cache_params=cache_params,
enable_redis_auth_cache=litellm_settings.get(
"enable_redis_auth_cache", False
)
is True,
)
if litellm.cache is not None:
verbose_proxy_logger.debug(
f"{blue_color_code}Set Cache on LiteLLM Proxy{reset_color_code}"
@ -3551,21 +3592,23 @@ class ProxyConfig:
verbose_proxy_logger.critical(
"LITELLM_MASTER_KEY is not set! All requests will be treated as INTERNAL_USER with no admin access. Set LITELLM_MASTER_KEY for production use."
)
### USER API KEY CACHE IN-MEMORY TTL ###
### USER API KEY CACHE TTL (in-memory + Redis when Redis auth sharing is enabled) ###
user_api_key_cache_ttl = general_settings.get(
"user_api_key_cache_ttl", None
)
if user_api_key_cache_ttl is not None:
ttl = float(user_api_key_cache_ttl)
# Mirror TTL on Redis as well when ``litellm_settings.enable_redis_auth_cache``
# attaches Redis to ``user_api_key_cache``; otherwise DualCache misses in
# memory fall back to a key that outlasts ``user_api_key_cache_ttl``.
user_api_key_cache.update_cache_ttl(
default_in_memory_ttl=float(user_api_key_cache_ttl),
default_redis_ttl=None, # user_api_key_cache uses in-memory TTL only; Redis not configured for key lookups
default_in_memory_ttl=ttl,
default_redis_ttl=ttl,
)
### PKCE MULTI-INSTANCE PREREQUISITE CHECK ###
# PKCE verifiers are stored in redis_usage_cache when available so they can
# be read back by any instance (not just the one that started the auth flow).
# user_api_key_cache is intentionally left in-memory-only to avoid routing
# all API-key lookups through Redis.
use_pkce = os.getenv("GENERIC_CLIENT_USE_PKCE", "false").lower() == "true"
if use_pkce and redis_usage_cache is None:
global _pkce_no_redis_warning_emitted
@ -6294,7 +6337,7 @@ class ProxyStartupEvent:
cls,
general_settings: dict,
prisma_client: Optional[PrismaClient],
user_api_key_cache: DualCache,
user_api_key_cache: UserApiKeyCache,
):
"""Initialize JWT auth on startup"""
if general_settings.get("litellm_jwtauth", None) is not None:
@ -6343,7 +6386,7 @@ class ProxyStartupEvent:
async def _warm_global_spend_cache(
cls,
litellm_proxy_admin_name: str,
user_api_key_cache: DualCache,
user_api_key_cache: UserApiKeyCache,
prisma_client: PrismaClient,
) -> None:
"""Warm global spend cache once at startup to reduce impact of first wave of requests."""
@ -6983,7 +7026,7 @@ class ProxyStartupEvent:
cls,
database_url: Optional[str],
proxy_logging_obj: ProxyLogging,
user_api_key_cache: DualCache,
user_api_key_cache: UserApiKeyCache,
) -> Optional[PrismaClient]:
"""
- Sets up prisma client
@ -10285,6 +10328,101 @@ def _paginate_models_response(
}
def _team_models_resolve_to_names(
team_models: List[str], access_groups: Dict[str, Any]
) -> List[str]:
"""Expand team model entries (including access group names) to concrete model names."""
resolved: List[str] = []
for name in team_models:
if name in access_groups:
resolved.extend(access_groups[name])
else:
resolved.append(name)
return resolved
async def _load_team_object_for_model_filter(
team_id: str, prisma_client: PrismaClient
) -> Optional[LiteLLM_TeamTable]:
"""Load team row from DB; returns None if missing or on error."""
try:
team_db_object = await prisma_client.db.litellm_teamtable.find_unique(
where={"team_id": team_id}
)
if team_db_object is None:
verbose_proxy_logger.warning(f"Team {team_id} not found in database")
return None
return LiteLLM_TeamTable(**team_db_object.model_dump())
except Exception as e:
verbose_proxy_logger.exception(f"Error fetching team {team_id}: {str(e)}")
return None
async def _gather_team_accessible_model_ids(
team_object: LiteLLM_TeamTable,
team_id: str,
prisma_client: PrismaClient,
llm_router: Router,
) -> Set[str]:
"""Collect model IDs the team can use from router config and DB."""
team_accessible_model_ids: Set[str] = set()
access_groups = llm_router.get_model_access_groups() if llm_router else {}
if (
not team_object.models
or SpecialModelNames.all_proxy_models.value in team_object.models
):
model_list = llm_router.get_model_list() if llm_router else []
if model_list is not None:
for model in model_list:
model_id = model.get("model_info", {}).get("id", None)
if model_id is None:
continue
team_model_id = model.get("model_info", {}).get("team_id", None)
if team_model_id is None or team_model_id == team_id:
team_accessible_model_ids.add(model_id)
else:
resolved_model_names: Set[str] = set()
for model_name in team_object.models:
if model_name in access_groups:
resolved_model_names.update(access_groups[model_name])
else:
resolved_model_names.add(model_name)
for model_name in resolved_model_names:
_models = (
llm_router.get_model_list(model_name=model_name, team_id=team_id)
if llm_router
else []
)
if _models is not None:
for model in _models:
model_id = model.get("model_info", {}).get("id", None)
if model_id is not None:
team_accessible_model_ids.add(model_id)
try:
if (
team_object.models
and SpecialModelNames.all_proxy_models.value not in team_object.models
):
_resolved_names = _team_models_resolve_to_names(
team_object.models, access_groups
)
db_models = await prisma_client.db.litellm_proxymodeltable.find_many(
where={"model_name": {"in": _resolved_names}}
)
for db_model in db_models:
if db_model.model_id:
team_accessible_model_ids.add(db_model.model_id)
except Exception as e:
verbose_proxy_logger.debug(
f"Error querying database models for team {team_id}: {str(e)}"
)
return team_accessible_model_ids
async def _filter_models_by_team_id(
all_models: List[Dict[str, Any]],
team_id: str,
@ -10307,78 +10445,13 @@ async def _filter_models_by_team_id(
Returns:
Filtered list of models
"""
# Get team from database
try:
team_db_object = await prisma_client.db.litellm_teamtable.find_unique(
where={"team_id": team_id}
)
if team_db_object is None:
verbose_proxy_logger.warning(f"Team {team_id} not found in database")
# If team doesn't exist, return empty list
return []
team_object = LiteLLM_TeamTable(**team_db_object.model_dump())
except Exception as e:
verbose_proxy_logger.exception(f"Error fetching team {team_id}: {str(e)}")
team_object = await _load_team_object_for_model_filter(team_id, prisma_client)
if team_object is None:
return []
# Get models accessible to this team (similar to _add_team_models_to_all_models)
team_accessible_model_ids: Set[str] = set()
if (
not team_object.models # empty list = all model access
or SpecialModelNames.all_proxy_models.value in team_object.models
):
# Team has access to all models
model_list = llm_router.get_model_list() if llm_router else []
if model_list is not None:
for model in model_list:
model_id = model.get("model_info", {}).get("id", None)
if model_id is None:
continue
# if team model id set, check if team id matches
team_model_id = model.get("model_info", {}).get("team_id", None)
can_add_model = False
if team_model_id is None:
can_add_model = True
elif team_model_id == team_id:
can_add_model = True
if can_add_model:
team_accessible_model_ids.add(model_id)
else:
# Team has access to specific models
for model_name in team_object.models:
_models = (
llm_router.get_model_list(model_name=model_name, team_id=team_id)
if llm_router
else []
)
if _models is not None:
for model in _models:
model_id = model.get("model_info", {}).get("id", None)
if model_id is not None:
team_accessible_model_ids.add(model_id)
# Also search database for models accessible to this team
# This complements the config search done above
try:
if (
team_object.models
and SpecialModelNames.all_proxy_models.value not in team_object.models
):
# Team has specific models - check database for those model names
db_models = await prisma_client.db.litellm_proxymodeltable.find_many(
where={"model_name": {"in": team_object.models}}
)
for db_model in db_models:
model_id = db_model.model_id
if model_id:
team_accessible_model_ids.add(model_id)
except Exception as e:
verbose_proxy_logger.debug(
f"Error querying database models for team {team_id}: {str(e)}"
)
team_accessible_model_ids = await _gather_team_accessible_model_ids(
team_object, team_id, prisma_client, llm_router
)
# Filter models based on direct_access or access_via_team_ids
# Models are already enriched with these fields before this function is called

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.

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 5char 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

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(

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

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

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

View file

@ -6,6 +6,13 @@ from typing_extensions import (
TypedDict,
)
from litellm.types.llms.openai import EmbeddingInput
# Gemini supports nested-list inputs (e.g. [["text", "image"]]) as an explicit
# opt-in for combined embeddings — a provider-specific extension of the
# OpenAI-faithful EmbeddingInput shape.
GeminiEmbeddingInput = Union[EmbeddingInput, List[List[str]]]
class FunctionResponse(TypedDict):
name: str

View file

@ -33405,6 +33405,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",
@ -34842,6 +34908,20 @@
"supports_tool_choice": true,
"source": "https://aws.amazon.com/bedrock/pricing/"
},
"zai.glm-5": {
"input_cost_per_token": 1e-06,
"litellm_provider": "bedrock_converse",
"max_input_tokens": 200000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "chat",
"output_cost_per_token": 3.2e-06,
"supports_function_calling": true,
"supports_reasoning": true,
"supports_system_messages": true,
"supports_tool_choice": true,
"source": "https://aws.amazon.com/bedrock/pricing/"
},
"zai.glm-4.7-flash": {
"input_cost_per_token": 7e-08,
"litellm_provider": "bedrock_converse",

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

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

View file

@ -477,3 +477,4 @@ def test_get_llm_provider_use_proxy_arg_true_with_direct_args():
assert provider == "litellm_proxy"
assert key == arg_api_key # Should use the argument key
assert base == arg_api_base # Should use the argument base

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

View file

@ -134,3 +134,62 @@ def test_is_assemblyai_route():
== False
)
assert handler.is_assemblyai_route("") == False
# --- Security: SSRF via transcript_id path traversal ---
def test_get_assembly_transcript_rejects_slash_in_id(assembly_handler):
with patch(
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.passthrough_endpoint_router.get_credentials",
return_value="test-key",
):
with pytest.raises(ValueError, match="disallowed characters"):
assembly_handler._get_assembly_transcript("../../admin/credentials")
def test_get_assembly_transcript_rejects_dotdot_in_id(assembly_handler):
with patch(
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.passthrough_endpoint_router.get_credentials",
return_value="test-key",
):
with pytest.raises(ValueError, match="disallowed characters"):
assembly_handler._get_assembly_transcript("..evil")
def test_get_assembly_transcript_rejects_fragment_in_id(assembly_handler):
with patch(
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.passthrough_endpoint_router.get_credentials",
return_value="test-key",
):
with pytest.raises(ValueError, match="disallowed characters"):
assembly_handler._get_assembly_transcript("abc#suffix")
def test_get_assembly_transcript_rejects_query_in_id(assembly_handler):
with patch(
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.passthrough_endpoint_router.get_credentials",
return_value="test-key",
):
with pytest.raises(ValueError, match="disallowed characters"):
assembly_handler._get_assembly_transcript("abc?x=1")
def test_get_assembly_transcript_allows_valid_id(
assembly_handler, mock_transcript_response
):
with patch(
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.passthrough_endpoint_router.get_credentials",
return_value="test-key",
):
with patch("httpx.get") as mock_get:
mock_get.return_value.json.return_value = mock_transcript_response
mock_get.return_value.raise_for_status.return_value = None
transcript = assembly_handler._get_assembly_transcript(
"abc123-valid-id_xyz"
)
assert transcript == mock_transcript_response
called_url = mock_get.call_args[0][0]
assert "abc123-valid-id_xyz" in called_url
assert ".." not in called_url

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,

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)

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"

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

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"

View file

@ -280,3 +280,55 @@ class TestDynamicProjectNameOnSpan:
if __name__ == "__main__":
unittest.main()
# --- Security: SSRF via prompt_version_id path traversal ---
def test_arize_phoenix_client_sanitize_id_rejects_traversal():
from litellm.integrations.arize.arize_phoenix_client import _sanitize_id
# dotdot without slashes
with pytest.raises(ValueError, match="path traversal"):
_sanitize_id("..something")
# full traversal (slash caught first)
with pytest.raises(ValueError, match="disallowed characters"):
_sanitize_id("../../projects")
def test_arize_phoenix_client_sanitize_id_rejects_slash():
from litellm.integrations.arize.arize_phoenix_client import _sanitize_id
with pytest.raises(ValueError, match="disallowed characters"):
_sanitize_id("valid/extra")
def test_arize_phoenix_client_sanitize_id_rejects_fragment():
from litellm.integrations.arize.arize_phoenix_client import _sanitize_id
with pytest.raises(ValueError, match="disallowed characters"):
_sanitize_id("abc#suffix")
def test_arize_phoenix_client_sanitize_id_rejects_query():
from litellm.integrations.arize.arize_phoenix_client import _sanitize_id
with pytest.raises(ValueError, match="disallowed characters"):
_sanitize_id("abc?x=1")
def test_arize_phoenix_client_sanitize_id_allows_uuid():
from litellm.integrations.arize.arize_phoenix_client import _sanitize_id
uid = "550e8400-e29b-41d4-a716-446655440000"
assert _sanitize_id(uid) == uid
def test_arize_phoenix_client_get_prompt_version_rejects_traversal():
from litellm.integrations.arize.arize_phoenix_client import ArizePhoenixClient
client = ArizePhoenixClient(
api_key="test-key", api_base="https://app.phoenix.arize.com"
)
with pytest.raises(ValueError, match="disallowed characters"):
client.get_prompt_version("../../projects")

View file

@ -11,6 +11,7 @@ sys.path.insert(
import litellm
from litellm.integrations.bitbucket import BitBucketPromptManager
from litellm.integrations.bitbucket.bitbucket_client import _sanitize_file_path
@patch("litellm.integrations.bitbucket.bitbucket_prompt_manager.BitBucketClient")
@ -370,3 +371,45 @@ def test_bitbucket_prompt_manager_list_templates(mock_client_class):
templates = manager.prompt_manager.list_templates()
assert isinstance(templates, list)
assert "test_prompt" in templates
# --- Security: path traversal / SSRF ---
def test_sanitize_file_path_rejects_traversal():
with pytest.raises(ValueError, match="path traversal"):
_sanitize_file_path("../../etc/passwd")
def test_sanitize_file_path_rejects_fragment():
with pytest.raises(ValueError, match="URL special characters"):
_sanitize_file_path("secret#.prompt")
def test_sanitize_file_path_rejects_query():
with pytest.raises(ValueError, match="URL special characters"):
_sanitize_file_path("secret?.prompt")
def test_sanitize_file_path_encodes_special_chars():
result = _sanitize_file_path("prompts/my prompt.prompt")
assert result == "prompts/my%20prompt.prompt"
def test_sanitize_file_path_allows_normal_paths():
assert _sanitize_file_path("prompts/my-prompt") == "prompts/my-prompt"
assert _sanitize_file_path("simple") == "simple"
def test_bitbucket_client_rejects_traversal_in_get_file_content():
from litellm.integrations.bitbucket.bitbucket_client import BitBucketClient
client = BitBucketClient(
{
"workspace": "ws",
"repository": "repo",
"access_token": "tok",
}
)
with pytest.raises(ValueError, match="path traversal"):
client.get_file_content("../../admin/credentials")

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

View file

@ -0,0 +1,290 @@
"""
Tests for Gemini batchEmbedContents transformation logic.
Covers:
- Text-only inputs (single and batch)
- Multimodal inputs (data URIs, GCS URLs, file references)
- Mixed text + multimodal inputs
- Response processing with correct indices
"""
import pytest
from litellm.llms.vertex_ai.gemini_embeddings.batch_embed_content_transformation import (
_build_part_for_input,
_is_multimodal_input,
process_response,
transform_openai_input_gemini_content,
transform_openai_input_gemini_embed_content,
)
from litellm.types.llms.vertex_ai import VertexAIBatchEmbeddingsResponseObject
from litellm.types.utils import EmbeddingResponse
IMAGE_DATA_URI = "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAgAAAAIAQMAAAD+wSzIAAAABlBMVEX///+/v7+jQ3Y5AAAADklEQVQI12P4AIX8EAgALgAD/aNpbtEAAAAASUVORK5CYII"
GCS_URL = "gs://my-bucket/image.png"
class TestIsMultimodalInput:
def test_text_only_string(self):
assert _is_multimodal_input("hello world") is False
def test_text_only_list(self):
assert _is_multimodal_input(["hello", "world"]) is False
def test_data_uri(self):
assert _is_multimodal_input([IMAGE_DATA_URI]) is True
def test_gcs_url(self):
assert _is_multimodal_input([GCS_URL]) is True
def test_file_reference(self):
assert _is_multimodal_input(["files/abc123"]) is True
def test_mixed_text_and_image(self):
assert _is_multimodal_input(["hello", IMAGE_DATA_URI]) is True
def test_nested_text_is_not_multimodal(self):
"""Nested list with text is not multimodal."""
assert _is_multimodal_input([["text_a", "text_b"]]) is False
def test_nested_list_with_image_is_multimodal(self):
assert _is_multimodal_input([["a red shoe", IMAGE_DATA_URI]]) is True
class TestBuildPartForInput:
def test_text_input(self):
part = _build_part_for_input("hello")
assert part["text"] == "hello"
assert part.get("inline_data") is None
def test_data_uri_input(self):
part = _build_part_for_input(IMAGE_DATA_URI)
assert part.get("text") is None
assert part["inline_data"] is not None
assert part["inline_data"]["mime_type"] == "image/png"
def test_gcs_url_input(self):
part = _build_part_for_input(GCS_URL)
assert part.get("text") is None
assert part["file_data"] is not None
assert part["file_data"]["mime_type"] == "image/png"
assert part["file_data"]["file_uri"] == GCS_URL
def test_file_reference_resolved(self):
resolved = {"files/abc": {"mime_type": "image/jpeg", "uri": "https://example.com/abc"}}
part = _build_part_for_input("files/abc", resolved_files=resolved)
assert part["file_data"] is not None
assert part["file_data"]["mime_type"] == "image/jpeg"
def test_file_reference_unresolved_raises(self):
with pytest.raises(ValueError, match="not resolved"):
_build_part_for_input("files/abc")
class TestTransformOpenaiInputGeminiContent:
"""Test that transform_openai_input_gemini_content creates separate requests per input."""
def test_single_text(self):
result = transform_openai_input_gemini_content(
input="hello", model="gemini-embedding-2-preview", optional_params={}
)
assert len(result["requests"]) == 1
assert result["requests"][0]["content"]["parts"][0]["text"] == "hello"
def test_multiple_texts(self):
result = transform_openai_input_gemini_content(
input=["hello", "world"], model="gemini-embedding-2-preview", optional_params={}
)
assert len(result["requests"]) == 2
assert result["requests"][0]["content"]["parts"][0]["text"] == "hello"
assert result["requests"][1]["content"]["parts"][0]["text"] == "world"
def test_multimodal_inputs_are_separate_requests(self):
"""Key regression test for #24209: each input becomes its own request."""
result = transform_openai_input_gemini_content(
input=["The food was delicious", IMAGE_DATA_URI],
model="gemini-embedding-2-preview",
optional_params={},
)
assert len(result["requests"]) == 2
# First request is text
assert result["requests"][0]["content"]["parts"][0]["text"] == "The food was delicious"
# Second request is image
assert result["requests"][1]["content"]["parts"][0]["inline_data"] is not None
def test_dimensions_mapped_to_output_dimensionality(self):
result = transform_openai_input_gemini_content(
input="hello",
model="gemini-embedding-2-preview",
optional_params={"dimensions": 256},
)
assert result["requests"][0]["outputDimensionality"] == 256
def test_model_name_prefixed(self):
result = transform_openai_input_gemini_content(
input="hello", model="gemini-embedding-2-preview", optional_params={}
)
assert result["requests"][0]["model"] == "models/gemini-embedding-2-preview"
def test_gcs_url_input(self):
result = transform_openai_input_gemini_content(
input=[GCS_URL], model="gemini-embedding-2-preview", optional_params={}
)
assert len(result["requests"]) == 1
assert result["requests"][0]["content"]["parts"][0]["file_data"] is not None
def test_mixed_text_image_gcs(self):
result = transform_openai_input_gemini_content(
input=["hello", IMAGE_DATA_URI, GCS_URL],
model="gemini-embedding-2-preview",
optional_params={},
)
assert len(result["requests"]) == 3
def test_nested_input_combined_embedding(self):
"""Nested list produces one request with multiple parts (combined embedding)."""
result = transform_openai_input_gemini_content(
input=[["a red shoe", IMAGE_DATA_URI]],
model="gemini-embedding-2-preview",
optional_params={},
)
assert len(result["requests"]) == 1
parts = result["requests"][0]["content"]["parts"]
assert len(parts) == 2
assert parts[0]["text"] == "a red shoe"
assert parts[1]["inline_data"] is not None
def test_mixed_nested_and_flat(self):
"""Mixed nested + flat produces correct number of requests."""
result = transform_openai_input_gemini_content(
input=[["text", IMAGE_DATA_URI], "standalone"],
model="gemini-embedding-2-preview",
optional_params={},
)
assert len(result["requests"]) == 2
# First: combined (2 parts)
assert len(result["requests"][0]["content"]["parts"]) == 2
# Second: standalone (1 part)
assert len(result["requests"][1]["content"]["parts"]) == 1
assert result["requests"][1]["content"]["parts"][0]["text"] == "standalone"
class TestTransformOpenaiInputGeminiEmbedContent:
"""Test transform_openai_input_gemini_embed_content (vertex_ai / embedContent path)."""
def test_text_and_image_combined(self):
result = transform_openai_input_gemini_embed_content(
input=["hello", IMAGE_DATA_URI],
model="gemini-embedding-2-preview",
optional_params={},
)
assert "content" in result
parts = result["content"]["parts"]
assert len(parts) == 2
assert parts[0]["text"] == "hello"
assert parts[1]["inline_data"] is not None
def test_gcs_url(self):
result = transform_openai_input_gemini_embed_content(
input=[GCS_URL],
model="gemini-embedding-2-preview",
optional_params={},
)
parts = result["content"]["parts"]
assert len(parts) == 1
assert parts[0]["file_data"]["file_uri"] == GCS_URL
def test_dimensions_mapped(self):
result = transform_openai_input_gemini_embed_content(
input="hello",
model="gemini-embedding-2-preview",
optional_params={"dimensions": 256},
)
assert result["outputDimensionality"] == 256
class TestProcessResponse:
"""Test that process_response sets correct indices."""
def test_single_embedding_index(self):
predictions: VertexAIBatchEmbeddingsResponseObject = {
"embeddings": [{"values": [0.1, 0.2]}]
}
model_response = EmbeddingResponse()
result = process_response(
input="hello",
model_response=model_response,
model="gemini-embedding-2-preview",
_predictions=predictions,
)
assert len(result.data) == 1
assert result.data[0]["index"] == 0
def test_multiple_embeddings_have_correct_indices(self):
"""Regression test: indices should be 0, 1, 2... not all 0."""
predictions: VertexAIBatchEmbeddingsResponseObject = {
"embeddings": [
{"values": [0.1, 0.2]},
{"values": [0.3, 0.4]},
{"values": [0.5, 0.6]},
]
}
model_response = EmbeddingResponse()
result = process_response(
input=["a", "b", "c"],
model_response=model_response,
model="gemini-embedding-2-preview",
_predictions=predictions,
)
assert len(result.data) == 3
assert result.data[0]["index"] == 0
assert result.data[1]["index"] == 1
assert result.data[2]["index"] == 2
def test_multimodal_mixed_input(self):
"""process_response works with mixed text + multimodal inputs."""
predictions: VertexAIBatchEmbeddingsResponseObject = {
"embeddings": [{"values": [0.1, 0.2]}, {"values": [0.3, 0.4]}]
}
result = process_response(
input=["hello", IMAGE_DATA_URI],
model_response=EmbeddingResponse(),
model="gemini-embedding-2-preview",
_predictions=predictions,
)
assert len(result.data) == 2
assert result.data[0]["index"] == 0
assert result.data[1]["index"] == 1
# Should count tokens only for the text element, not the image
assert result.usage.prompt_tokens > 0
def test_nested_input_token_counting(self):
"""Nested list: only plain-text sub-elements should be counted."""
predictions: VertexAIBatchEmbeddingsResponseObject = {
"embeddings": [{"values": [0.1, 0.2]}]
}
result = process_response(
input=[["a red shoe", IMAGE_DATA_URI]],
model_response=EmbeddingResponse(),
model="gemini-embedding-2-preview",
_predictions=predictions,
)
assert len(result.data) == 1
assert result.usage.prompt_tokens > 0
def test_nested_empty_list_raises(self):
with pytest.raises(ValueError, match="must not be empty"):
transform_openai_input_gemini_content(
input=[[]],
model="gemini-embedding-2-preview",
optional_params={},
)
def test_nested_non_string_element_raises(self):
with pytest.raises(ValueError, match="must be strings"):
transform_openai_input_gemini_content(
input=[[["doubly", "nested"]]],
model="gemini-embedding-2-preview",
optional_params={},
)

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

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"

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

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

View file

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

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

View file

@ -231,6 +231,50 @@ class TestTokenUtilities:
result = get_stored_api_key()
assert result is None
def test_get_stored_api_key_base_url_match(self):
"""Stored key is returned when expected_base_url matches stored origin"""
token_data = {"key": "sk-prod", "base_url": "https://real-proxy.com"}
with patch(
"litellm.litellm_core_utils.cli_token_utils.load_cli_token",
return_value=token_data,
):
assert (
get_stored_api_key(expected_base_url="https://real-proxy.com")
== "sk-prod"
)
def test_get_stored_api_key_base_url_match_trailing_slash(self):
"""Trailing slash on expected_base_url is normalised before comparison"""
token_data = {"key": "sk-prod", "base_url": "https://real-proxy.com"}
with patch(
"litellm.litellm_core_utils.cli_token_utils.load_cli_token",
return_value=token_data,
):
assert (
get_stored_api_key(expected_base_url="https://real-proxy.com/")
== "sk-prod"
)
def test_get_stored_api_key_base_url_mismatch(self):
"""Stored key is NOT returned when expected_base_url differs from stored origin"""
token_data = {"key": "sk-prod", "base_url": "https://real-proxy.com"}
with patch(
"litellm.litellm_core_utils.cli_token_utils.load_cli_token",
return_value=token_data,
):
assert get_stored_api_key(expected_base_url="https://evil.com") is None
def test_get_stored_api_key_old_token_no_base_url(self):
"""Old tokens without a base_url field are rejected when origin check is requested"""
token_data = {"key": "sk-old-token"}
with patch(
"litellm.litellm_core_utils.cli_token_utils.load_cli_token",
return_value=token_data,
):
assert (
get_stored_api_key(expected_base_url="https://real-proxy.com") is None
)
class TestLoginCommand:
"""Test login CLI command"""

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

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

View file

@ -778,3 +778,373 @@ def test_get_callback_identifier_custom_logger_registry_and_fallback():
result = get_callback_identifier(my_callback_function)
# Should fall back to callback_name() which returns __name__
assert result == "my_callback_function"
# ---------------------------------------------------------------------------
# /health response shape: model-access scoping and display-field allowlist
# ---------------------------------------------------------------------------
# These tests pin the contract that the /health response (a) only includes
# deployments the calling key is allowed to see, and (b) does not return
# provider routing fields like api_base / api_version. They guard against
# regressions that would widen the response shape.
@pytest.mark.asyncio
async def test_health_endpoint_filters_model_list_by_user_access():
"""
health_endpoint() should restrict _llm_model_list to deployments whose
model_name appears in user_api_key_dict.models before running the health
check. A key scoped to ["model-a"] should only see model-a in the result,
not other deployments configured on the proxy.
"""
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.health_endpoints._health_endpoints import health_endpoint
full_model_list = [
{
"model_name": "model-a",
"litellm_params": {
"model": "openai/gpt-4o",
"api_base": "https://example-a.test",
},
"model_info": {"id": "id-a"},
},
{
"model_name": "model-b",
"litellm_params": {
"model": "openai/gpt-4o",
"api_base": "https://example-b.test",
"api_version": "2024-10-21",
},
"model_info": {"id": "id-b"},
},
]
user_api_key_dict = UserAPIKeyAuth(
api_key="hashed-test-key",
models=["model-a"],
)
captured: dict = {}
async def fake_perform(**kwargs):
captured["model_list"] = kwargs["model_list"]
return {
"healthy_endpoints": [],
"unhealthy_endpoints": [],
"healthy_count": 0,
"unhealthy_count": 0,
}
with (
patch("litellm.proxy.proxy_server.llm_model_list", full_model_list),
patch("litellm.proxy.proxy_server.llm_router", None),
patch("litellm.proxy.proxy_server.prisma_client", None),
patch("litellm.proxy.proxy_server.use_background_health_checks", False),
patch("litellm.proxy.proxy_server.user_model", None),
patch("litellm.proxy.proxy_server.health_check_results", {}),
patch("litellm.proxy.proxy_server.health_check_details", True),
patch("litellm.proxy.proxy_server.health_check_concurrency", 1),
patch(
"litellm.proxy.health_endpoints._health_endpoints._perform_health_check_and_save",
side_effect=fake_perform,
),
):
from fastapi import Response
await health_endpoint(response=Response(), user_api_key_dict=user_api_key_dict)
assert (
"model_list" in captured
), "health_endpoint did not call _perform_health_check_and_save"
returned_names = {m["model_name"] for m in captured["model_list"]}
assert returned_names == {
"model-a"
}, f"health_endpoint did not scope model_list to caller access: {returned_names}"
@pytest.mark.asyncio
async def test_health_endpoint_filters_background_cache_by_user_access():
"""
When background_health_checks is enabled, health_endpoint() should also
scope the cached result to the caller's allowed models rather than
returning the cache verbatim.
"""
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.health_endpoints._health_endpoints import health_endpoint
full_model_list = [
{
"model_name": "model-a",
"litellm_params": {
"model": "openai/gpt-4o",
"api_base": "https://example-a.test",
},
"model_info": {"id": "id-a"},
},
{
"model_name": "model-b",
"litellm_params": {
"model": "openai/gpt-4o",
"api_base": "https://example-b.test",
},
"model_info": {"id": "id-b"},
},
]
cached_results = {
"healthy_endpoints": [
{
"model": "openai/gpt-4o",
"model_id": "id-a",
"api_base": "https://example-a.test",
},
{
"model": "openai/gpt-4o",
"model_id": "id-b",
"api_base": "https://example-b.test",
},
],
"unhealthy_endpoints": [],
"healthy_count": 2,
"unhealthy_count": 0,
}
user_api_key_dict = UserAPIKeyAuth(
api_key="hashed-test-key",
models=["model-a"],
)
with (
patch("litellm.proxy.proxy_server.llm_model_list", full_model_list),
patch("litellm.proxy.proxy_server.llm_router", None),
patch("litellm.proxy.proxy_server.prisma_client", None),
patch("litellm.proxy.proxy_server.use_background_health_checks", True),
patch("litellm.proxy.proxy_server.user_model", None),
patch("litellm.proxy.proxy_server.health_check_results", cached_results),
patch("litellm.proxy.proxy_server.health_check_details", True),
patch("litellm.proxy.proxy_server.health_check_concurrency", 1),
):
from fastapi import Response
result = await health_endpoint(
response=Response(), user_api_key_dict=user_api_key_dict
)
# Sanity: the source cache had two entries before scoping; the scoping
# step is what reduces it to one. (This guards against the test passing
# vacuously when the cache filter drops everything because cached
# entries lack the model_id key — both entries carry model_id above.)
assert len(cached_results["healthy_endpoints"]) == 2
assert all(
ep.get("model_id") for ep in cached_results["healthy_endpoints"]
), "test fixture invariant: every cached entry must carry a model_id"
# The non-admin caller must not see api_base on the returned cache entries.
returned = result.get("healthy_endpoints", [])
assert (
len(returned) == 1
), f"expected exactly one cached entry after scoping, got {len(returned)}"
assert returned[0]["model_id"] == "id-a"
assert "api_base" not in returned[0]
assert result["healthy_count"] == 1
assert result["unhealthy_count"] == 0
@pytest.mark.asyncio
async def test_health_endpoint_admin_sees_routing_fields_non_admin_does_not():
"""
A proxy admin should still see ``api_base`` and ``api_version`` in the
/health response so they can tell which Vertex region / Azure resource
+ API version is healthy. A non-admin caller must not both fields
should be stripped, and the response should carry a notice header so
non-admin clients can detect the change programmatically.
"""
from fastapi import Response
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
from litellm.proxy.health_endpoints._health_endpoints import health_endpoint
full_model_list = [
{
"model_name": "model-a",
"litellm_params": {
"model": "openai/gpt-4o",
"api_base": "https://example-a.test",
},
"model_info": {"id": "id-a"},
},
]
cached_results = {
"healthy_endpoints": [
{
"model": "openai/gpt-4o",
"model_id": "id-a",
"api_base": "https://us-central1-aiplatform.googleapis.com/v1/projects/p",
"api_version": "2024-10-21",
},
],
"unhealthy_endpoints": [],
"healthy_count": 1,
"unhealthy_count": 0,
}
admin_key = UserAPIKeyAuth(
api_key="hashed-admin-key",
models=["model-a"],
user_role=LitellmUserRoles.PROXY_ADMIN,
)
non_admin_key = UserAPIKeyAuth(
api_key="hashed-user-key",
models=["model-a"],
)
common_patches = [
patch("litellm.proxy.proxy_server.llm_model_list", full_model_list),
patch("litellm.proxy.proxy_server.llm_router", None),
patch("litellm.proxy.proxy_server.prisma_client", None),
patch("litellm.proxy.proxy_server.use_background_health_checks", True),
patch("litellm.proxy.proxy_server.user_model", None),
patch("litellm.proxy.proxy_server.health_check_results", cached_results),
patch("litellm.proxy.proxy_server.health_check_details", True),
patch("litellm.proxy.proxy_server.health_check_concurrency", 1),
]
for p in common_patches:
p.start()
try:
admin_response = Response()
non_admin_response = Response()
admin_result = await health_endpoint(
response=admin_response, user_api_key_dict=admin_key
)
non_admin_result = await health_endpoint(
response=non_admin_response, user_api_key_dict=non_admin_key
)
finally:
for p in common_patches:
p.stop()
admin_eps = admin_result.get("healthy_endpoints", [])
non_admin_eps = non_admin_result.get("healthy_endpoints", [])
assert len(admin_eps) == 1
assert (
admin_eps[0]["api_base"]
== "https://us-central1-aiplatform.googleapis.com/v1/projects/p"
), "admin must see the full api_base so they can identify the region"
assert (
admin_eps[0]["api_version"] == "2024-10-21"
), "admin must see api_version so they can distinguish provider deployments"
assert len(non_admin_eps) == 1
assert "api_base" not in non_admin_eps[0]
assert "api_version" not in non_admin_eps[0]
# Non-admin response must advertise that api_base/api_version were
# withheld so clients that previously parsed them can detect the change.
assert (
non_admin_response.headers.get("Litellm-Health-Field-Notice")
== "api_base and api_version are admin-only on this endpoint"
)
assert "Litellm-Health-Field-Notice" not in admin_response.headers
# Stripping must produce a copy — the shared cache must still carry the
# routing fields so the next admin caller can read them.
cached_first = cached_results["healthy_endpoints"][0]
assert (
cached_first["api_base"]
== "https://us-central1-aiplatform.googleapis.com/v1/projects/p"
)
assert cached_first["api_version"] == "2024-10-21"
@pytest.mark.asyncio
async def test_health_endpoint_warns_when_scoped_models_lack_model_id():
"""
When a scoped key's accessible models exist on the proxy but none of the
matching deployments expose a ``model_info.id``, the cache filter drops
everything. The response should include a structured ``warnings`` field
so the caller can distinguish "no deployments configured" from
"deployments excluded due to missing model_info.id".
"""
from fastapi import Response
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.health_endpoints._health_endpoints import health_endpoint
full_model_list = [
{
"model_name": "model-a",
"litellm_params": {
"model": "openai/gpt-4o",
"api_base": "https://example-a.test",
},
# Intentionally no model_info.id — this is the misconfiguration
# the warnings field is meant to flag.
"model_info": {},
},
]
cached_results = {
"healthy_endpoints": [
{
"model": "openai/gpt-4o",
"model_id": "id-a",
"api_base": "https://example-a.test",
},
],
"unhealthy_endpoints": [],
"healthy_count": 1,
"unhealthy_count": 0,
}
user_api_key_dict = UserAPIKeyAuth(
api_key="hashed-user-key",
models=["model-a"],
)
with (
patch("litellm.proxy.proxy_server.llm_model_list", full_model_list),
patch("litellm.proxy.proxy_server.llm_router", None),
patch("litellm.proxy.proxy_server.prisma_client", None),
patch("litellm.proxy.proxy_server.use_background_health_checks", True),
patch("litellm.proxy.proxy_server.user_model", None),
patch("litellm.proxy.proxy_server.health_check_results", cached_results),
patch("litellm.proxy.proxy_server.health_check_details", True),
patch("litellm.proxy.proxy_server.health_check_concurrency", 1),
):
result = await health_endpoint(
response=Response(), user_api_key_dict=user_api_key_dict
)
assert result["healthy_count"] == 0
assert result["unhealthy_count"] == 0
assert "warnings" in result, (
"empty cache result must surface a warnings field so the caller "
"can distinguish 'no deployments' from 'deployments excluded'"
)
assert any("model_info.id" in w for w in result["warnings"])
def test_clean_endpoint_data_strips_credentials_keeps_routing_fields():
"""
_clean_endpoint_data() drops credentials but leaves api_base /
api_version intact the per-caller hide/show happens in the endpoint
layer based on user role, not in the cleaning helper. This guarantees
proxy admins continue to see those fields in the /health response.
"""
from litellm.proxy.health_check import _clean_endpoint_data
raw = {
"model": "openai/gpt-4o",
"api_key": "sk-test",
"api_base": "https://example.test/v1",
"api_version": "2024-10-21",
"aws_access_key_id": "AKIAEXAMPLE",
}
cleaned = _clean_endpoint_data(raw, details=True)
assert "api_key" not in cleaned
assert "aws_access_key_id" not in cleaned
assert cleaned.get("api_base") == "https://example.test/v1"
assert cleaned.get("api_version") == "2024-10-21"

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

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

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.

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"

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

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

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 == {}

View file

@ -6,8 +6,8 @@ import {
deriveErrorMessage,
handleError,
} from "@/components/networking";
import { all_admin_roles } from "@/utils/roles";
import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized";
import { all_admin_roles } from "@/utils/roles";
// ── Types ────────────────────────────────────────────────────────────────────
@ -81,7 +81,6 @@ export const useProjects = () => {
return useQuery<ProjectResponse[]>({
queryKey: projectKeys.list({}),
queryFn: async () => fetchProjects(accessToken!),
enabled:
Boolean(accessToken) && all_admin_roles.includes(userRole || ""),
enabled: Boolean(accessToken) && all_admin_roles.includes(userRole!),
});
};

View file

@ -169,8 +169,8 @@ const UsagePage: React.FC<UsagePageProps> = ({ teams, organizations }) => {
}
}, [isAdmin, userID]);
// For non-admins, always pass their own user_id
const effectiveUserId = isAdmin ? selectedUserId : userID || null;
// For non-admins or "my-usage" view, always pass their own user_id
const effectiveUserId = usageView === "my-usage" || !isAdmin ? userID || null : selectedUserId;
const startTime = useMemo(() => (dateValue.from ? new Date(dateValue.from) : null), [dateValue.from]);
const endTime = useMemo(() => (dateValue.to ? new Date(dateValue.to) : null), [dateValue.to]);
@ -477,10 +477,10 @@ const UsagePage: React.FC<UsagePageProps> = ({ teams, organizations }) => {
}
/>
)}
{/* Your Usage Panel */}
{usageView === "global" && (
{/* Your Usage / Global Usage Panel */}
{(usageView === "global" || usageView === "my-usage") && (
<>
{isAdmin && (
{isAdmin && usageView === "global" && (
<div className="mb-4">
<Text className="mb-2">Filter by user</Text>
<Select

View file

@ -11,7 +11,7 @@ import {
} from "@ant-design/icons";
import { Badge, Select } from "antd";
import React from "react";
export type UsageOption = "global" | "organization" | "team" | "customer" | "tag" | "agent" | "user" | "user-agent-activity";
export type UsageOption = "global" | "my-usage" | "organization" | "team" | "customer" | "tag" | "agent" | "user" | "user-agent-activity";
export interface UsageViewSelectProps {
value: UsageOption;
onChange: (value: UsageOption) => void;
@ -43,6 +43,13 @@ const OPTIONS: OptionConfig[] = [
descriptionForNonAdmin: "View your usage",
icon: <GlobalOutlined style={{ fontSize: "16px" }} />,
},
{
value: "my-usage",
label: "Your Usage",
description: "View your own usage",
icon: <UserOutlined style={{ fontSize: "16px" }} />,
adminOnly: true,
},
{
value: "organization",
label: "Organization Usage",