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

This commit is contained in:
mateo-berri 2026-05-19 01:46:31 +00:00
commit ec767b7e4c
104 changed files with 6530 additions and 354 deletions

View file

@ -292,7 +292,7 @@ curl -X POST 'http://0.0.0.0:4000/v1/chat/completions' \
| [CompactifAI (`compactifai`)](https://docs.litellm.ai/docs/providers/compactifai) | ✅ | ✅ | ✅ | | | | | | | |
| [Custom (`custom`)](https://docs.litellm.ai/docs/providers/custom_llm_server) | ✅ | ✅ | ✅ | | | | | | | |
| [Custom OpenAI (`custom_openai`)](https://docs.litellm.ai/docs/providers/openai_compatible) | ✅ | ✅ | ✅ | | | ✅ | ✅ | ✅ | ✅ | |
| [Dashscope (`dashscope`)](https://docs.litellm.ai/docs/providers/dashscope) | ✅ | ✅ | ✅ | | | | | | | |
| [Dashscope (`dashscope`)](https://docs.litellm.ai/docs/providers/dashscope) | ✅ | ✅ | ✅ | | | | | | | |
| [Databricks (`databricks`)](https://docs.litellm.ai/docs/providers/databricks) | ✅ | ✅ | ✅ | | | | | | | |
| [DataRobot (`datarobot`)](https://docs.litellm.ai/docs/providers/datarobot) | ✅ | ✅ | ✅ | | | | | | | |
| [Deepgram (`deepgram`)](https://docs.litellm.ai/docs/providers/deepgram) | ✅ | ✅ | ✅ | | | ✅ | | | | |

View file

@ -12,6 +12,10 @@ spec:
name: {{ include "litellm.fullname" . }}
minReplicas: {{ .Values.autoscaling.minReplicas }}
maxReplicas: {{ .Values.autoscaling.maxReplicas }}
{{- if .Values.autoscaling.behavior }}
behavior:
{{- toYaml .Values.autoscaling.behavior | nindent 4 }}
{{- end }}
metrics:
{{- if .Values.autoscaling.targetCPUUtilizationPercentage }}
- type: Resource

View file

@ -0,0 +1,36 @@
suite: "hpa with behavior"
templates:
- hpa.yaml
tests:
- it: "renders behavior when set"
set:
autoscaling.enabled: true
autoscaling.behavior:
scaleUp:
stabilizationWindowSeconds: 60
policies:
- type: Pods
value: 2
periodSeconds: 60
scaleDown:
stabilizationWindowSeconds: 90
policies:
- type: Pods
value: 1
periodSeconds: 60
asserts:
- isKind: { of: HorizontalPodAutoscaler }
- equal: { path: spec.behavior.scaleUp.stabilizationWindowSeconds, value: 60 }
- equal: { path: spec.behavior.scaleDown.stabilizationWindowSeconds, value: 90 }
---
suite: "hpa without behavior"
templates:
- hpa.yaml
tests:
- it: "does not render behavior when not set"
set:
autoscaling.enabled: true
asserts:
- isKind: { of: HorizontalPodAutoscaler }
- isNull: { path: spec.behavior }

View file

@ -184,6 +184,7 @@ autoscaling:
maxReplicas: 100
targetCPUUtilizationPercentage: 80
# targetMemoryUtilizationPercentage: 80
# behavior: {}
# Autoscaling with keda is mutually exclusive with hpa
keda:

View file

@ -1,6 +1,6 @@
[project]
name = "litellm-enterprise"
version = "0.1.40"
version = "0.1.41"
description = "Package for LiteLLM Enterprise features"
readme = "README.md"
requires-python = ">=3.9"
@ -26,7 +26,7 @@ required-version = ">=0.10.9"
module-root = ""
[tool.commitizen]
version = "0.1.40"
version = "0.1.41"
version_files = [
"pyproject.toml:^version",
"../pyproject.toml:litellm-enterprise==",

View file

@ -0,0 +1,4 @@
-- AlterTable
-- Adds the admin-toggleable pause flag used by the router's blocked filter and the
-- credential lookup helpers; defaults to false so existing rows behave unchanged.
ALTER TABLE "LiteLLM_ProxyModelTable" ADD COLUMN IF NOT EXISTS "blocked" BOOLEAN NOT NULL DEFAULT false;

View file

@ -48,9 +48,10 @@ model LiteLLM_CredentialsTable {
// Models on proxy
model LiteLLM_ProxyModelTable {
model_id String @id @default(uuid())
model_name String
model_name String
litellm_params Json
model_info Json?
model_info Json?
blocked Boolean @default(false)
created_at DateTime @default(now()) @map("created_at")
created_by String
updated_at DateTime @default(now()) @updatedAt @map("updated_at")

View file

@ -416,6 +416,7 @@ custom_prometheus_metadata_labels: List[str] = []
custom_prometheus_tags: List[str] = []
prometheus_metrics_config: Optional[List] = None
prometheus_emit_stream_label: bool = False
prometheus_user_budget_label_include_email_alias: bool = False
prometheus_end_user_metrics_max_series_per_metric: Optional[int] = 10000
prometheus_end_user_metrics_ttl_seconds: Optional[float] = 3600.0
prometheus_end_user_metrics_cleanup_interval_seconds: Optional[float] = 60.0
@ -1880,6 +1881,12 @@ if TYPE_CHECKING:
from .llms.dashscope.chat.transformation import (
DashScopeChatConfig as DashScopeChatConfig,
)
from .llms.dashscope.embed.transformation import (
DashScopeEmbeddingConfig as DashScopeEmbeddingConfig,
)
from .llms.dashscope.rerank.transformation import (
DashScopeRerankConfig as DashScopeRerankConfig,
)
from .llms.moonshot.chat.transformation import (
MoonshotChatConfig as MoonshotChatConfig,
)

View file

@ -100,6 +100,8 @@ def _get_redis_cluster_kwargs(client=None):
"azure_tenant_id",
"azure_client_secret",
"max_connections",
"socket_timeout",
"socket_connect_timeout",
}
return available_args

View file

@ -87,6 +87,16 @@ class CachingHandlerResponse(BaseModel):
in_memory_cache_obj = InMemoryCache()
def _is_chat_completion_cached_dict(cached_result: dict) -> bool:
cached_id = cached_result.get("id")
if isinstance(cached_id, str) and cached_id.startswith("chatcmpl"):
return True
obj = cached_result.get("object")
if isinstance(obj, str):
return obj.startswith("chat.completion")
return "choices" in cached_result
def _should_defer_streaming_cache_hit_callbacks(*, kwargs: Dict[str, Any]) -> bool:
"""
When stream=True, do not run success callbacks at cache-hit time.
@ -861,27 +871,47 @@ class LLMCachingHandler:
elif (call_type == "aresponses" or call_type == "responses") and isinstance(
cached_result, dict
):
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,
)
use_chat_completion_cache = _is_chat_completion_cached_dict(cached_result)
if use_chat_completion_cache:
if kwargs.get("stream", False) is True:
bridge_call_type = (
CallTypes.acompletion.value
if call_type == "aresponses"
else CallTypes.completion.value
)
cached_result = self._convert_cached_stream_response(
cached_result=cached_result,
call_type=bridge_call_type,
logging_obj=logging_obj,
model=model,
)
else:
cached_result = convert_to_model_response_object(
response_object=cached_result,
model_response_object=ModelResponse(),
)
else:
cached_result = response_obj
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

@ -37,6 +37,15 @@ class ResponsesToCompletionBridgeHandler:
stream = litellm_params.get("stream", False)
return bool(stream)
@staticmethod
def _is_preformatted_cached_chat_stream(result: Any) -> bool:
from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
return (
isinstance(result, CustomStreamWrapper)
and result.custom_llm_provider == "cached_response"
)
@staticmethod
def _coerce_response_object(
response_obj: Any,
@ -177,6 +186,8 @@ class ResponsesToCompletionBridgeHandler:
**request_data,
)
from litellm.types.utils import ModelResponse
stream = self._resolve_stream_flag(optional_params, litellm_params)
if isinstance(result, ResponsesAPIResponse):
return self.transformation_handler.transform_response(
@ -192,6 +203,8 @@ class ResponsesToCompletionBridgeHandler:
api_key=kwargs.get("api_key"),
json_mode=kwargs.get("json_mode"),
)
elif isinstance(result, ModelResponse):
return result
elif not stream:
responses_api_response = self._collect_response_from_stream(result)
return self.transformation_handler.transform_response(
@ -208,6 +221,10 @@ class ResponsesToCompletionBridgeHandler:
json_mode=kwargs.get("json_mode"),
)
else:
if self._is_preformatted_cached_chat_stream(result):
return self._apply_post_stream_processing(
result, model, custom_llm_provider
)
completion_stream = self.transformation_handler.get_model_response_iterator(
streaming_response=result, # type: ignore
sync_stream=True,
@ -256,6 +273,8 @@ class ResponsesToCompletionBridgeHandler:
aresponses=True,
)
from litellm.types.utils import ModelResponse
stream = self._resolve_stream_flag(optional_params, litellm_params)
if isinstance(result, ResponsesAPIResponse):
return self.transformation_handler.transform_response(
@ -271,6 +290,8 @@ class ResponsesToCompletionBridgeHandler:
api_key=kwargs.get("api_key"),
json_mode=kwargs.get("json_mode"),
)
elif isinstance(result, ModelResponse):
return result
elif not stream:
responses_api_response = await self._collect_response_from_stream_async(
result
@ -289,6 +310,10 @@ class ResponsesToCompletionBridgeHandler:
json_mode=kwargs.get("json_mode"),
)
else:
if self._is_preformatted_cached_chat_stream(result):
return self._apply_post_stream_processing(
result, model, custom_llm_provider
)
completion_stream = self.transformation_handler.get_model_response_iterator(
streaming_response=result, # type: ignore
sync_stream=False,

View file

@ -1141,6 +1141,14 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator):
event_type = parsed_chunk.get("type")
if isinstance(event_type, ResponsesAPIStreamEvents):
event_type = event_type.value
if parsed_chunk.get("object") == "chat.completion.chunk" or (
event_type is None
and isinstance(parsed_chunk.get("choices"), list)
and parsed_chunk.get("choices")
):
return ModelResponseStream(**parsed_chunk)
verbose_logger.debug(f"Chat provider: Processing event type: {event_type}")
if event_type == "response.created":

View file

@ -173,17 +173,45 @@ def _cost_per_token_custom_pricing_helper(
prompt_tokens: float = 0,
completion_tokens: float = 0,
response_time_ms: Optional[float] = 0.0,
cached_tokens: float = 0,
cache_creation_tokens: float = 0,
### CUSTOM PRICING ###
custom_cost_per_token: Optional[CostPerToken] = None,
custom_cost_per_second: Optional[float] = None,
) -> Optional[Tuple[float, float]]:
"""Internal helper function for calculating cost, if custom pricing given"""
"""Internal helper function for calculating cost, if custom pricing given.
prompt_tokens is assumed to include both cached_tokens and cache_creation_tokens
(OpenAI-compatible convention). Anthropic-style usage where prompt_tokens excludes
cache tokens is handled at the caller (cost_per_token) before invoking this helper.
"""
if custom_cost_per_token is None and custom_cost_per_second is None:
return None
if custom_cost_per_token is not None:
input_cost = custom_cost_per_token["input_cost_per_token"] * prompt_tokens
output_cost = custom_cost_per_token["output_cost_per_token"] * completion_tokens
input_cost_per_token = custom_cost_per_token["input_cost_per_token"]
output_cost_per_token = custom_cost_per_token["output_cost_per_token"]
cache_read_input_token_cost = custom_cost_per_token.get(
"cache_read_input_token_cost",
input_cost_per_token,
)
cache_creation_input_token_cost = custom_cost_per_token.get(
"cache_creation_input_token_cost",
input_cost_per_token,
)
regular_prompt_tokens = max(
prompt_tokens - cached_tokens - cache_creation_tokens,
0,
)
input_cost = (
regular_prompt_tokens * input_cost_per_token
+ cached_tokens * cache_read_input_token_cost
+ cache_creation_tokens * cache_creation_input_token_cost
)
output_cost = completion_tokens * output_cost_per_token
return input_cost, output_cost
elif custom_cost_per_second is not None:
output_cost = custom_cost_per_second * response_time_ms / 1000 # type: ignore
@ -323,10 +351,56 @@ def cost_per_token( # noqa: PLR0915
)
## CUSTOM PRICING ##
# Normalize cache token counts across providers:
# - OpenAI-compatible: usage.prompt_tokens_details.cached_tokens
# (prompt_tokens already INCLUDES cached_tokens)
# - Anthropic: usage.cache_read_input_tokens / cache_creation_input_tokens
# (prompt_tokens does NOT include these — adjust before calling helper)
_cache_read_tokens: float = 0
_cache_creation_tokens: float = 0
_is_anthropic_style = False
if usage_object is not None:
_pt_details = getattr(usage_object, "prompt_tokens_details", None)
if _pt_details is not None:
_cache_read_tokens = float(getattr(_pt_details, "cached_tokens", 0) or 0)
# OpenAI-compatible providers report cache-write tokens under
# either `cache_write_tokens` (kimi-k2) or `cache_creation_tokens`.
# Mirror db_spend_update_writer to stay symmetric.
_cache_creation_tokens = float(
getattr(_pt_details, "cache_write_tokens", 0)
or getattr(_pt_details, "cache_creation_tokens", 0)
or 0
)
_anthropic_read = getattr(usage_object, "cache_read_input_tokens", None)
_anthropic_create = getattr(usage_object, "cache_creation_input_tokens", None)
if _anthropic_read is not None or _anthropic_create is not None:
_is_anthropic_style = True
if _anthropic_read is not None:
_cache_read_tokens = float(_anthropic_read)
if _anthropic_create is not None:
_cache_creation_tokens = float(_anthropic_create)
if not _cache_read_tokens and cache_read_input_tokens:
_cache_read_tokens = float(cache_read_input_tokens)
_is_anthropic_style = True
if not _cache_creation_tokens and cache_creation_input_tokens:
_cache_creation_tokens = float(cache_creation_input_tokens)
_is_anthropic_style = True
# Anthropic reports prompt_tokens as input_tokens (excluding cache tokens).
# Adjust so the helper's "prompt_tokens includes cache tokens" invariant holds.
_normalized_prompt_tokens = float(prompt_tokens)
if _is_anthropic_style:
_normalized_prompt_tokens += _cache_read_tokens + _cache_creation_tokens
response_cost = _cost_per_token_custom_pricing_helper(
prompt_tokens=prompt_tokens,
prompt_tokens=_normalized_prompt_tokens,
completion_tokens=completion_tokens,
response_time_ms=response_time_ms,
cached_tokens=_cache_read_tokens,
cache_creation_tokens=_cache_creation_tokens,
custom_cost_per_second=custom_cost_per_second,
custom_cost_per_token=custom_cost_per_token,
)

View file

@ -918,9 +918,11 @@ class GuardrailRaisedException(Exception):
guardrail_name: Optional[str] = None,
message: str = "",
should_wrap_with_default_message: bool = True,
status_code: int = 400,
):
default_message = f"Guardrail raised an exception, Guardrail: {guardrail_name}, Message: {message}"
self.guardrail_name = guardrail_name
self.status_code = status_code
self.message = default_message if should_wrap_with_default_message else message
super().__init__(self.message)
@ -930,12 +932,14 @@ class BlockedPiiEntityError(Exception):
self,
entity_type: str,
guardrail_name: Optional[str] = None,
status_code: int = 400,
):
"""
Raised when a blocked entity is detected by a guardrail.
"""
self.entity_type = entity_type
self.guardrail_name = guardrail_name
self.status_code = status_code
self.message = f"Blocked entity detected: {entity_type} by Guardrail: {guardrail_name}. This entity is not allowed to be used in this request."
super().__init__(self.message)

View file

@ -43,7 +43,11 @@ if TYPE_CHECKING:
dc = DualCache()
from litellm.exceptions import ModifyResponseException as ModifyResponseException
from litellm.exceptions import (
BlockedPiiEntityError,
GuardrailRaisedException,
ModifyResponseException,
)
class CustomGuardrail(CustomLogger):
@ -737,12 +741,15 @@ class CustomGuardrail(CustomLogger):
(this was logged previously as an API failure - guardrail_failed_to_respond).
Guardrails signal intentional blocks by raising:
- GuardrailRaisedException (generic guardrail API, tool permission)
- BlockedPiiEntityError (Presidio PII detection)
- HTTPException with status 400 (content policy violation)
- ModifyResponseException (passthrough mode violation)
"""
if isinstance(e, ModifyResponseException):
return True
if isinstance(e, (GuardrailRaisedException, BlockedPiiEntityError)):
return True
if (
HTTPException is not None
and isinstance(e, HTTPException)

View file

@ -1107,8 +1107,13 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
"mcp_tool_call_metadata",
"vector_store_request_metadata",
]:
if md.get(key) is not None:
common_attrs[f"metadata.{key}"] = str(md[key])
value = md.get(key)
if value is None:
continue
if isinstance(value, (dict, list)):
common_attrs[f"metadata.{key}"] = safe_dumps(value)
else:
common_attrs[f"metadata.{key}"] = str(value)
# get hidden params
hidden_params = getattr(std_log, "hidden_params", None) or (std_log or {}).get(

View file

@ -3540,6 +3540,10 @@ class PrometheusLogger(CustomLogger):
user_object.budget_reset_at = user_info.budget_reset_at
if user_object.max_budget is None and user_info.max_budget is not None:
user_object.max_budget = user_info.max_budget
if user_info.user_email is not None:
user_object.user_email = user_info.user_email
if user_info.user_alias is not None:
user_object.user_alias = user_info.user_alias
return user_object
@ -3556,6 +3560,8 @@ class PrometheusLogger(CustomLogger):
"""
enum_values = UserAPIKeyLabelValues(
user=user.user_id,
user_email=user.user_email or "",
user_alias=user.user_alias or "",
)
_labels = prometheus_label_factory(

View file

@ -20,6 +20,7 @@ from typing import (
cast,
)
import litellm
from litellm import verbose_logger
from litellm.router_utils.batch_utils import InMemoryFile
from litellm.types.llms.openai import (
@ -1170,9 +1171,16 @@ def migrate_file_to_image_url(
ChatCompletionImageUrlObject,
)
file_id = message["file"].get("file_id")
file_data = message["file"].get("file_data")
format = message["file"].get("format")
file_sub = message.get("file")
if file_sub is None:
raise litellm.BadRequestError(
message="Content block has type='file' but is missing the required 'file' field",
model=None,
llm_provider=None,
)
file_id = file_sub.get("file_id")
file_data = file_sub.get("file_data")
format = file_sub.get("format")
if not file_id and not file_data:
raise ValueError("file_id and file_data are both None")
image_url_object = ChatCompletionImageObject(

View file

@ -2057,9 +2057,16 @@ def anthropic_process_openai_file_message(
AnthropicMessagesContainerUploadParam,
]:
file_message = cast(ChatCompletionFileObject, message)
file_data = file_message["file"].get("file_data")
file_id = file_message["file"].get("file_id")
format = file_message["file"].get("format")
file_sub = file_message.get("file")
if file_sub is None:
raise litellm.BadRequestError(
message="Content block has type='file' but is missing the required 'file' field",
model=None,
llm_provider="anthropic",
)
file_data = file_sub.get("file_data")
file_id = file_sub.get("file_id")
format = file_sub.get("format")
if file_data:
image_chunk = convert_to_anthropic_image_obj(
openai_image_url=file_data,
@ -4879,7 +4886,13 @@ class BedrockConverseMessagesProcessor:
@staticmethod
def _process_file_message(message: ChatCompletionFileObject) -> BedrockContentBlock:
file_message = message["file"]
file_message = message.get("file")
if file_message is None:
raise litellm.BadRequestError(
message="Content block has type='file' but is missing the required 'file' field",
model=None,
llm_provider="bedrock",
)
file_data = file_message.get("file_data")
file_id = file_message.get("file_id")
@ -4900,7 +4913,13 @@ class BedrockConverseMessagesProcessor:
async def _async_process_file_message(
message: ChatCompletionFileObject,
) -> BedrockContentBlock:
file_message = message["file"]
file_message = message.get("file")
if file_message is None:
raise litellm.BadRequestError(
message="Content block has type='file' but is missing the required 'file' field",
model=None,
llm_provider="bedrock",
)
file_data = file_message.get("file_data")
file_id = file_message.get("file_id")
format = file_message.get("format")

View file

@ -5,6 +5,7 @@ from typing import Any, Dict, List, Literal, Optional, Union, cast
from httpx import Headers, Response
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
from litellm.llms.base_llm.batches.transformation import BaseBatchesConfig
from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.secret_managers.main import get_secret_str
@ -263,9 +264,32 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig):
cancelling_at=None,
cancelled_at=None,
request_counts=None,
metadata=original_request.get("metadata", {}),
metadata=self._get_openai_compatible_batch_metadata(
original_request.get("metadata", {})
),
)
@staticmethod
def _get_openai_compatible_batch_metadata(metadata: Any) -> Dict[str, str]:
"""
OpenAI Batch metadata only accepts string values.
"""
if not isinstance(metadata, dict):
return {}
sanitized_metadata: Dict[str, str] = {}
for key, value in metadata.items():
if key == "standard_logging_guardrail_information" or value is None:
continue
str_key = str(key)
if isinstance(value, str):
sanitized_metadata[str_key] = value
else:
sanitized_metadata[str_key] = safe_dumps(value)
return sanitized_metadata
def transform_retrieve_batch_request(
self,
batch_id: str,

View file

@ -22,7 +22,7 @@ class BedrockCohereEmbeddingConfig:
) -> dict:
for k, v in non_default_params.items():
if k == "encoding_format":
optional_params["embedding_types"] = v
optional_params["embedding_types"] = v if isinstance(v, list) else [v]
elif k == "dimensions":
optional_params["output_dimension"] = v
return optional_params

View file

@ -0,0 +1,28 @@
"""
Common utilities for the DashScope LLM provider.
"""
from typing import Optional
import httpx
from litellm.llms.base_llm.chat.transformation import BaseLLMException
class DashScopeError(BaseLLMException):
"""Exception class for DashScope provider errors."""
def __init__(
self,
status_code: int,
message: str,
headers: Optional[httpx.Headers] = None,
):
self.status_code = status_code
self.message = message
self.headers = headers or httpx.Headers()
super().__init__(
status_code=status_code,
message=message,
headers=dict(self.headers),
)

View file

@ -0,0 +1,7 @@
"""
DashScope Embedding Module
"""
from .transformation import DashScopeEmbeddingConfig
__all__ = ["DashScopeEmbeddingConfig"]

View file

@ -0,0 +1,191 @@
"""
Transformation logic from OpenAI /v1/embeddings format to DashScope's /v1/embeddings format.
Supports
- text-embedding-v4
- text-embedding-v3
Endpoint
- https://dashscope.aliyuncs.com/compatible-mode/v1/embeddings
Docs - https://help.aliyun.com/zh/model-studio/text-embedding-synchronous-api
"""
from typing import List, Optional, Union
import httpx
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.llms.base_llm.embedding.transformation import BaseEmbeddingConfig
from litellm.secret_managers.main import get_secret_str
from litellm.types.llms.openai import AllEmbeddingInputValues, AllMessageValues
from litellm.types.utils import EmbeddingResponse, Usage
from ..common_utils import DashScopeError
DEFAULT_API_BASE = "https://dashscope.aliyuncs.com/compatible-mode/v1"
class DashScopeEmbeddingConfig(BaseEmbeddingConfig):
"""
Reference: https://help.aliyun.com/zh/model-studio/text-embedding-synchronous-api
DashScope exposes an OpenAI-compatible /v1/embeddings endpoint, so the
request and response shapes are nearly identical to OpenAI's.
"""
def __init__(self) -> None:
pass
def get_supported_openai_params(self, model: str) -> List[str]:
# DashScope's compatible-mode embeddings API accepts the same params as OpenAI.
# `dimensions` / `encoding_format` are only honored by text-embedding-v3 / v4;
# earlier versions silently ignore them server-side.
return ["dimensions", "encoding_format", "user"]
def map_openai_params(
self,
non_default_params: dict,
optional_params: dict,
model: str,
drop_params: bool = False,
) -> dict:
supported = self.get_supported_openai_params(model)
for k, v in non_default_params.items():
if v is None:
continue
if k in supported:
optional_params[k] = v
# unsupported params are dropped when drop_params=True;
# the upstream _check_valid_arg already raised UnsupportedParamsError
# for drop_params=False before this method is called.
return optional_params
def validate_environment(
self,
headers: dict,
model: str,
messages: List[AllMessageValues],
optional_params: dict,
litellm_params: dict,
api_key: Optional[str] = None,
api_base: Optional[str] = None,
) -> dict:
if api_key is None:
api_key = get_secret_str("DASHSCOPE_API_KEY")
if api_key is None:
raise ValueError(
"DashScope API key is required. Set 'DASHSCOPE_API_KEY' env var or pass api_key explicitly."
)
default_headers = {
"Content-Type": "application/json",
"Authorization": f"Bearer {api_key}",
}
return {**default_headers, **headers}
def get_complete_url(
self,
api_base: Optional[str],
api_key: Optional[str],
model: str,
optional_params: dict,
litellm_params: dict,
stream: Optional[bool] = None,
) -> str:
base = api_base or get_secret_str("DASHSCOPE_API_BASE") or DEFAULT_API_BASE
base = base.rstrip("/")
if base.endswith("/embeddings"):
return base
return f"{base}/embeddings"
def transform_embedding_request(
self,
model: str,
input: AllEmbeddingInputValues,
optional_params: dict,
headers: dict,
) -> dict:
data: dict = {
"model": model,
"input": input,
}
for key in ("dimensions", "encoding_format", "user"):
value = optional_params.get(key)
if value is not None:
data[key] = value
return data
def transform_embedding_response(
self,
model: str,
raw_response: httpx.Response,
model_response: EmbeddingResponse,
logging_obj: LiteLLMLoggingObj,
api_key: Optional[str],
request_data: dict,
optional_params: dict,
litellm_params: dict,
) -> EmbeddingResponse:
try:
response_json = raw_response.json()
except Exception as e:
raise DashScopeError(
status_code=raw_response.status_code,
message=f"Failed to parse DashScope response as JSON: {str(e)}",
)
logging_obj.post_call(
input=request_data.get("input"),
api_key=api_key,
additional_args={"complete_input_dict": request_data},
original_response=response_json,
)
if "error" in response_json:
error = response_json["error"]
message = (
error.get("message", str(error))
if isinstance(error, dict)
else str(error)
)
raise DashScopeError(
status_code=raw_response.status_code,
message=message,
)
model_response.object = "list"
model_response.data = response_json.get("data", [])
model_response.model = response_json.get("model", model)
usage = response_json.get("usage") or {}
prompt_tokens = usage.get("prompt_tokens", 0)
total_tokens = usage.get("total_tokens", prompt_tokens)
setattr(
model_response,
"usage",
Usage(
prompt_tokens=prompt_tokens,
completion_tokens=0,
total_tokens=total_tokens,
),
)
if "id" in response_json:
setattr(model_response, "id", response_json["id"])
return model_response
def get_error_class(
self,
error_message: str,
status_code: int,
headers: Union[dict, httpx.Headers],
) -> BaseLLMException:
if isinstance(headers, dict):
headers = httpx.Headers(headers)
return DashScopeError(
status_code=status_code,
message=error_message,
headers=headers,
)

View file

@ -0,0 +1,7 @@
"""
DashScope Rerank Module
"""
from .transformation import DashScopeRerankConfig
__all__ = ["DashScopeRerankConfig"]

View file

@ -0,0 +1,241 @@
"""
Transformation logic for DashScope's OpenAI-compatible /v1/reranks API.
Supports
- qwen3-rerank
(Other DashScope rerankers gte-rerank-v2 / qwen3-vl-rerank share the same
endpoint but have not been validated against this transformer. Behavior with
those models is undefined.)
Endpoint
- https://dashscope.aliyuncs.com/compatible-api/v1/reranks
Note: chat/embed live under `/compatible-mode/v1/`, but DashScope's rerank
route is exposed under `/compatible-api/v1/reranks` per the docs. Override
with `DASHSCOPE_API_BASE_RERANK` to point at a different host or path.
Empirically, qwen3-rerank accepts `return_documents=true` and echoes
`results[].document.text` back, even though the public docs list the flag
as supported only for gte-rerank-v2 / qwen3-vl-rerank.
Docs - https://help.aliyun.com/zh/model-studio/text-rerank-api
"""
from typing import Any, Dict, List, Optional, Union
import httpx
from litellm._uuid import uuid
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.llms.base_llm.rerank.transformation import BaseRerankConfig
from litellm.secret_managers.main import get_secret_str
from litellm.types.rerank import (
OptionalRerankParams,
RerankBilledUnits,
RerankResponse,
RerankResponseMeta,
RerankTokens,
)
from ..common_utils import DashScopeError
DEFAULT_RERANK_URL = "https://dashscope.aliyuncs.com/compatible-api/v1/reranks"
class DashScopeRerankConfig(BaseRerankConfig):
"""
Reference: https://help.aliyun.com/zh/model-studio/text-rerank-api
Targets DashScope's qwen3-rerank model. Request fields: model, query,
documents, top_n, return_documents. Response: results[].index,
results[].relevance_score, optionally results[].document.text (when
return_documents=true), plus a top-level usage.total_tokens counter.
"""
def __init__(self) -> None:
pass
def get_complete_url(
self,
api_base: Optional[str],
model: str,
optional_params: Optional[dict] = None,
) -> str:
if api_base is None:
api_base = get_secret_str("DASHSCOPE_API_BASE_RERANK") or DEFAULT_RERANK_URL
if api_base == DEFAULT_RERANK_URL:
return DEFAULT_RERANK_URL
cleaned = api_base.rstrip("/")
if cleaned.endswith("/reranks") or cleaned.endswith("/rerank"):
return cleaned
if cleaned.endswith("/v1"):
return f"{cleaned}/reranks"
# Unknown base: append /reranks rather than silently ignoring the caller's api_base.
return f"{cleaned}/reranks"
def validate_environment(
self,
headers: dict,
model: str,
api_key: Optional[str] = None,
optional_params: Optional[dict] = None,
) -> dict:
if api_key is None:
api_key = get_secret_str("DASHSCOPE_API_KEY")
if api_key is None:
raise ValueError(
"DashScope API key is required. Set 'DASHSCOPE_API_KEY' env var or pass api_key explicitly."
)
default_headers = {
"Authorization": f"Bearer {api_key}",
"accept": "application/json",
"content-type": "application/json",
}
return {**default_headers, **headers}
def get_supported_cohere_rerank_params(self, model: str) -> list:
return ["query", "documents", "top_n", "return_documents"]
def map_cohere_rerank_params(
self,
non_default_params: Optional[dict],
model: str,
drop_params: bool,
query: str,
documents: List[Union[str, Dict[str, Any]]],
custom_llm_provider: Optional[str] = None,
top_n: Optional[int] = None,
rank_fields: Optional[List[str]] = None,
return_documents: Optional[bool] = True,
max_chunks_per_doc: Optional[int] = None,
max_tokens_per_doc: Optional[int] = None,
) -> Dict:
# qwen3-rerank accepts query/documents/top_n/return_documents. The
# rest (rank_fields, max_*_per_doc) are silently dropped.
params: OptionalRerankParams = OptionalRerankParams(
query=query,
documents=documents,
)
if top_n is not None:
params["top_n"] = top_n
if return_documents is not None:
params["return_documents"] = return_documents
return dict(params)
def transform_rerank_request(
self,
model: str,
optional_rerank_params: Dict,
headers: dict,
litellm_params: Optional[dict] = None,
) -> dict:
if "query" not in optional_rerank_params:
raise ValueError("query is required for DashScope rerank")
if "documents" not in optional_rerank_params:
raise ValueError("documents is required for DashScope rerank")
request: Dict[str, Any] = {
"model": model,
"query": optional_rerank_params["query"],
"documents": optional_rerank_params["documents"],
}
if optional_rerank_params.get("top_n") is not None:
request["top_n"] = optional_rerank_params["top_n"]
if optional_rerank_params.get("return_documents") is not None:
request["return_documents"] = optional_rerank_params["return_documents"]
return request
def transform_rerank_response(
self,
model: str,
raw_response: httpx.Response,
model_response: RerankResponse,
logging_obj: LiteLLMLoggingObj,
api_key: Optional[str] = None,
request_data: Optional[dict] = None,
optional_params: Optional[dict] = None,
litellm_params: Optional[dict] = None,
) -> RerankResponse:
request_data = request_data or {}
optional_params = optional_params or {}
litellm_params = litellm_params or {}
try:
response_json = raw_response.json()
except Exception:
raise DashScopeError(
status_code=raw_response.status_code,
message=raw_response.text,
)
logging_obj.post_call(
input=request_data.get("query"),
api_key=api_key,
additional_args={"complete_input_dict": request_data},
original_response=response_json,
)
# DashScope error envelope: {"code": "...", "message": "...", "request_id": "..."}
if "code" in response_json and "results" not in response_json:
raise DashScopeError(
status_code=raw_response.status_code,
message=response_json.get("message", str(response_json)),
)
results = response_json.get("results")
if results is None:
raise DashScopeError(
status_code=raw_response.status_code,
message=f"No results in DashScope rerank response: {response_json}",
)
# qwen3-rerank returns:
# {"index": int, "relevance_score": float}
# plus, when return_documents=true was sent:
# "document": {"text": "..."}
# which already matches LiteLLM's RerankResponseDocument shape.
transformed_results: List[dict] = []
for r in results:
item: Dict[str, Any] = {
"index": r["index"],
"relevance_score": r["relevance_score"],
}
doc = r.get("document")
if isinstance(doc, dict):
item["document"] = doc
elif isinstance(doc, str):
# Defensive: spec says dict, but normalize string-shaped echoes.
item["document"] = {"text": doc}
transformed_results.append(item)
usage = response_json.get("usage") or {}
total_tokens = usage.get("total_tokens")
billed_units = RerankBilledUnits(total_tokens=total_tokens)
tokens = RerankTokens(input_tokens=total_tokens)
meta = RerankResponseMeta(billed_units=billed_units, tokens=tokens)
return RerankResponse(
id=response_json.get("id") or str(uuid.uuid4()),
results=transformed_results, # type: ignore
meta=meta,
)
def get_error_class(
self,
error_message: str,
status_code: int,
headers: Union[dict, httpx.Headers],
) -> BaseLLMException:
if isinstance(headers, dict):
headers = httpx.Headers(headers)
return DashScopeError(
status_code=status_code,
message=error_message,
headers=headers,
)

View file

@ -2,13 +2,15 @@
Translates from OpenAI's `/v1/chat/completions` to DeepSeek's `/v1/chat/completions`
"""
from typing import Any, Coroutine, List, Literal, Optional, Tuple, Union, overload
from typing import Any, Coroutine, List, Literal, Optional, Tuple, Union, cast, overload
import litellm
from litellm.litellm_core_utils.prompt_templates.common_utils import (
handle_messages_with_content_list_to_str_conversion,
)
from litellm.secret_managers.main import get_secret_str
from litellm.types.llms.openai import AllMessageValues
from litellm.utils import supports_reasoning
from ...openai.chat.gpt_transformation import OpenAIGPTConfig
@ -62,6 +64,48 @@ class DeepSeekChatConfig(OpenAIGPTConfig):
return optional_params
def _fill_reasoning_content(
self, messages: List[AllMessageValues]
) -> List[AllMessageValues]:
"""
DeepSeek thinking mode requires `reasoning_content` to be passed back on
every assistant message in multi-turn conversations. If it is missing,
the API returns:
"The reasoning_content in the thinking mode must be passed back to the API."
For each assistant message that is missing `reasoning_content`:
1. Promote it from `provider_specific_fields["reasoning_content"]` if present
(LiteLLM stores provider-specific response fields there).
2. Otherwise inject a single space the minimum value the API accepts.
"""
result: List[AllMessageValues] = []
for msg in messages:
if msg.get("role") == "assistant" and not msg.get("reasoning_content"):
patched = dict(cast(dict, msg))
provider_fields = patched.get("provider_specific_fields") or {}
stored = provider_fields.get("reasoning_content")
if stored:
patched["reasoning_content"] = stored
cleaned = dict(provider_fields)
cleaned.pop("reasoning_content", None)
patched["provider_specific_fields"] = cleaned
else:
litellm.verbose_logger.warning(
"DeepSeek thinking mode: assistant message is missing "
"`reasoning_content` and none was saved in "
"`provider_specific_fields`. A single-space placeholder "
"is being injected to satisfy API validation, but the "
"model will receive a blank reasoning chain for this turn, "
"which may silently degrade multi-turn response quality. "
"Preserve `reasoning_content` from the original assistant "
"response when building multi-turn conversation history."
)
patched["reasoning_content"] = " "
result.append(cast(AllMessageValues, patched))
else:
result.append(msg)
return result
@overload
def _transform_messages(
self, messages: List[AllMessageValues], model: str, is_async: Literal[True]
@ -91,6 +135,66 @@ class DeepSeekChatConfig(OpenAIGPTConfig):
messages=messages, model=model, is_async=False
)
def _thinking_mode_active(self, model: str, optional_params: dict) -> bool:
"""
Returns True only when thinking mode is actually active for this request:
- model supports reasoning (capability check)
- user explicitly passed thinking={"type": "enabled"} (opt-in check)
"""
return (
supports_reasoning(model=model, custom_llm_provider="deepseek")
and (optional_params.get("thinking") or {}).get("type") == "enabled"
)
def transform_request(
self,
model: str,
messages: List[AllMessageValues],
optional_params: dict,
litellm_params: dict,
headers: dict,
) -> dict:
"""
Ensures `reasoning_content` is forwarded on assistant messages for
multi-turn thinking-mode conversations (issue #28045).
Only runs when thinking mode is actually active - guarded by both
supports_reasoning() (model capability) and optional_params["thinking"]
(user explicitly enabled it), preventing spurious injection on models
like deepseek-v3.2 that support thinking as opt-in but not always-on.
"""
if self._thinking_mode_active(model=model, optional_params=optional_params):
messages = self._fill_reasoning_content(messages)
return super().transform_request(
model=model,
messages=messages,
optional_params=optional_params,
litellm_params=litellm_params,
headers=headers,
)
async def async_transform_request(
self,
model: str,
messages: List[AllMessageValues],
optional_params: dict,
litellm_params: dict,
headers: dict,
) -> dict:
"""
Async equivalent of transform_request applies the same reasoning_content
fix for multi-turn thinking-mode conversations.
"""
if self._thinking_mode_active(model=model, optional_params=optional_params):
messages = self._fill_reasoning_content(messages)
return await super().async_transform_request(
model=model,
messages=messages,
optional_params=optional_params,
litellm_params=litellm_params,
headers=headers,
)
def _get_openai_compatible_provider_info(
self, api_base: Optional[str], api_key: Optional[str]
) -> Tuple[Optional[str], Optional[str]]:

View file

@ -0,0 +1,133 @@
"""
DeepSeek Anthropic-compatible messages transformation config.
"""
from typing import Any, Dict, List, Optional, Tuple
import litellm
from litellm.llms.anthropic.experimental_pass_through.messages.transformation import (
AnthropicMessagesConfig,
)
from litellm.secret_managers.main import get_secret_str
from litellm.types.router import GenericLiteLLMParams
class DeepSeekAnthropicMessagesConfig(AnthropicMessagesConfig):
"""
DeepSeek exposes an Anthropic-compatible Messages API at
https://api.deepseek.com/anthropic.
It accepts the native Anthropic Messages conversation shape, including
thinking blocks in assistant history, but rejects Anthropic's explicit
custom-tool discriminator (`{"type": "custom"}`).
"""
@property
def custom_llm_provider(self) -> Optional[str]:
return "deepseek"
@staticmethod
def get_api_key(api_key: Optional[str] = None) -> Optional[str]:
return api_key or get_secret_str("DEEPSEEK_API_KEY") or litellm.api_key
@staticmethod
def get_api_base(api_base: Optional[str] = None) -> str:
return (
api_base
or get_secret_str("DEEPSEEK_ANTHROPIC_API_BASE")
or get_secret_str("DEEPSEEK_API_BASE")
or "https://api.deepseek.com/anthropic"
)
def validate_anthropic_messages_environment(
self,
headers: dict,
model: str,
messages: List[Any],
optional_params: dict,
litellm_params: dict,
api_key: Optional[str] = None,
api_base: Optional[str] = None,
) -> Tuple[dict, Optional[str]]:
dynamic_api_key = self.get_api_key(api_key=api_key)
if (
"x-api-key" not in headers
and "authorization" not in headers
and dynamic_api_key is not None
):
headers["x-api-key"] = dynamic_api_key
if "anthropic-version" not in headers:
headers["anthropic-version"] = "2023-06-01"
if "content-type" not in headers:
headers["content-type"] = "application/json"
headers = self._update_headers_with_anthropic_beta(
headers=headers,
optional_params=optional_params,
custom_llm_provider=self.custom_llm_provider or "deepseek",
)
return headers, api_base
def get_complete_url(
self,
api_base: Optional[str],
api_key: Optional[str],
model: str,
optional_params: dict,
litellm_params: dict,
stream: Optional[bool] = None,
) -> str:
base_url = self.get_api_base(api_base=api_base).rstrip("/")
if base_url.endswith("/v1/messages") and "/anthropic/" in base_url:
return base_url
if base_url.endswith("/v1/messages"):
base_url = base_url[: -len("/v1/messages")]
if base_url.endswith("/v1"):
base_url = base_url[: -len("/v1")]
if base_url.endswith("/beta"):
base_url = base_url[: -len("/beta")]
if not base_url.endswith("/anthropic") and "/anthropic/" not in base_url:
base_url = f"{base_url}/anthropic"
return f"{base_url}/v1/messages"
@staticmethod
def _sanitize_tools_for_deepseek(tools: Any) -> Any:
if not isinstance(tools, list):
return tools
sanitized_tools = []
for tool in tools:
if isinstance(tool, dict) and tool.get("type") == "custom":
sanitized_tool = dict(tool)
sanitized_tool.pop("type", None)
sanitized_tools.append(sanitized_tool)
else:
sanitized_tools.append(tool)
return sanitized_tools
def transform_anthropic_messages_request(
self,
model: str,
messages: List[Dict],
anthropic_messages_optional_request_params: Dict,
litellm_params: GenericLiteLLMParams,
headers: dict,
) -> Dict:
anthropic_messages_request = super().transform_anthropic_messages_request(
model=model,
messages=messages,
anthropic_messages_optional_request_params=anthropic_messages_optional_request_params,
litellm_params=litellm_params,
headers=headers,
)
if "tools" in anthropic_messages_request:
anthropic_messages_request["tools"] = self._sanitize_tools_for_deepseek(
anthropic_messages_request["tools"]
)
return anthropic_messages_request

View file

@ -1,5 +1,7 @@
from typing import List, Optional, cast
import litellm
from litellm.litellm_core_utils.prompt_templates.factory import (
convert_generic_image_chunk_to_openai_image_obj,
convert_to_anthropic_image_obj,
@ -101,7 +103,10 @@ class GoogleAIStudioGeminiConfig(VertexGeminiConfig):
return supported_params
def _transform_messages(
self, messages: List[AllMessageValues], model: Optional[str] = None
self,
messages: List[AllMessageValues],
model: Optional[str] = None,
litellm_params: Optional[dict] = None,
) -> List[ContentType]:
"""
Google AI Studio Gemini does not support HTTP/HTTPS URLs for files.
@ -141,14 +146,23 @@ class GoogleAIStudioGeminiConfig(VertexGeminiConfig):
img_element["image_url"] = converted_image_url # type: ignore
elif element.get("type") == "file":
file_element = cast(ChatCompletionFileObject, element)
file_id = file_element["file"].get("file_id")
_file_field = file_element.get("file")
if _file_field is None:
raise litellm.BadRequestError(
message="Content block has type='file' but is missing the required 'file' field",
model=model,
llm_provider="gemini",
)
file_id = _file_field.get("file_id")
if file_id and ("http://" in file_id or "https://" in file_id):
# Convert HTTP/HTTPS file URL to base64 data
try:
base64_data = convert_url_to_base64(file_id)
file_element["file"]["file_data"] = base64_data # type: ignore
file_element["file"].pop("file_id", None) # type: ignore
_file_field["file_data"] = base64_data # type: ignore
_file_field.pop("file_id", None) # type: ignore
except Exception:
# If conversion fails, leave as is and let the API handle it
pass
return _gemini_convert_messages_with_history(messages=messages, model=model)
return _gemini_convert_messages_with_history(
messages=messages, model=model, litellm_params=litellm_params
)

View file

@ -287,7 +287,13 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig):
content_item["image_url"] = new_image_url_obj
elif content_item.get("type") == "file":
content_item = cast(ChatCompletionFileObject, content_item)
file_obj = content_item["file"]
file_obj = content_item.get("file")
if file_obj is None:
raise litellm.BadRequestError(
message="Content block has type='file' but is missing the required 'file' field",
model=None,
llm_provider="openai",
)
new_file_obj = ChatCompletionFileObjectFile(
**{ # type: ignore
k: v

View file

@ -6,13 +6,16 @@ Why separate file? Make it easy to see how transformation works
import json
import os
from typing import TYPE_CHECKING, Dict, List, Literal, Optional, Tuple, Union, cast
import re
from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Tuple, Union, cast
from urllib.parse import quote
import httpx
from pydantic import BaseModel
import litellm
from litellm._logging import verbose_logger
from litellm.litellm_core_utils.asyncify import asyncify
from litellm.litellm_core_utils.prompt_templates.common_utils import (
_get_image_mime_type_from_url,
)
@ -57,6 +60,45 @@ from ..common_utils import (
get_supports_system_message,
)
# Typed as Any to avoid introducing a module-load-time cyclic import to
# vertex_llm_base. The instance is lazily constructed by _get_vertex_base()
# the first time GCS metadata needs to be fetched.
_GCS_METADATA_VERTEX_BASE: Optional[Any] = None
# Shared sync client for GCS JSON API metadata reads so proxy/SSL settings
# from litellm's HTTP stack apply (see Greptile review on PR #27278).
_GCS_METADATA_HTTP_HANDLER: Optional[HTTPHandler] = None
_GEMINI_MIME_TYPE_ALIASES: Dict[str, str] = {
"image/jpg": "image/jpeg",
}
def _apply_gemini_mime_type_aliases(mime_type: str) -> str:
"""Normalize known MIME aliases only; does not consult the file-type registry.
Also strips MIME parameters (e.g. ``; charset=utf-8``) so that values
sourced from GCS object metadata (``contentType``) validate correctly.
"""
normalized = mime_type.split(";", 1)[0].strip().lower()
return _GEMINI_MIME_TYPE_ALIASES.get(normalized, normalized)
def _get_vertex_base() -> Any:
"""Lazily return the shared VertexBase instance to avoid a module-load-time cyclic import."""
global _GCS_METADATA_VERTEX_BASE
if _GCS_METADATA_VERTEX_BASE is None:
from ..vertex_llm_base import VertexBase
_GCS_METADATA_VERTEX_BASE = VertexBase()
return _GCS_METADATA_VERTEX_BASE
def _get_gcs_metadata_http_handler() -> HTTPHandler:
global _GCS_METADATA_HTTP_HANDLER
if _GCS_METADATA_HTTP_HANDLER is None:
_GCS_METADATA_HTTP_HANDLER = HTTPHandler(timeout=5.0)
return _GCS_METADATA_HTTP_HANDLER
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj
@ -171,12 +213,299 @@ def _apply_gemini_metadata(
return cast(PartType, part_dict)
def _parse_gs_uri(gs_uri: str) -> Tuple[str, str]:
if not gs_uri.startswith("gs://"):
raise ValueError(f"Invalid gs URI: {gs_uri}")
uri_without_scheme = gs_uri[5:] # drop gs://
uri_parts = uri_without_scheme.split("/", 1)
if len(uri_parts) != 2 or not uri_parts[0] or not uri_parts[1]:
raise ValueError(f"Invalid gs URI: {gs_uri}")
return uri_parts[0], uri_parts[1]
def _is_valid_gcs_bucket_name(bucket: str) -> bool:
"""
Validate bucket name against core GCS naming constraints.
"""
bucket_length = len(bucket)
max_bucket_length = 222 if "." in bucket else 63
if bucket_length < 3 or bucket_length > max_bucket_length:
return False
if "." in bucket and any(
len(label) == 0 or len(label) > 63 for label in bucket.split(".")
):
return False
if not re.fullmatch(r"[a-z0-9][a-z0-9._-]*[a-z0-9]", bucket):
return False
if ".." in bucket:
return False
if re.fullmatch(r"\d+\.\d+\.\d+\.\d+", bucket):
return False
return True
def _gs_uri_requires_content_type_metadata(url: str) -> bool:
"""
True when _process_gemini_media would call _get_gcs_object_content_type
(extension-less gs:// and no explicit format passed into that helper).
"""
if "gs://" not in url:
return False
extension_with_dot = os.path.splitext(url)[-1]
extension = extension_with_dot[1:] if extension_with_dot else ""
return len(extension) == 0
def _image_url_payload_may_need_sync_gcs_metadata_fetch(
raw_image_url: Any,
) -> bool:
"""
True when this image_url value (content-part image_url or assistant ``images[]``
entry) can trigger a blocking GCS metadata read for MIME resolution.
"""
fmt: Optional[str] = None
url: Optional[str] = None
if isinstance(raw_image_url, dict):
url = raw_image_url.get("url") # type: ignore[assignment]
if not isinstance(url, str):
return False
fmt = (
raw_image_url.get("format")
or raw_image_url.get("mime_type")
or raw_image_url.get("content_type")
)
elif isinstance(raw_image_url, str):
url = raw_image_url
else:
return False
if "gs://" not in url or fmt:
return False
return _gs_uri_requires_content_type_metadata(url)
def _openai_messages_may_need_sync_gcs_metadata_fetch(
messages: List[AllMessageValues],
) -> bool:
"""
Heuristic: True if any message part can trigger a blocking GCS JSON
metadata read inside _transform_request_body (extension-less gs:// without
explicit MIME hints). Covers user/system ``content`` parts and assistant
``images`` (same paths as ``_gemini_convert_messages_with_history``). Used
to decide whether ``async_transform_request_body`` should offload the sync
transform via ``asyncify``.
"""
for raw in messages:
msg: Any = raw
if not isinstance(msg, dict) and hasattr(msg, "model_dump"):
msg = msg.model_dump(exclude_none=False)
if not isinstance(msg, dict):
continue
images_field = msg.get("images")
if isinstance(images_field, list):
for image_item in images_field:
if not isinstance(image_item, dict):
continue
if _image_url_payload_may_need_sync_gcs_metadata_fetch(
image_item.get("image_url")
):
return True
content = msg.get("content")
if not isinstance(content, list):
continue
for item in content:
if not isinstance(item, dict):
continue
itype = item.get("type")
if itype == "image_url":
if _image_url_payload_may_need_sync_gcs_metadata_fetch(
item.get("image_url")
):
return True
elif itype == "file":
file_obj = item.get("file")
if not isinstance(file_obj, dict):
continue
fmt = (
file_obj.get("format")
or file_obj.get("mime_type")
or file_obj.get("content_type")
)
passed = file_obj.get("file_id") or file_obj.get("file_data")
if (
isinstance(passed, str)
and "gs://" in passed
and not fmt
and _gs_uri_requires_content_type_metadata(passed)
):
return True
return False
def _get_gcs_object_content_type(
image_url: str,
vertex_project: Optional[str] = None,
vertex_credentials: Optional[Any] = None,
) -> Optional[str]:
"""
Resolve content type from GCS object metadata.
Only attaches a Bearer token when the caller explicitly supplies Vertex
credentials, to avoid using the server's default Google credentials on
the Gemini API-key (Google AI Studio) path and being used as an oracle
for private GCS object metadata. Without explicit credentials we only
issue an anonymous request, which only succeeds for publicly-readable
objects.
"""
try:
bucket, object_name = _parse_gs_uri(image_url)
except ValueError:
return None
if not _is_valid_gcs_bucket_name(bucket):
return None
headers: Dict[str, str] = {}
explicit_vertex_auth_provided = (
vertex_project is not None or vertex_credentials is not None
)
if explicit_vertex_auth_provided:
try:
access_token, _ = _get_vertex_base().get_access_token(
credentials=vertex_credentials,
project_id=vertex_project,
)
headers["Authorization"] = f"Bearer {access_token}"
except Exception as e:
raise litellm.BadRequestError(
message=(
"Unable to fetch GCS metadata with provided Vertex credentials/project. "
f"Original error: {str(e)}"
),
model=None,
llm_provider="vertex_ai",
)
# Build the URL via httpx.URL with a fixed scheme/host and URL-encode both
# bucket and object so CodeQL does not flag the interpolation as a
# potential SSRF that could resolve to an arbitrary host.
encoded_bucket = quote(bucket, safe="")
encoded_object = quote(object_name, safe="")
metadata_url = httpx.URL(
scheme="https",
host="storage.googleapis.com",
path=f"/storage/v1/b/{encoded_bucket}/o/{encoded_object}",
params={"fields": "contentType"},
)
try:
response = _get_gcs_metadata_http_handler().get(
url=str(metadata_url),
headers=headers or None,
)
except httpx.RequestError as e:
if explicit_vertex_auth_provided:
raise litellm.BadRequestError(
message=(
"Unable to reach GCS JSON API for object metadata with provided "
f"Vertex credentials. {type(e).__name__}: {e}"
),
model=None,
llm_provider="vertex_ai",
) from e
return None
if response.is_error:
if explicit_vertex_auth_provided:
preview = (response.text or "")[:1024]
raise litellm.BadRequestError(
message=(
"Unable to read GCS object metadata with provided Vertex credentials. "
f"HTTP {response.status_code}. Response body (truncated): {preview!r}"
),
model=None,
llm_provider="vertex_ai",
)
return None
try:
payload = response.json()
except ValueError as e:
if explicit_vertex_auth_provided:
raise litellm.BadRequestError(
message=(
"GCS metadata response was not valid JSON when using provided "
f"Vertex credentials (HTTP {response.status_code}). Error: {e}"
),
model=None,
llm_provider="vertex_ai",
) from e
return None
if not isinstance(payload, dict):
if explicit_vertex_auth_provided:
raise litellm.BadRequestError(
message=(
"GCS metadata response was not a JSON object when using provided "
f"Vertex credentials (HTTP {response.status_code})."
),
model=None,
llm_provider="vertex_ai",
)
return None
content_type = payload.get("contentType")
if isinstance(content_type, str) and len(content_type) > 0:
return content_type
if explicit_vertex_auth_provided:
preview = (response.text or "")[:1024]
raise litellm.BadRequestError(
message=(
"GCS metadata JSON did not include a non-empty contentType field when "
f"using provided Vertex credentials (HTTP {response.status_code}). "
f"Body (truncated): {preview!r}"
),
model=None,
llm_provider="vertex_ai",
)
return None
def _normalize_and_validate_gemini_mime_type(
mime_type: str, model: Optional[str]
) -> str:
# Import lazily to avoid a module-level cyclic-import alert with
# litellm.types.files.
from litellm.types.files import get_file_extension_from_mime_type
normalized_mime_type = _apply_gemini_mime_type_aliases(mime_type)
try:
file_extension = get_file_extension_from_mime_type(normalized_mime_type)
file_type = get_file_type_from_extension(file_extension)
except ValueError:
raise litellm.BadRequestError(
message=f"File type not supported by gemini - {normalized_mime_type}",
model=model,
llm_provider="vertex_ai",
)
if not is_gemini_1_5_accepted_file_type(file_type):
raise litellm.BadRequestError(
message=f"File type not supported by gemini - {file_type}",
model=model,
llm_provider="vertex_ai",
)
return get_file_mime_type_for_file_type(file_type)
def _process_gemini_media(
image_url: str,
format: Optional[str] = None,
media_resolution_enum: Optional[Dict[str, str]] = None,
model: Optional[str] = None,
video_metadata: Optional[Dict[str, Any]] = None,
vertex_project: Optional[str] = None,
vertex_credentials: Optional[Any] = None,
) -> PartType:
"""
Given a media URL (image, audio, or video), return the appropriate PartType for Gemini
@ -193,20 +522,63 @@ def _process_gemini_media(
try:
# GCS URIs
if "gs://" in image_url:
# Figure out file type
extension_with_dot = os.path.splitext(image_url)[-1] # Ex: ".png"
extension = extension_with_dot[1:] # Ex: "png"
explicit_gcs_format = False
if not format:
file_type = get_file_type_from_extension(extension)
mime_type: Optional[str] = None
# For extension-less gs:// URIs, we cannot infer from path.
# If callers pass `format`/`mime_type`, this branch is skipped.
if extension:
file_type = get_file_type_from_extension(extension)
# Validate the file type is supported by Gemini
if not is_gemini_1_5_accepted_file_type(file_type):
raise Exception(f"File type not supported by gemini - {file_type}")
# Validate the file type is supported by Gemini
if not is_gemini_1_5_accepted_file_type(file_type):
raise litellm.BadRequestError(
message=f"File type not supported by gemini - {file_type}",
model=model,
llm_provider="vertex_ai",
)
mime_type = get_file_mime_type_for_file_type(file_type)
mime_type = get_file_mime_type_for_file_type(file_type)
else:
mime_type = _get_gcs_object_content_type(
image_url=image_url,
vertex_project=vertex_project,
vertex_credentials=vertex_credentials,
)
if mime_type is None:
raise litellm.BadRequestError(
message=(
f"Unable to determine mime type for gs URI: {image_url}. "
"This gs:// URI has no file extension and GCS metadata "
"lookup failed. Set it explicitly using image_url.format "
"(or image_url.mime_type/content_type) or "
"message.content[].file.format."
),
model=model,
llm_provider="vertex_ai",
)
else:
mime_type = format
explicit_gcs_format = True
if mime_type is None:
raise litellm.BadRequestError(
message=f"File type not supported by gemini - {image_url}",
model=model,
llm_provider="vertex_ai",
)
if explicit_gcs_format:
# Callers who pass format/mime_type explicitly for gs:// URIs
# rely on pass-through to Gemini (pre-PR behavior). Only apply
# known MIME aliases; skip litellm's file-type registry.
mime_type = _apply_gemini_mime_type_aliases(mime_type)
else:
mime_type = _normalize_and_validate_gemini_mime_type(
mime_type=mime_type,
model=model,
)
file_data = FileDataType(mime_type=mime_type, file_uri=image_url)
part: PartType = {"file_data": file_data}
return _apply_gemini_metadata(
@ -258,8 +630,6 @@ def _snake_to_camel(snake_str: str) -> str:
def _camel_to_snake(camel_str: str) -> str:
"""Convert camelCase to snake_case"""
import re
return re.sub(r"(?<!^)(?=[A-Z])", "_", camel_str).lower()
@ -311,6 +681,7 @@ def check_if_part_exists_in_parts(
def _gemini_convert_messages_with_history( # noqa: PLR0915
messages: List[AllMessageValues],
model: Optional[str] = None,
litellm_params: Optional[dict] = None,
) -> List[ContentType]:
"""
Converts given messages from OpenAI format to Gemini format
@ -326,6 +697,16 @@ def _gemini_convert_messages_with_history( # noqa: PLR0915
msg_i = 0
tool_call_responses = []
vertex_project = None
vertex_credentials = None
if litellm_params:
vertex_project = litellm_params.get("vertex_project") or litellm_params.get(
"vertex_ai_project"
)
vertex_credentials = litellm_params.get(
"vertex_credentials"
) or litellm_params.get("vertex_ai_credentials")
try:
while msg_i < len(messages):
user_content: List[PartType] = []
@ -351,20 +732,42 @@ def _gemini_convert_messages_with_history( # noqa: PLR0915
img_element = element
format: Optional[str] = None
media_resolution_enum: Optional[Dict[str, str]] = None
if isinstance(img_element["image_url"], dict):
image_url = img_element["image_url"]["url"]
format = img_element["image_url"].get("format")
detail = img_element["image_url"].get("detail")
raw_image_url = img_element.get("image_url")
if raw_image_url is None:
raise litellm.BadRequestError(
message="Invalid message content: element type is 'image_url' but 'image_url' field is missing ",
model=model,
llm_provider="vertex_ai",
)
if isinstance(raw_image_url, dict):
image_url = raw_image_url.get("url")
if image_url is None:
raise litellm.BadRequestError(
message="Invalid message content: element type is 'image_url' but 'url' field is missing inside 'image_url' ",
model=model,
llm_provider="vertex_ai",
)
# TypedDict does not declare mime_type/content_type;
# read via Dict[str, Any] for caller-provided MIME fields.
image_url_dict = cast(Dict[str, Any], raw_image_url)
format = (
image_url_dict.get("format")
or image_url_dict.get("mime_type")
or image_url_dict.get("content_type")
)
detail = image_url_dict.get("detail")
media_resolution_enum = (
_convert_detail_to_media_resolution_enum(detail)
)
else:
image_url = img_element["image_url"]
image_url = raw_image_url
_part = _process_gemini_media(
image_url=image_url,
format=format,
media_resolution_enum=media_resolution_enum,
model=model,
vertex_project=vertex_project,
vertex_credentials=vertex_credentials,
)
_parts.append(_part)
elif element["type"] == "input_audio":
@ -390,15 +793,31 @@ def _gemini_convert_messages_with_history( # noqa: PLR0915
image_url=openai_image_str,
format=audio_format_modified,
model=model,
vertex_project=vertex_project,
vertex_credentials=vertex_credentials,
)
_parts.append(_part)
elif element["type"] == "file":
file_element = cast(ChatCompletionFileObject, element)
file_id = file_element["file"].get("file_id")
format = file_element["file"].get("format")
file_data = file_element["file"].get("file_data")
detail = file_element["file"].get("detail")
video_metadata = file_element["file"].get("video_metadata")
_file_field = file_element.get("file")
if _file_field is None:
raise litellm.BadRequestError(
message="Content block has type='file' but is missing the required 'file' field",
model=model,
llm_provider="vertex_ai",
)
# TypedDict does not declare mime_type/content_type;
# read via Dict[str, Any] for caller-provided MIME fields.
file_dict = cast(Dict[str, Any], _file_field)
file_id = file_dict.get("file_id")
format = (
file_dict.get("format")
or file_dict.get("mime_type")
or file_dict.get("content_type")
)
file_data = file_dict.get("file_data")
detail = file_dict.get("detail")
video_metadata = file_dict.get("video_metadata")
passed_file = file_id or file_data
if passed_file is None:
raise Exception(
@ -417,13 +836,23 @@ def _gemini_convert_messages_with_history( # noqa: PLR0915
model=model,
media_resolution_enum=media_resolution_enum,
video_metadata=video_metadata,
vertex_project=vertex_project,
vertex_credentials=vertex_credentials,
)
_parts.append(_part)
except Exception:
raise Exception(
"Unable to determine mime type for file_id: {}, set this explicitly using message[{}].content[{}].file.format".format(
file_id, msg_i, element_idx
)
except litellm.BadRequestError:
raise
except Exception as e:
raise litellm.BadRequestError(
message=(
f"Unable to determine mime type for file: "
f"{file_id or 'provided data'}, set this explicitly "
f"using message[{msg_i}].content[{element_idx}].file.format "
f"(or file.mime_type/content_type). "
f"Original error: {str(e)}"
),
model=model,
llm_provider="vertex_ai",
)
user_content.extend(_parts)
elif _message_content is not None and isinstance(_message_content, str):
@ -528,7 +957,11 @@ def _gemini_convert_messages_with_history( # noqa: PLR0915
image_url_obj = image_item.get("image_url")
if isinstance(image_url_obj, dict):
assistant_image_url = image_url_obj.get("url")
format = image_url_obj.get("format")
format = (
image_url_obj.get("format")
or image_url_obj.get("mime_type")
or image_url_obj.get("content_type")
)
detail = image_url_obj.get("detail")
media_resolution_enum = (
_convert_detail_to_media_resolution_enum(detail)
@ -539,6 +972,8 @@ def _gemini_convert_messages_with_history( # noqa: PLR0915
format=format,
media_resolution_enum=media_resolution_enum,
model=model,
vertex_project=vertex_project,
vertex_credentials=vertex_credentials,
)
assistant_content.append(_part)
@ -713,11 +1148,11 @@ def _transform_request_body( # noqa: PLR0915
try:
if custom_llm_provider == "gemini":
content = litellm.GoogleAIStudioGeminiConfig()._transform_messages(
messages=messages, model=model
messages=messages, model=model, litellm_params=litellm_params
)
else:
content = litellm.VertexGeminiConfig()._transform_messages(
messages=messages, model=model
messages=messages, model=model, litellm_params=litellm_params
)
tools: Optional[Tools] = optional_params.pop("tools", None)
tool_choice: Optional[ToolConfig] = optional_params.pop("tool_choice", None)
@ -893,6 +1328,20 @@ async def async_transform_request_body(
vertex_auth_header=vertex_auth_header,
)
if _openai_messages_may_need_sync_gcs_metadata_fetch(messages):
# _transform_request_body may issue a sync httpx.get (up to 5s timeout)
# via _get_gcs_object_content_type to fetch GCS object metadata. Run the
# whole sync transformation on a worker thread so it does not block the
# async event loop.
return await asyncify(_transform_request_body)(
messages=messages,
model=model,
custom_llm_provider=custom_llm_provider,
litellm_params=litellm_params,
cached_content=cached_content,
optional_params=optional_params,
)
return _transform_request_body(
messages=messages,
model=model,

View file

@ -2533,9 +2533,14 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
return model_response
def _transform_messages(
self, messages: List[AllMessageValues], model: Optional[str] = None
self,
messages: List[AllMessageValues],
model: Optional[str] = None,
litellm_params: Optional[dict] = None,
) -> List[ContentType]:
return _gemini_convert_messages_with_history(messages=messages, model=model)
return _gemini_convert_messages_with_history(
messages=messages, model=model, litellm_params=litellm_params
)
def get_error_class(
self, error_message: str, status_code: int, headers: Union[Dict, httpx.Headers]
@ -3139,6 +3144,31 @@ class ModelResponseIterator:
self.cumulative_tool_call_index: int = 0
self.has_seen_tool_calls: bool = False
@staticmethod
def _check_streaming_error(chunk: dict) -> None:
"""Detect embedded errors (e.g. 429 RESOURCE_EXHAUSTED) in streaming chunks and raise VertexAIError."""
if "error" not in chunk:
return
error_data = chunk["error"]
if not isinstance(error_data, dict):
raise VertexAIError(
status_code=500,
message=f"Unexpected error format in mid-stream chunk: {error_data}",
)
raw_code = error_data.get("code", 500)
if raw_code is None:
raw_code = 500
try:
error_code = int(raw_code)
except (TypeError, ValueError):
error_code = 500
error_message = error_data.get("message", "Unknown error")
error_status = error_data.get("status", "UNKNOWN")
raise VertexAIError(
status_code=error_code,
message=f"{error_status} - {error_message}",
)
def _apply_stream_candidates(
self,
_candidates: List[Candidates],
@ -3256,6 +3286,11 @@ class ModelResponseIterator:
def chunk_parser(self, chunk: dict) -> Optional["ModelResponseStream"]:
try:
verbose_logger.debug(f"RAW GEMINI CHUNK: {chunk}")
# Detect mid-stream error chunks (e.g. 429 RESOURCE_EXHAUSTED).
# Vertex AI can return errors as HTTP 200 but with an "error" field in the SSE body.
self._check_streaming_error(chunk)
from litellm.types.utils import ModelResponseStream
processed_chunk = GenerateContentResponseBody(**chunk) # type: ignore

View file

@ -292,22 +292,15 @@ class VertexAIPartnerModels(VertexBase):
Returns:
Dict containing token count information
"""
try:
import vertexai
except Exception as e:
raise VertexAIError(
status_code=400,
message=f"""vertexai import failed please run `pip install -U "google-cloud-aiplatform>=1.38"`. Got error: {e}""",
)
if not (
hasattr(vertexai, "preview") or hasattr(vertexai.preview, "language_models")
):
raise VertexAIError(
status_code=400,
message="""Upgrade vertex ai. Run `pip install "google-cloud-aiplatform>=1.38"`""",
)
# Note: we intentionally do not import `vertexai` (the Gemini SDK shipped
# by `google-cloud-aiplatform`) on this path. Partner models such as
# Claude on Vertex use the Anthropic Messages API protocol directly via
# `:rawPredict`, and `VertexAIPartnerModelsTokenCounter` reaches that
# endpoint with an authenticated httpx client — it never touches the
# Gemini SDK. Requiring `google-cloud-aiplatform>=1.38` here turned a
# SDK-free Anthropic-protocol call into a hard dependency on the Gemini
# SDK (see #28084), breaking `/v1/messages/count_tokens` for Claude-on-
# Vertex on any LiteLLM install without that extra. Stay SDK-free.
try:
from litellm.llms.vertex_ai.vertex_ai_partner_models.count_tokens.handler import (
VertexAIPartnerModelsTokenCounter,

View file

@ -5720,6 +5720,33 @@ def embedding( # noqa: PLR0915
aembedding=aembedding,
headers=headers,
)
elif custom_llm_provider == "dashscope":
dashscope_key = (
api_key or litellm.api_key or get_secret_str("DASHSCOPE_API_KEY")
)
if dashscope_key is None:
raise ValueError(
"Missing API key for DashScope. Set DASHSCOPE_API_KEY environment variable or pass api_key parameter."
)
if extra_headers is not None and isinstance(extra_headers, dict):
headers = extra_headers
else:
headers = {}
response = base_llm_http_handler.embedding(
model=model,
input=input,
timeout=timeout,
custom_llm_provider=custom_llm_provider,
logging_obj=logging,
api_base=api_base,
optional_params=optional_params,
litellm_params={},
model_response=EmbeddingResponse(),
api_key=dashscope_key,
client=client,
aembedding=aembedding,
headers=headers,
)
elif custom_llm_provider == "ovhcloud":
api_key = api_key or litellm.api_key or get_secret_str("OVHCLOUD_API_KEY")
api_base = (

View file

@ -4549,6 +4549,7 @@ class PrismaCompatibleUpdateDBModel(TypedDict, total=False):
model_name: str
litellm_params: str
model_info: str
blocked: bool
updated_at: str
updated_by: str

View file

@ -69,6 +69,35 @@ else:
ProxyLogging = Any
def _extract_cache_read_tokens(usage_obj: dict) -> int:
"""
Anthropic: top-level cache_read_input_tokens field.
OpenAI-compatible (moonshotai, openai, deepseek, etc.): prompt_tokens_details.cached_tokens.
"""
explicit = usage_obj.get("cache_read_input_tokens", 0) or 0
if explicit:
return int(explicit)
details = usage_obj.get("prompt_tokens_details") or {}
return int(details.get("cached_tokens", 0) or 0)
def _extract_cache_creation_tokens(usage_obj: dict) -> int:
"""
Anthropic: top-level cache_creation_input_tokens field.
OpenAI-compatible (kimi-k2 etc.): prompt_tokens_details.cache_write_tokens
or prompt_tokens_details.cache_creation_tokens.
"""
explicit = usage_obj.get("cache_creation_input_tokens", 0) or 0
if explicit:
return int(explicit)
details = usage_obj.get("prompt_tokens_details") or {}
return int(
details.get("cache_write_tokens", 0)
or details.get("cache_creation_tokens", 0)
or 0
)
class DBSpendUpdateWriter:
"""
Module responsible for
@ -1992,12 +2021,8 @@ class DBSpendUpdateWriter:
api_requests=1,
successful_requests=1 if request_status == "success" else 0,
failed_requests=1 if request_status != "success" else 0,
cache_read_input_tokens=usage_obj.get("cache_read_input_tokens", 0)
or 0,
cache_creation_input_tokens=usage_obj.get(
"cache_creation_input_tokens", 0
)
or 0,
cache_read_input_tokens=_extract_cache_read_tokens(usage_obj),
cache_creation_input_tokens=_extract_cache_creation_tokens(usage_obj),
)
return daily_transaction
except Exception as e:

View file

@ -1,6 +1,7 @@
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
from fastapi import HTTPException, status
from pydantic import BaseModel
from litellm._logging import verbose_proxy_logger
from litellm.caching import DualCache
@ -53,6 +54,37 @@ def require_caller_user_id_for_non_admin(
return user_api_key_dict.user_id
def _check_passthrough_routes_caller_permission(
data: BaseModel,
user_api_key_dict: UserAPIKeyAuth,
*,
entity: str = "key",
) -> None:
"""
Only proxy admins may set `allowed_passthrough_routes` (top-level or under
`metadata`) it short-circuits the role-based route gate, so keys and teams
must be gated identically.
"""
# view-only admins excluded by design; blocked upstream from writes anyway
if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value:
return
if getattr(data, "allowed_passthrough_routes", None):
raise HTTPException(
status_code=403,
detail={
"error": f"Only proxy admins can set `allowed_passthrough_routes` on a {entity}."
},
)
metadata = getattr(data, "metadata", None)
if isinstance(metadata, dict) and metadata.get("allowed_passthrough_routes"):
raise HTTPException(
status_code=403,
detail={
"error": f"Only proxy admins can set `metadata.allowed_passthrough_routes` on a {entity}."
},
)
def _is_user_team_admin(
user_api_key_dict: UserAPIKeyAuth, team_obj: LiteLLM_TeamTable
) -> bool:

View file

@ -55,6 +55,7 @@ from litellm.proxy.common_utils.rbac_utils import check_org_admin_can_generate_k
from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time
from litellm.proxy.hooks.key_management_event_hooks import KeyManagementEventHooks
from litellm.proxy.management_endpoints.common_utils import (
_check_passthrough_routes_caller_permission,
_is_user_org_admin_for_team,
_is_user_team_admin,
_set_object_metadata_field,
@ -548,36 +549,6 @@ def _check_allowed_routes_caller_permission(
)
def _check_passthrough_routes_caller_permission(
data: BaseModel,
user_api_key_dict: UserAPIKeyAuth,
) -> None:
"""
Only proxy admins may set `allowed_passthrough_routes` on a key, either at
the top level of the request or nested under `metadata`.
The route gate evaluates passthrough access ahead of the standard role
gate, so the field is restricted to admins to keep that ordering safe.
"""
if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value:
return
if getattr(data, "allowed_passthrough_routes", None):
raise HTTPException(
status_code=403,
detail={
"error": "Only proxy admins can set `allowed_passthrough_routes` on a key."
},
)
metadata = getattr(data, "metadata", None)
if isinstance(metadata, dict) and metadata.get("allowed_passthrough_routes"):
raise HTTPException(
status_code=403,
detail={
"error": "Only proxy admins can set `metadata.allowed_passthrough_routes` on a key."
},
)
async def validate_team_id_used_in_service_account_request(
team_id: Optional[str],
prisma_client: Optional[PrismaClient],

View file

@ -150,6 +150,9 @@ def update_db_model(
model_info[key] = value.isoformat()
prisma_compatible_model_dict["model_info"] = json.dumps(model_info)
if updated_patch.blocked is not None:
prisma_compatible_model_dict["blocked"] = updated_patch.blocked
return prisma_compatible_model_dict
@ -230,6 +233,20 @@ async def patch_model(
premium_user=premium_user,
)
# Pause/resume (`blocked`) is a proxy-admin-only privilege. Team admins
# passed the auth check above for team-scoped models, but they must not
# be able to unblock (or block) a model their proxy admin has paused.
if (
patch_data.blocked is not None
and user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN
):
raise ProxyException(
message="Only proxy admins can change a model's blocked flag.",
type=ProxyErrorTypes.auth_error.value,
code=status.HTTP_403_FORBIDDEN,
param="blocked",
)
# Handle team model updates with proper alias management
update_data = await _update_team_model_in_db(
db_model=db_model,

View file

@ -73,6 +73,7 @@ from litellm.proxy.auth.auth_checks import (
)
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.management_endpoints.common_utils import (
_check_passthrough_routes_caller_permission,
_is_user_org_admin_for_team,
_is_user_team_admin,
_set_object_metadata_field,
@ -1049,6 +1050,10 @@ async def new_team( # noqa: PLR0915
Member(role="admin", user_id=user_api_key_dict.user_id)
)
_check_passthrough_routes_caller_permission(
data, user_api_key_dict, entity="team"
)
## ADD TO MODEL TABLE
_model_id = None
if data.model_aliases is not None and isinstance(data.model_aliases, dict):
@ -1646,6 +1651,10 @@ async def update_team( # noqa: PLR0915
user_api_key_dict=user_api_key_dict,
)
_check_passthrough_routes_caller_permission(
data, user_api_key_dict, entity="team"
)
if data.soft_budget is not None:
max_budget_to_check = (
data.max_budget

View file

@ -4727,6 +4727,7 @@ class ProxyConfig:
if _id is not None:
model.model_info["id"] = _id
model.model_info["db_model"] = True
model.model_info["blocked"] = bool(getattr(model, "blocked", False))
if premium_user is True:
# seeing "created_at", "updated_at", "created_by", "updated_by" is a LiteLLM Enterprise Feature
@ -8089,6 +8090,11 @@ async def model_list(
proxy_logging_obj=proxy_logging_obj,
)
# Compute once — used in both branches below to hide paused models from the listing.
blocked_names = (
llm_router.get_fully_blocked_model_names() if llm_router is not None else set()
)
# If scope=expand and user has admin privileges, return all proxy models
if should_expand_scope:
# Get all proxy models as if user is a proxy admin
@ -8121,6 +8127,10 @@ async def model_list(
only_model_access_groups=only_model_access_groups or False,
)
# Hide paused models from the public listing (admins manage them via /model/info)
if blocked_names:
all_models = [m for m in all_models if m not in blocked_names]
# Build response data with all proxy models
model_data = []
for model in all_models:
@ -8154,6 +8164,10 @@ async def model_list(
user_api_key_cache=user_api_key_cache,
)
# Hide paused models from the public listing (admins manage them via /model/info)
if blocked_names:
all_models = [m for m in all_models if m not in blocked_names]
# Build response data
model_data = []
for model in all_models:

View file

@ -430,7 +430,11 @@ async def route_request( # noqa: PLR0915 - Complex routing function, refactorin
deployment = llm_router.get_deployment_by_model_group_name(
model_group_name=model
)
if deployment and deployment.litellm_params:
if (
deployment
and deployment.litellm_params
and not llm_router._is_deployment_blocked(deployment)
):
deployment_creds = deployment.litellm_params.model_dump(
exclude_none=True
)

View file

@ -48,9 +48,10 @@ model LiteLLM_CredentialsTable {
// Models on proxy
model LiteLLM_ProxyModelTable {
model_id String @id @default(uuid())
model_name String
model_name String
litellm_params Json
model_info Json?
model_info Json?
blocked Boolean @default(false)
created_at DateTime @default(now()) @map("created_at")
created_by String
updated_at DateTime @default(now()) @updatedAt @map("updated_at")

View file

@ -30,6 +30,7 @@ from typing import (
List,
Literal,
Optional,
Set,
Tuple,
Union,
cast,
@ -6957,12 +6958,11 @@ class Router:
unhealthy_deployments = _get_cooldown_deployments(
litellm_router_instance=self, parent_otel_span=parent_otel_span
)
healthy_deployments: list = []
for deployment in _all_deployments:
if deployment["model_info"]["id"] in unhealthy_deployments:
continue
else:
healthy_deployments.append(deployment)
unhealthy_set = set(unhealthy_deployments)
healthy_deployments: list = [
d for d in _all_deployments if d["model_info"]["id"] not in unhealthy_set
]
healthy_deployments = self._filter_blocked_deployments(healthy_deployments)
return healthy_deployments, _all_deployments
@ -6990,10 +6990,12 @@ class Router:
)
# Convert to set for O(1) lookup instead of O(n)
unhealthy_deployments_set = set(unhealthy_deployments)
healthy_deployments: list = []
for deployment in _all_deployments:
if deployment["model_info"]["id"] not in unhealthy_deployments_set:
healthy_deployments.append(deployment)
healthy_deployments: list = [
d
for d in _all_deployments
if d["model_info"]["id"] not in unhealthy_deployments_set
]
healthy_deployments = self._filter_blocked_deployments(healthy_deployments)
return healthy_deployments, _all_deployments
def routing_strategy_pre_call_checks(self, deployment: dict):
@ -8142,10 +8144,14 @@ class Router:
def get_deployment_credentials(self, model_id: str) -> Optional[dict]:
"""
Returns -> dict of credentials for a given model id
Returns -> dict of credentials for a given model id.
Returns None if the deployment is paused via `LiteLLM_ProxyModelTable.blocked`,
so file/batch/passthrough callers that resolve credentials directly cannot keep
using a paused deployment.
"""
deployment = self.get_deployment(model_id=model_id)
if deployment is None:
if deployment is None or self._is_deployment_blocked(deployment):
return None
return CredentialLiteLLMParams(
**deployment.litellm_params.model_dump(exclude_none=True)
@ -8190,7 +8196,9 @@ class Router:
Returns:
Dictionary containing api_key, api_base, custom_llm_provider, etc.
Returns None if model not found.
Returns None if model not found, or if the resolved deployment is
paused via `LiteLLM_ProxyModelTable.blocked` (so passthrough callers
cannot bypass an admin pause by resolving credentials directly).
Example:
credentials = router.get_deployment_credentials_with_provider("gpt-4o-litellm")
@ -8216,7 +8224,7 @@ class Router:
elif isinstance(deployment_dict, Deployment):
deployment = deployment_dict
if deployment is None:
if deployment is None or self._is_deployment_blocked(deployment):
return None
# Get basic credentials
@ -9243,6 +9251,29 @@ class Router:
return model_names
def get_fully_blocked_model_names(self) -> Set[str]:
"""
Returns the set of model_names where every backing deployment has `blocked=True`.
Used by `/v1/models` to hide paused models from client listings while still
surfacing them on admin endpoints (e.g. `/model/info`). A model with at least
one non-blocked deployment is still serviceable and remains visible.
"""
deployments = self.get_model_list() or []
blocked_by_name: Dict[str, bool] = {}
for deployment in deployments:
name = deployment.get("model_name") or ""
if not name:
continue
is_blocked = (deployment.get("model_info") or {}).get("blocked") is True
if name in blocked_by_name:
blocked_by_name[name] = blocked_by_name[name] and is_blocked
else:
blocked_by_name[name] = is_blocked
return {
name for name, fully_blocked in blocked_by_name.items() if fully_blocked
}
def _get_team_specific_model(
self, deployment: DeploymentTypedDict, team_id: Optional[str] = None
) -> Optional[str]:
@ -10131,6 +10162,12 @@ class Router:
)
if isinstance(healthy_deployments, dict):
if (healthy_deployments.get("model_info") or {}).get("blocked") is True:
raise litellm.ServiceUnavailableError(
message=f"Model '{model}' is administratively paused. Contact your proxy admin to unblock it.",
model=model,
llm_provider="",
)
return healthy_deployments
# Health-check-based filtering (before cooldown)
@ -10164,6 +10201,8 @@ class Router:
)
healthy_deployments = _pre_cooldown_deployments
healthy_deployments = self._filter_blocked_deployments(healthy_deployments)
healthy_deployments = await self.async_callback_filter_deployments(
model=model,
healthy_deployments=healthy_deployments,
@ -10387,6 +10426,12 @@ class Router:
# 3. If specific deployment returned, verify if it supports pass-through
if isinstance(healthy_deployments, dict):
if (healthy_deployments.get("model_info") or {}).get("blocked") is True:
raise litellm.ServiceUnavailableError(
message=f"Model '{model}' is administratively paused. Contact your proxy admin to unblock it.",
model=model,
llm_provider="",
)
litellm_params = healthy_deployments.get("litellm_params", {})
if litellm_params.get("use_in_pass_through"):
return healthy_deployments
@ -10555,6 +10600,12 @@ class Router:
)
if isinstance(healthy_deployments, dict):
if (healthy_deployments.get("model_info") or {}).get("blocked") is True:
raise litellm.ServiceUnavailableError(
message=f"Model '{model}' is administratively paused. Contact your proxy admin to unblock it.",
model=model,
llm_provider="",
)
return healthy_deployments
parent_otel_span: Optional[Span] = _get_parent_otel_span_from_kwargs(
@ -10585,6 +10636,8 @@ class Router:
)
healthy_deployments = _pre_cooldown_deployments
healthy_deployments = self._filter_blocked_deployments(healthy_deployments)
# filter pre-call checks
if self.enable_pre_call_checks and messages is not None:
healthy_deployments = self._pre_call_checks(
@ -10704,6 +10757,12 @@ class Router:
# 2. If the returned is a specific deployment (Dict), verify and return directly
if isinstance(healthy_deployments, dict):
if (healthy_deployments.get("model_info") or {}).get("blocked") is True:
raise litellm.ServiceUnavailableError(
message=f"Model '{model}' is administratively paused. Contact your proxy admin to unblock it.",
model=model,
llm_provider="",
)
litellm_params = healthy_deployments.get("litellm_params", {})
if litellm_params.get("use_in_pass_through"):
return healthy_deployments
@ -10743,6 +10802,9 @@ class Router:
healthy_deployments=pass_through_deployments,
cooldown_deployments=cooldown_deployments,
)
pass_through_deployments = self._filter_blocked_deployments(
pass_through_deployments
)
# 5. Apply pre-call checks (if enabled)
if self.enable_pre_call_checks and messages is not None:
@ -10832,6 +10894,36 @@ class Router:
if deployment["model_info"]["id"] not in cooldown_set
]
def _filter_blocked_deployments(
self, healthy_deployments: List[Dict]
) -> List[Dict]:
"""
Filters out deployments that an admin has paused via `LiteLLM_ProxyModelTable.blocked`.
Applied alongside the cooldown filter on every routing entry point that calls
`_common_checks_available_deployment` directly the primary sync/async path,
the sync pass-through path, and the retry / health-check helpers so paused
deployments never serve a request. The async pass-through path inherits this
filter through its delegation to `async_get_healthy_deployments`.
"""
return [
deployment
for deployment in healthy_deployments
if (deployment.get("model_info") or {}).get("blocked") is not True
]
@staticmethod
def _is_deployment_blocked(deployment: "Deployment") -> bool:
"""
Returns True when a `Deployment` Pydantic instance carries the admin-paused
flag. Used by credential-lookup helpers so passthrough file / batch endpoints
cannot bypass the pause by resolving credentials directly.
"""
model_info = getattr(deployment, "model_info", None)
if model_info is None:
return False
return getattr(model_info, "blocked", None) is True
async def _async_filter_health_check_unhealthy_deployments(
self,
healthy_deployments: List[Dict],

View file

@ -160,6 +160,7 @@ class UserAPIKeyLabelNames(Enum):
END_USER = "end_user"
USER = "user"
USER_EMAIL = "user_email"
USER_ALIAS = "user_alias"
API_KEY_HASH = "hashed_api_key"
API_KEY_ALIAS = "api_key_alias"
TEAM = "team"
@ -533,17 +534,9 @@ class PrometheusMetricLabels:
UserAPIKeyLabelNames.USER.value,
]
litellm_user_max_budget_metric = [
UserAPIKeyLabelNames.USER.value,
]
litellm_user_max_budget_metric = litellm_remaining_user_budget_metric
litellm_user_budget_remaining_hours_metric = [
UserAPIKeyLabelNames.USER.value,
]
litellm_user_budget_remaining_hours_metric = [
UserAPIKeyLabelNames.USER.value,
]
litellm_user_budget_remaining_hours_metric = litellm_remaining_user_budget_metric
litellm_remaining_api_key_requests_for_model = [
UserAPIKeyLabelNames.API_KEY_HASH.value,
@ -730,6 +723,22 @@ class PrometheusMetricLabels:
):
custom_labels.append(UserAPIKeyLabelNames.STREAM.value)
_user_budget_metrics = {
"litellm_remaining_user_budget_metric",
"litellm_user_max_budget_metric",
"litellm_user_budget_remaining_hours_metric",
}
if (
label_name in _user_budget_metrics
and litellm.prometheus_user_budget_label_include_email_alias is True
):
for label in [
UserAPIKeyLabelNames.USER_EMAIL.value,
UserAPIKeyLabelNames.USER_ALIAS.value,
]:
if label not in default_labels and label not in custom_labels:
custom_labels.append(label)
if label_name in PrometheusMetricLabels._org_label_metrics:
for label in [
UserAPIKeyLabelNames.ORG_ID.value,
@ -759,6 +768,7 @@ class UserAPIKeyLabelValues:
end_user: Optional[str] = None
user: Optional[str] = None
user_email: Optional[str] = None
user_alias: Optional[str] = None
hashed_api_key: Optional[str] = None
api_key_alias: Optional[str] = None
team: Optional[str] = None

View file

@ -133,6 +133,9 @@ class ModelInfo(BaseModel):
# the model_name that can be used by the team when making LLM calls
team_public_model_name: Optional[str] = None
# admin-toggled pause flag; mirrors LiteLLM_ProxyModelTable.blocked
blocked: Optional[bool] = None
def __init__(self, id: Optional[Union[str, int]] = None, **params):
if id is None:
id = str(uuid.uuid4()) # Generate a UUID if id is None or not provided
@ -323,6 +326,7 @@ class updateDeployment(BaseModel):
model_name: Optional[str] = None
litellm_params: Optional[updateLiteLLMParams] = None
model_info: Optional[ModelInfo] = None
blocked: Optional[bool] = None
model_config = ConfigDict(protected_namespaces=())

View file

@ -107,9 +107,13 @@ class LiteLLMCommonStrings(Enum):
SupportedCacheControls = ["ttl", "s-maxage", "no-cache", "no-store"]
class CostPerToken(TypedDict):
input_cost_per_token: float
output_cost_per_token: float
class CostPerToken(TypedDict, total=False):
# Required base rates — kept under total=False so we can mark them
# Required individually while leaving the cache rates NotRequired.
input_cost_per_token: Required[float]
output_cost_per_token: Required[float]
cache_read_input_token_cost: float
cache_creation_input_token_cost: float
class ProviderField(TypedDict):

View file

@ -8402,6 +8402,12 @@ class ProviderConfigManager:
)
return VolcEngineEmbeddingConfig()
elif litellm.LlmProviders.DASHSCOPE == provider:
from litellm.llms.dashscope.embed.transformation import (
DashScopeEmbeddingConfig,
)
return DashScopeEmbeddingConfig()
elif litellm.LlmProviders.OVHCLOUD == provider:
return litellm.OVHCloudEmbeddingConfig()
elif litellm.LlmProviders.SNOWFLAKE == provider:
@ -8481,6 +8487,12 @@ class ProviderConfigManager:
return litellm.VoyageRerankConfig()
elif litellm.LlmProviders.WATSONX == provider:
return litellm.IBMWatsonXRerankConfig()
elif litellm.LlmProviders.DASHSCOPE == provider:
from litellm.llms.dashscope.rerank.transformation import (
DashScopeRerankConfig,
)
return DashScopeRerankConfig()
return litellm.CohereRerankConfig()
@staticmethod
@ -8527,6 +8539,12 @@ class ProviderConfigManager:
)
return MinimaxMessagesConfig()
elif litellm.LlmProviders.DEEPSEEK == provider:
from litellm.llms.deepseek.messages.transformation import (
DeepSeekAnthropicMessagesConfig,
)
return DeepSeekAnthropicMessagesConfig()
return None
@staticmethod

View file

@ -1,6 +1,6 @@
[project]
name = "litellm"
version = "1.85.0"
version = "1.86.0"
description = "Library to easily interface with LLM API providers"
readme = "README.md"
requires-python = ">=3.10, <3.14"
@ -57,7 +57,7 @@ proxy = [
"azure-storage-blob==12.28.0",
"mcp==1.26.0",
"litellm-proxy-extras==0.4.72",
"litellm-enterprise==0.1.40",
"litellm-enterprise==0.1.41",
"RestrictedPython==8.1",
"rich==13.9.4",
"polars==1.38.1",
@ -251,7 +251,7 @@ source-exclude = [
profile = "black"
[tool.commitizen]
version = "1.85.0"
version = "1.86.0"
version_files = [
"pyproject.toml:^version",
]

View file

@ -48,9 +48,10 @@ model LiteLLM_CredentialsTable {
// Models on proxy
model LiteLLM_ProxyModelTable {
model_id String @id @default(uuid())
model_name String
model_name String
litellm_params Json
model_info Json?
model_info Json?
blocked Boolean @default(false)
created_at DateTime @default(now()) @map("created_at")
created_by String
updated_at DateTime @default(now()) @updatedAt @map("updated_at")

View file

@ -36,6 +36,75 @@ SAFE_BODY_MATCHER_NAME = "safe_body"
KEY_FINGERPRINT_MATCHER_NAME = "key_fingerprint"
KEY_FINGERPRINT_HEADER = "x-litellm-key-fp"
VCR_DIAG_DIR_ENV = "LITELLM_VCR_DIAG_DIR"
VCR_DIAG_DIR_DEFAULT = "test-results/vcr-diagnostics"
def _vcr_diag_dir() -> str:
return os.environ.get(VCR_DIAG_DIR_ENV) or VCR_DIAG_DIR_DEFAULT
def vcr_diag_write_line(msg: str) -> None:
try:
directory = _vcr_diag_dir()
os.makedirs(directory, exist_ok=True)
path = os.path.join(directory, f"{os.getpid()}.log")
with open(path, "a", encoding="utf-8") as fh:
fh.write(msg.rstrip("\n") + "\n")
except OSError:
pass
def reset_vcr_diag_dir() -> None:
if os.environ.get("PYTEST_XDIST_WORKER"):
return
directory = _vcr_diag_dir()
if not os.path.isdir(directory):
return
try:
names = os.listdir(directory)
except OSError:
return
for name in names:
if name.endswith(".log"):
try:
os.remove(os.path.join(directory, name))
except OSError:
pass
def emit_vcr_diagnostic_log(terminalreporter) -> None:
directory = _vcr_diag_dir()
if not os.path.isdir(directory):
return
try:
files = sorted(f for f in os.listdir(directory) if f.endswith(".log"))
except OSError:
return
if not files:
return
terminalreporter.write_sep("=", "VCR DIAGNOSTIC LOG", bold=True)
terminalreporter.write_line(
f" source dir: {directory} (also archived as a CI artifact)"
)
for name in files:
path = os.path.join(directory, name)
try:
with open(path, "r", encoding="utf-8") as fh:
content = fh.read()
except OSError as exc:
terminalreporter.write_line(
f" [failed to read {name}: {type(exc).__name__}: {exc}]"
)
continue
if not content.strip():
continue
terminalreporter.write_sep("-", name, bold=False)
for line in content.splitlines():
terminalreporter.write_line(line)
terminalreporter.write_sep("=", bold=True)
# Intentionally narrower than ``FILTERED_REQUEST_HEADERS``: AWS SigV4 headers
# carry secrets but their values rotate on every call, so fingerprinting them
# would defeat caching.
@ -91,6 +160,32 @@ VCR_IMAGE_B64_PLACEHOLDER = "dGVzdA=="
VCR_FIXED_MULTIPART_BOUNDARY = "vcr-static-boundary"
def pin_httpx_multipart_boundary(monkeypatch) -> None:
try:
import httpx._multipart as _httpx_multipart
except ImportError:
return
_original_init = _httpx_multipart.MultipartStream.__init__
def _init_with_fixed_boundary(self, data, files, boundary=None, **kwargs):
if boundary is None:
boundary = VCR_FIXED_MULTIPART_BOUNDARY.encode("ascii")
return _original_init(self, data=data, files=files, boundary=boundary, **kwargs)
monkeypatch.setattr(
_httpx_multipart.MultipartStream, "__init__", _init_with_fixed_boundary
)
@pytest.fixture(scope="session", autouse=True)
def _pin_multipart_boundary():
monkeypatch = pytest.MonkeyPatch()
pin_httpx_multipart_boundary(monkeypatch)
yield
monkeypatch.undo()
def _scrub_response(response):
if not isinstance(response, dict):
return response
@ -139,9 +234,17 @@ def _strip_image_b64_payloads(response):
preserves all those checks while shrinking cassettes by ~99%.
"""
if not isinstance(response, dict):
vcr_diag_write_line(
f"[vcr-strip-b64] response is {type(response).__name__!r}, not "
"dict; skipping b64 scrub"
)
return response
body = response.get("body")
if not isinstance(body, dict):
vcr_diag_write_line(
f"[vcr-strip-b64] response['body'] is {type(body).__name__!r}, "
"not dict; skipping b64 scrub"
)
return response
raw = body.get("string")
if raw is None:
@ -151,12 +254,20 @@ def _strip_image_b64_payloads(response):
try:
text = bytes(raw).decode("utf-8")
except UnicodeDecodeError:
vcr_diag_write_line(
"[vcr-strip-b64] response body bytes are not valid UTF-8; "
"skipping b64 scrub"
)
return response
was_bytes = True
elif isinstance(raw, str):
text = raw
was_bytes = False
else:
vcr_diag_write_line(
f"[vcr-strip-b64] response['body']['string'] is "
f"{type(raw).__name__!r}, not bytes/str; skipping b64 scrub"
)
return response
try:
@ -186,6 +297,35 @@ def _before_record_response(response):
return filter_non_2xx_response(_scrub_response(_strip_image_b64_payloads(response)))
def _canonical_body(request) -> tuple[bytes, str]:
pre_type = type(getattr(request, "body", None)).__name__
_materialize_iterable_body(request)
body = getattr(request, "body", None)
if body is None:
return b"", pre_type
if isinstance(body, bytes):
return body, pre_type
if isinstance(body, bytearray):
return bytes(body), pre_type
if isinstance(body, str):
return body.encode("utf-8"), pre_type
if isinstance(body, (dict, list)):
try:
return (
json.dumps(body, sort_keys=True, separators=(",", ":")).encode("utf-8"),
pre_type,
)
except (TypeError, ValueError):
pass
method = getattr(request, "method", "?")
uri = getattr(request, "uri", getattr(request, "url", "?"))
vcr_diag_write_line(
f"[vcr-canonical-body] FALLBACK: {method} {uri} body type "
f"{type(body).__name__!r} not coerced to bytes; comparing as b''"
)
return b"", pre_type
def _safe_body_matcher(r1, r2) -> None:
"""Compare request bodies as bytes; never invokes ``json.loads``.
@ -195,27 +335,47 @@ def _safe_body_matcher(r1, r2) -> None:
This matcher is strictly more conservative the only equivalence
it gives up vs. the default is "JSON key order doesn't matter".
"""
body1 = getattr(r1, "body", None)
body2 = getattr(r2, "body", None)
body1, pre1 = _canonical_body(r1)
body2, pre2 = _canonical_body(r2)
if body1 == body2:
return
def _to_bytes(b):
if b is None:
return b""
if isinstance(b, bytes):
return b
if isinstance(b, str):
return b.encode("utf-8")
return None
n1 = _to_bytes(body1)
n2 = _to_bytes(body2)
if n1 is not None and n2 is not None and n1 == n2:
return
_emit_body_mismatch_diagnostic(r1, r2, body1, body2, pre1, pre2)
raise AssertionError("request bodies differ")
def _emit_body_mismatch_diagnostic(r1, r2, body1, body2, pre1, pre2) -> None:
def _describe(label, asbytes, pre_type):
return (
f" {label}: pre_canonical_type={pre_type!r} length={len(asbytes)} "
f"sha256={hashlib.sha256(asbytes).hexdigest()} "
f"preview={asbytes[:120]!r}"
)
method_a = getattr(r1, "method", "?")
method_b = getattr(r2, "method", "?")
url_a = getattr(r1, "uri", getattr(r1, "url", "?"))
url_b = getattr(r2, "uri", getattr(r2, "url", "?"))
lines = [
"[vcr-safe-body-matcher] request body mismatch",
f" request[a]: {method_a} {url_a}",
f" request[b]: {method_b} {url_b}",
_describe("body[a]", body1, pre1),
_describe("body[b]", body2, pre2),
]
if body1 != body2:
offset = next(
(i for i in range(min(len(body1), len(body2))) if body1[i] != body2[i]),
min(len(body1), len(body2)),
)
start = max(0, offset - 100)
end_a = min(len(body1), offset + 100)
end_b = min(len(body2), offset + 100)
lines.append(f" first divergent byte offset: {offset}")
lines.append(f" window[a] @ {start}..{end_a}: {body1[start:end_a]!r}")
lines.append(f" window[b] @ {start}..{end_b}: {body2[start:end_b]!r}")
vcr_diag_write_line("\n".join(lines))
def _iter_header_values(headers, name: str):
if headers is None:
return
@ -271,6 +431,13 @@ def _compute_key_fingerprint(request) -> str:
stable = _stable_key_value(header_name, text)
parts.append(f"{header_name}={stable}")
if not parts:
method = getattr(request, "method", "?")
uri = getattr(request, "uri", getattr(request, "url", "?"))
vcr_diag_write_line(
f"[vcr-key-fingerprint] no API key header found on {method} "
f"{uri}; falling back to 'no-key'. If this request should have "
"carried auth, something earlier in the pipeline stripped it."
)
return "no-key"
digest = hashlib.sha256("\n".join(parts).encode("utf-8")).hexdigest()
return digest[:16]
@ -360,6 +527,13 @@ def _normalize_multipart_boundary(request) -> None:
elif isinstance(body, str):
new_body = body.replace(current_boundary, VCR_FIXED_MULTIPART_BOUNDARY)
else:
vcr_diag_write_line(
f"[vcr-multipart-normalize] body normalization SKIPPED: "
f"body type {type(body).__name__!r} is not bytes/bytearray/str. "
f"content-type={content_type_value!r}. "
f"Recorded body will retain the random boundary substring "
f"and the safe_body matcher will miss on the next run."
)
return
try:
@ -389,6 +563,7 @@ def _before_record_request(request):
headers = getattr(request, "headers", None)
if headers is None:
return request
_materialize_iterable_body(request)
if not any(_iter_header_values(headers, KEY_FINGERPRINT_HEADER)):
fingerprint = _compute_key_fingerprint(request)
try:
@ -400,6 +575,56 @@ def _before_record_request(request):
return request
def _materialize_iterable_body(request) -> None:
body = getattr(request, "body", None)
if body is None or isinstance(body, (bytes, bytearray, str)):
return
if not hasattr(body, "__next__"):
return
try:
chunks = list(body)
except TypeError:
return
out = _coalesce_chunks_to_bytes(chunks)
if out is None:
method = getattr(request, "method", "?")
uri = getattr(request, "uri", getattr(request, "url", "?"))
first_type = type(chunks[0]).__name__ if chunks else "empty"
vcr_diag_write_line(
f"[vcr-materialize] FALLBACK: {method} {uri} chunk type "
f"{first_type!r} not coerced to bytes; storing b''"
)
out = b""
try:
request.body = out
except (AttributeError, TypeError):
pass
for attr in ("_was_iter", "_was_file"):
try:
setattr(request, attr, False)
except (AttributeError, TypeError):
pass
def _coalesce_chunks_to_bytes(chunks):
if not chunks:
return b""
first = chunks[0]
try:
if isinstance(first, int):
return bytes(chunks)
if isinstance(first, (bytes, bytearray)):
return b"".join(c if isinstance(c, bytes) else bytes(c) for c in chunks)
if isinstance(first, str):
return "".join(chunks).encode("utf-8")
except (TypeError, ValueError):
return None
return None
def _key_fingerprint_matcher(r1, r2) -> None:
def _fp(req):
for value in _iter_header_values(
@ -410,7 +635,17 @@ def _key_fingerprint_matcher(r1, r2) -> None:
return value if isinstance(value, str) else str(value)
return "no-key"
if _fp(r1) != _fp(r2):
fp1, fp2 = _fp(r1), _fp(r2)
if fp1 != fp2:
method_a = getattr(r1, "method", "?")
method_b = getattr(r2, "method", "?")
url_a = getattr(r1, "uri", getattr(r1, "url", "?"))
url_b = getattr(r2, "uri", getattr(r2, "url", "?"))
vcr_diag_write_line(
"[vcr-key-fingerprint-matcher] API key fingerprints differ\n"
f" request[a]: {method_a} {url_a} fingerprint={fp1!r}\n"
f" request[b]: {method_b} {url_b} fingerprint={fp2!r}"
)
raise AssertionError("API key fingerprints differ")

View file

@ -159,9 +159,20 @@ def make_redis_persister(
raise CassetteNotFoundError() from exc
if data is None:
raise CassetteNotFoundError()
if isinstance(data, bytes):
data = data.decode("utf-8")
return deserialize(data, serializer)
try:
if isinstance(data, bytes):
data = data.decode("utf-8")
return deserialize(data, serializer)
except Exception as exc:
_record_cache_failure("load", exc)
msg = (
f"VCR redis load failed for {cassette_path}; cached "
f"payload is corrupt, treating as cache miss: "
f"{type(exc).__name__}: {exc}"
)
_log.warning(msg)
warnings.warn(msg, VCRCassetteCacheWarning, stacklevel=2)
raise CassetteNotFoundError() from exc
@staticmethod
def save_cassette(cassette_path, cassette_dict, serializer):

View file

@ -5,14 +5,17 @@ import pytest
sys.path.insert(0, os.path.abspath("../.."))
from tests._vcr_conftest_common import ( # noqa: E402
from tests._vcr_conftest_common import ( # noqa: E402,F401
VerboseReporterState,
_pin_multipart_boundary,
apply_vcr_auto_marker_to_items,
emit_cassette_cache_session_banner,
emit_vcr_classification_summary,
emit_vcr_diagnostic_log,
install_live_call_probe,
record_vcr_outcome,
register_persister_if_enabled,
reset_vcr_diag_dir,
vcr_config_dict,
)
@ -44,6 +47,7 @@ def _vcr_outcome_gate(request, vcr):
def pytest_configure(config):
_verbose_state.remember_pluginmanager(config)
reset_vcr_diag_dir()
def pytest_runtest_logreport(report):
@ -57,3 +61,4 @@ def pytest_collection_modifyitems(config, items):
def pytest_terminal_summary(terminalreporter, exitstatus, config):
emit_cassette_cache_session_banner(terminalreporter)
emit_vcr_classification_summary(terminalreporter)
emit_vcr_diagnostic_log(terminalreporter)

View file

@ -23,12 +23,21 @@ pwd = os.path.dirname(os.path.realpath(__file__))
print(pwd)
file_path = os.path.join(pwd, "gettysburg.wav")
audio_file = open(file_path, "rb")
file2_path = os.path.join(pwd, "eagle.wav")
audio_file2 = open(file2_path, "rb")
with open(file_path, "rb") as _f:
_GETTYSBURG_BYTES = _f.read()
with open(file2_path, "rb") as _f:
_EAGLE_BYTES = _f.read()
def _audio_file():
return ("gettysburg.wav", _GETTYSBURG_BYTES, "audio/wav")
def _audio_file2():
return ("eagle.wav", _EAGLE_BYTES, "audio/wav")
load_dotenv()
@ -44,7 +53,7 @@ async def _run_transcription(
):
transcript = await litellm.atranscription(
model=model,
file=audio_file,
file=_audio_file(),
api_key=api_key,
api_base=api_base,
response_format=response_format,
@ -101,7 +110,7 @@ async def test_transcription_caching():
response_1 = await litellm.atranscription(
model="whisper-1",
file=audio_file,
file=_audio_file(),
)
await asyncio.sleep(5)
@ -110,7 +119,7 @@ async def test_transcription_caching():
response_2 = await litellm.atranscription(
model="whisper-1",
file=audio_file,
file=_audio_file(),
)
print("response_1", response_1)
@ -122,7 +131,7 @@ async def test_transcription_caching():
response_3 = await litellm.atranscription(
model="whisper-1",
file=audio_file2,
file=_audio_file2(),
)
print("response_3", response_3)
print("response3 hidden params", response_3._hidden_params)
@ -146,7 +155,7 @@ async def test_whisper_log_pre_call():
with patch.object(custom_logger, "log_pre_api_call") as mock_log_pre_call:
await litellm.atranscription(
model="whisper-1",
file=audio_file,
file=_audio_file(),
)
mock_log_pre_call.assert_called_once()
@ -165,7 +174,7 @@ async def test_whisper_log_pre_call():
with patch.object(custom_logger, "log_pre_api_call") as mock_log_pre_call:
await litellm.atranscription(
model="whisper-1",
file=audio_file,
file=_audio_file(),
)
mock_log_pre_call.assert_called_once()
@ -177,7 +186,7 @@ async def test_gpt_4o_transcribe():
from unittest.mock import patch, MagicMock
await litellm.atranscription(
model="openai/gpt-4o-transcribe", file=audio_file, response_format="json"
model="openai/gpt-4o-transcribe", file=_audio_file(), response_format="json"
)
@ -187,7 +196,9 @@ async def test_gpt_4o_transcribe_model_mapping():
# Test GPT-4o mini transcribe
response = await litellm.atranscription(
model="openai/gpt-4o-mini-transcribe", file=audio_file, response_format="json"
model="openai/gpt-4o-mini-transcribe",
file=_audio_file(),
response_format="json",
)
# Check that the response contains the correct model in hidden params
@ -198,7 +209,7 @@ async def test_gpt_4o_transcribe_model_mapping():
# Test GPT-4o transcribe
response2 = await litellm.atranscription(
model="openai/gpt-4o-transcribe", file=audio_file, response_format="json"
model="openai/gpt-4o-transcribe", file=_audio_file(), response_format="json"
)
# Check that the response contains the correct model in hidden params
@ -209,7 +220,7 @@ async def test_gpt_4o_transcribe_model_mapping():
# Test traditional whisper-1 still works
response3 = await litellm.atranscription(
model="openai/whisper-1", file=audio_file, response_format="json"
model="openai/whisper-1", file=_audio_file(), response_format="json"
)
# Check that the response contains the correct model in hidden params
@ -262,7 +273,7 @@ async def test_azure_transcribe_model_mapping():
# Make the transcription call
response = await litellm.atranscription(
model="azure/whisper-1",
file=audio_file,
file=_audio_file(),
response_format="json",
api_key="test-api-key",
api_base="https://my-endpoint-europe-berri-992.openai.azure.com/",

View file

@ -16,14 +16,17 @@ sys.path.insert(
) # Adds the parent directory to the system path
import litellm
from tests._vcr_conftest_common import ( # noqa: E402
from tests._vcr_conftest_common import ( # noqa: E402,F401
VerboseReporterState,
_pin_multipart_boundary,
apply_vcr_auto_marker_to_items,
emit_cassette_cache_session_banner,
emit_vcr_classification_summary,
emit_vcr_diagnostic_log,
install_live_call_probe,
record_vcr_outcome,
register_persister_if_enabled,
reset_vcr_diag_dir,
vcr_config_dict,
)
@ -55,6 +58,7 @@ def _vcr_outcome_gate(request, vcr):
def pytest_configure(config):
_verbose_state.remember_pluginmanager(config)
reset_vcr_diag_dir()
def pytest_runtest_logreport(report):
@ -160,3 +164,4 @@ def pytest_collection_modifyitems(config, items):
def pytest_terminal_summary(terminalreporter, exitstatus, config):
emit_cassette_cache_session_banner(terminalreporter)
emit_vcr_classification_summary(terminalreporter)
emit_vcr_diagnostic_log(terminalreporter)

View file

@ -9,14 +9,17 @@ sys.path.insert(
) # Adds the parent directory to the system path
import litellm # noqa: E402,F401
from tests._vcr_conftest_common import ( # noqa: E402
from tests._vcr_conftest_common import ( # noqa: E402,F401
VerboseReporterState,
_pin_multipart_boundary,
apply_vcr_auto_marker_to_items,
emit_cassette_cache_session_banner,
emit_vcr_classification_summary,
emit_vcr_diagnostic_log,
install_live_call_probe,
record_vcr_outcome,
register_persister_if_enabled,
reset_vcr_diag_dir,
vcr_config_dict,
)
@ -58,6 +61,7 @@ def _vcr_outcome_gate(request, vcr):
def pytest_configure(config):
_verbose_state.remember_pluginmanager(config)
reset_vcr_diag_dir()
def pytest_runtest_logreport(report):
@ -71,3 +75,4 @@ def pytest_collection_modifyitems(config, items):
def pytest_terminal_summary(terminalreporter, exitstatus, config):
emit_cassette_cache_session_banner(terminalreporter)
emit_vcr_classification_summary(terminalreporter)
emit_vcr_diagnostic_log(terminalreporter)

View file

@ -103,12 +103,6 @@ class BaseLLMImageEditTest(ABC):
pwd = os.path.dirname(os.path.realpath(__file__))
# Image fixtures must be regenerated per access — module-level
# ``open(...)`` handles get consumed after a single multipart upload, leaving
# subsequent tests in the same process to send empty bodies. That non-determinism
# (a) blows the recorded cassette past ``MAX_EPISODES_PER_CASSETTE`` so the
# persister refuses to save (see ``tests/_vcr_redis_persister.py``), and
# (b) re-bills the live image edit endpoint on every CI run.
def _read_image_bytes(filename: str) -> bytes:
with open(os.path.join(pwd, filename), "rb") as f:
return f.read()
@ -119,32 +113,20 @@ _LITELLM_SITE_BYTES = _read_image_bytes("litellm_site.png")
def _make_test_images() -> list:
"""Return a fresh pair of image streams seeded with the fixture bytes.
return [_ISHAAN_GITHUB_BYTES, _LITELLM_SITE_BYTES]
Use this everywhere you'd previously have used the module-level
``TEST_IMAGES``. Each call returns brand new ``BytesIO`` objects whose
file pointers start at 0, so multipart uploads encode the full image
bytes on every test invocation. Parametrized and ``flaky``-retried
test methods call ``get_base_image_edit_call_args`` once per
invocation, so a fresh stream per call is sufficient the factory
must not auto-rewind on EOF or the SDK's multipart writer will read
the same bytes forever (worker OOM).
"""
def _make_single_test_image() -> bytes:
return _ISHAAN_GITHUB_BYTES
def get_test_images_as_bytesio():
return [
BytesIO(_ISHAAN_GITHUB_BYTES),
BytesIO(_LITELLM_SITE_BYTES),
]
def _make_single_test_image() -> BytesIO:
return BytesIO(_ISHAAN_GITHUB_BYTES)
def get_test_images_as_bytesio():
"""Helper function to get test images as BytesIO objects"""
return _make_test_images()
class TestOpenAIImageEditGPTImage1(BaseLLMImageEditTest):
"""
Concrete implementation of BaseLLMImageEditTest for OpenAI image edits.
@ -710,10 +692,9 @@ async def test_multiple_image_edit_with_different_formats():
try:
prompt = "Create a cohesive artistic style across all images"
# Test with mixed BytesIO and file objects
mixed_images = [
_make_single_test_image(), # File object
get_test_images_as_bytesio()[1], # BytesIO object
_make_single_test_image(),
get_test_images_as_bytesio()[1],
]
result = await aimage_edit(

View file

@ -12,14 +12,17 @@ sys.path.insert(
) # Adds the parent directory to the system path
import litellm # noqa: E402,F401
from tests._vcr_conftest_common import ( # noqa: E402
from tests._vcr_conftest_common import ( # noqa: E402,F401
VerboseReporterState,
_pin_multipart_boundary,
apply_vcr_auto_marker_to_items,
emit_cassette_cache_session_banner,
emit_vcr_classification_summary,
emit_vcr_diagnostic_log,
install_live_call_probe,
record_vcr_outcome,
register_persister_if_enabled,
reset_vcr_diag_dir,
vcr_config_dict,
)
@ -86,6 +89,7 @@ def _vcr_outcome_gate(request, vcr):
def pytest_configure(config):
_verbose_state.remember_pluginmanager(config)
reset_vcr_diag_dir()
def pytest_runtest_logreport(report):
@ -116,3 +120,4 @@ def pytest_collection_modifyitems(config, items):
def pytest_terminal_summary(terminalreporter, exitstatus, config):
emit_cassette_cache_session_banner(terminalreporter)
emit_vcr_classification_summary(terminalreporter)
emit_vcr_diagnostic_log(terminalreporter)

View file

@ -13,14 +13,17 @@ sys.path.insert(
import litellm # noqa: E402
from tests._vcr_conftest_common import ( # noqa: E402
from tests._vcr_conftest_common import ( # noqa: E402,F401
VerboseReporterState,
_pin_multipart_boundary,
apply_vcr_auto_marker_to_items,
emit_cassette_cache_session_banner,
emit_vcr_classification_summary,
emit_vcr_diagnostic_log,
install_live_call_probe,
record_vcr_outcome,
register_persister_if_enabled,
reset_vcr_diag_dir,
vcr_config_dict,
)
@ -52,6 +55,7 @@ def _vcr_outcome_gate(request, vcr):
def pytest_configure(config):
_verbose_state.remember_pluginmanager(config)
reset_vcr_diag_dir()
def pytest_runtest_logreport(report):
@ -116,3 +120,4 @@ def pytest_collection_modifyitems(config, items):
def pytest_terminal_summary(terminalreporter, exitstatus, config):
emit_cassette_cache_session_banner(terminalreporter)
emit_vcr_classification_summary(terminalreporter)
emit_vcr_diagnostic_log(terminalreporter)

View file

@ -18,14 +18,17 @@ sys.path.insert(
import litellm # noqa: E402
from tests._vcr_conftest_common import ( # noqa: E402
from tests._vcr_conftest_common import ( # noqa: E402,F401
VerboseReporterState,
_pin_multipart_boundary,
apply_vcr_auto_marker_to_items,
emit_cassette_cache_session_banner,
emit_vcr_classification_summary,
emit_vcr_diagnostic_log,
install_live_call_probe,
record_vcr_outcome,
register_persister_if_enabled,
reset_vcr_diag_dir,
vcr_config_dict,
)
@ -73,6 +76,7 @@ def _vcr_outcome_gate(request, vcr):
def pytest_configure(config):
_verbose_state.remember_pluginmanager(config)
reset_vcr_diag_dir()
def pytest_runtest_logreport(report):
@ -82,6 +86,7 @@ def pytest_runtest_logreport(report):
def pytest_terminal_summary(terminalreporter, exitstatus, config):
emit_cassette_cache_session_banner(terminalreporter)
emit_vcr_classification_summary(terminalreporter)
emit_vcr_diagnostic_log(terminalreporter)
# ---------------------------------------------------------------------------

View file

@ -101,7 +101,9 @@ async def test_openai_realtime_direct_call_no_intent():
try:
await litellm._arealtime(
model="openai/gpt-4o-realtime-preview",
# OpenAI shut down the gpt-4o-realtime-preview family (incl. the
# undated alias) on 2026-05-07; gpt-realtime is the GA successor.
model="openai/gpt-realtime",
websocket=websocket_client,
api_key=os.environ.get("OPENAI_API_KEY"),
timeout=60,
@ -249,14 +251,16 @@ async def test_openai_realtime_direct_call_with_intent():
websocket_client = RealTimeWebSocketClient()
caught_exception = None
# OpenAI shut down the gpt-4o-realtime-preview family (incl. the undated
# alias) on 2026-05-07; gpt-realtime is the GA successor.
query_params: RealtimeQueryParams = {
"model": "openai/gpt-4o-realtime-preview",
"model": "openai/gpt-realtime",
"intent": "chat",
}
try:
await litellm._arealtime(
model="openai/gpt-4o-realtime-preview",
model="openai/gpt-realtime",
websocket=websocket_client,
api_key=os.environ.get("OPENAI_API_KEY"),
query_params=query_params,

View file

@ -21,7 +21,10 @@ class TestOpenAIRealtime(BaseRealtimeTest):
"""
def get_model(self) -> str:
return "gpt-4o-realtime-preview"
# OpenAI shut down the entire gpt-4o-realtime-preview family
# (including the undated alias) on 2026-05-07. gpt-realtime is the
# current GA realtime model.
return "gpt-realtime"
def get_api_key_env_var(self) -> str:
return "OPENAI_API_KEY"

View file

@ -26,9 +26,7 @@ from litellm.litellm_core_utils.realtime_streaming import RealTimeStreaming
from litellm.types.guardrails import GuardrailEventHooks
OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY")
OPENAI_REALTIME_URL = (
"wss://api.openai.com/v1/realtime?model=gpt-4o-realtime-preview-2024-12-17"
)
OPENAI_REALTIME_URL = "wss://api.openai.com/v1/realtime?model=gpt-realtime"
pytestmark = pytest.mark.skipif(
not OPENAI_API_KEY,
@ -192,10 +190,35 @@ async def test_text_message_blocked_by_guardrail_no_ai_response():
len(transcript_deltas) >= 1
), f"Expected guardrail message in transcript delta, got: {event_types}"
# 3. No *real* AI response should have been generated.
# The guardrail may produce its own response (e.g. "Content blocked: ...")
# via response.cancel + conversation.item.create + response.create.
# We allow the guardrail's own block message but NOT original AI content.
# 3. No *real* AI response to the blocked content should have been
# generated. The original user message is blocked BEFORE it is
# forwarded to OpenAI, so the only thing the model ever sees is the
# guardrail's "say exactly: <block message>" prompt
# (see realtime_streaming.py). Two safe outcomes are possible:
# - the model voices the block message verbatim (older realtime
# snapshots did this -> text contains "blocked"), or
# - the model declines to repeat it (gpt-realtime tends to refuse
# verbatim-repeat instructions, e.g. "I'm sorry, but I can't
# repeat that message.").
# Both mean the blocked prompt itself was never answered, so we
# accept either. The hard invariant is that the blocked phrase must
# never leak into AI output, and the model must not have produced a
# normal answer to the user (which would have neither a block nor a
# refusal marker).
safe_markers = (
"block",
"guardrail",
"content filter",
"policy",
"can't repeat",
"cannot repeat",
"won't repeat",
"can't assist",
"can't help",
"unable to",
"i'm sorry",
"i am sorry",
)
done_events = [e for e in client_events if e.get("type") == "response.done"]
for done in done_events:
output = done.get("response", {}).get("output", [])
@ -205,11 +228,19 @@ async def test_text_message_blocked_by_guardrail_no_ai_response():
for c in item.get("content", [])
]
real_ai_text = " ".join(ai_texts).strip()
# Allow guardrail-generated block messages (contain "Content blocked" or "blocked")
if real_ai_text:
assert (
"blocked" in real_ai_text.lower()
or "guardrail" in real_ai_text.lower()
BLOCKED_PHRASE not in real_ai_text
), f"Blocked phrase leaked into AI response: {real_ai_text!r}"
normalized_ai_text = (
real_ai_text.lower()
.replace("\u2019", "'")
.replace("\u2018", "'")
.replace("\u201c", '"')
.replace("\u201d", '"')
)
assert any(
marker in normalized_ai_text for marker in safe_markers
), f"AI responded with non-guardrail content even though message was blocked: {real_ai_text!r}"
finally:

View file

@ -176,3 +176,113 @@ def test_completion_cost_deepseek():
pass
except Exception as e:
pytest.fail(f"Error occurred: {e}")
def test_deepseek_fill_reasoning_content_multiturn():
"""
Unit test for _fill_reasoning_content.
Reproduces issue #28045: DeepSeek thinking mode fails in multi-turn conversations
because reasoning_content is not passed back to the API.
"""
from litellm.llms.deepseek.chat.transformation import DeepSeekChatConfig
config = DeepSeekChatConfig()
# Case 1: assistant message already has reasoning_content — should be left as-is
messages_with_rc = [
{"role": "user", "content": "Hello"},
{"role": "assistant", "content": "Hi", "reasoning_content": "I thought about it"},
{"role": "user", "content": "Follow up"},
]
result = config._fill_reasoning_content(messages_with_rc)
assert result[1]["reasoning_content"] == "I thought about it"
# Case 2: assistant message has reasoning_content in provider_specific_fields — should be promoted
messages_with_psf = [
{"role": "user", "content": "Hello"},
{
"role": "assistant",
"content": "Hi",
"provider_specific_fields": {"reasoning_content": "stored thinking"},
},
{"role": "user", "content": "Follow up"},
]
result = config._fill_reasoning_content(messages_with_psf)
assert result[1]["reasoning_content"] == "stored thinking"
# Should be removed from provider_specific_fields to avoid duplication
assert "reasoning_content" not in result[1].get("provider_specific_fields", {})
# Case 3: assistant message has no reasoning_content anywhere — should inject placeholder
messages_no_rc = [
{"role": "user", "content": "Hello"},
{"role": "assistant", "content": "Hi"},
{"role": "user", "content": "Follow up"},
]
result = config._fill_reasoning_content(messages_no_rc)
assert result[1]["reasoning_content"] == " "
# Case 4: non-assistant messages should never be touched
messages_user_only = [
{"role": "user", "content": "Hello"},
{"role": "system", "content": "You are helpful"},
]
result = config._fill_reasoning_content(messages_user_only)
assert "reasoning_content" not in result[0]
assert "reasoning_content" not in result[1]
def test_deepseek_fill_reasoning_content_guard_in_transform_request():
"""
_fill_reasoning_content must only run when BOTH conditions are true:
1. supports_reasoning() is True for the model
2. thinking mode is explicitly enabled in optional_params ({"type": "enabled"})
This prevents spurious injection on models like deepseek-v3.2 that support
thinking as opt-in but not always-on. Addresses oss-pr-review-agent feedback
on PR #28057.
"""
from litellm.llms.deepseek.chat.transformation import DeepSeekChatConfig
config = DeepSeekChatConfig()
messages = [
{"role": "user", "content": "Hello"},
{"role": "assistant", "content": "Hi"},
{"role": "user", "content": "Follow up"},
]
# Case 1: reasoning model + thinking enabled -> injection should happen
result = config.transform_request(
model="deepseek-reasoner",
messages=messages,
optional_params={"thinking": {"type": "enabled"}},
litellm_params={},
headers={},
)
assert result["messages"][1].get("reasoning_content") == " ", (
"reasoning_content should be injected when thinking is enabled"
)
# Case 2: reasoning model + thinking NOT in optional_params -> no injection
result = config.transform_request(
model="deepseek-reasoner",
messages=messages,
optional_params={},
litellm_params={},
headers={},
)
assert "reasoning_content" not in result["messages"][1], (
"reasoning_content should not be injected when thinking is not enabled"
)
# Case 3: non-reasoning model + thinking enabled -> no injection
result = config.transform_request(
model="deepseek-chat",
messages=messages,
optional_params={"thinking": {"type": "enabled"}},
litellm_params={},
headers={},
)
assert "reasoning_content" not in result["messages"][1], (
"reasoning_content should not be injected for non-reasoning models"
)

View file

@ -262,3 +262,44 @@ class TestNvidiaNim(BaseLLMRerankTest):
def get_expected_cost(self) -> float:
"""Nvidia NIM rerank models are free (cost = 0.0)"""
return 0.0
@pytest.mark.asyncio()
@pytest.mark.parametrize("sync_mode", [True, False])
async def test_basic_rerank(self, sync_mode, monkeypatch):
"""
Override the base live rerank test with a mocked HTTP layer.
NVIDIA reached end-of-life for the hosted
nvidia/llama-3.2-nv-rerankqa-1b-v2 rerank API on 2026-05-18 and
published no replacement model, so a live call now returns HTTP 410
("Gone"). NVIDIA's hosted catalog rotates on a schedule, so pointing
at another live model would only defer the same failure. Mock the
transport instead (same pattern as
test_nvidia_nim_rerank_ranking_endpoint above) so the request/response
transformation and cost calculation stay covered offline.
"""
monkeypatch.setenv("NVIDIA_NIM_API_KEY", "fake-api-key")
mock_response = MagicMock()
mock_response.status_code = 200
mock_response.headers = {}
mock_response.text = ""
mock_response.json.return_value = {
"rankings": [
{"index": 0, "logit": 0.95},
{"index": 1, "logit": 0.75},
],
"usage": {"total_tokens": 7},
}
with (
patch(
"litellm.llms.custom_httpx.http_handler.HTTPHandler.post",
return_value=mock_response,
),
patch(
"litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post",
return_value=mock_response,
),
):
await super().test_basic_rerank(sync_mode=sync_mode)

View file

@ -22,14 +22,17 @@ sys.path.insert(
) # Adds the parent directory to the system path
import litellm
from tests._vcr_conftest_common import ( # noqa: E402
from tests._vcr_conftest_common import ( # noqa: E402,F401
VerboseReporterState,
_pin_multipart_boundary,
apply_vcr_auto_marker_to_items,
emit_cassette_cache_session_banner,
emit_vcr_classification_summary,
emit_vcr_diagnostic_log,
install_live_call_probe,
record_vcr_outcome,
register_persister_if_enabled,
reset_vcr_diag_dir,
vcr_config_dict,
)
@ -84,6 +87,7 @@ def _vcr_outcome_gate(request, vcr):
def pytest_configure(config):
_verbose_state.remember_pluginmanager(config)
reset_vcr_diag_dir()
def pytest_runtest_logreport(report):
@ -93,6 +97,7 @@ def pytest_runtest_logreport(report):
def pytest_terminal_summary(terminalreporter, exitstatus, config):
emit_cassette_cache_session_banner(terminalreporter)
emit_vcr_classification_summary(terminalreporter)
emit_vcr_diagnostic_log(terminalreporter)
# ---------------------------------------------------------------------------

View file

@ -25,6 +25,7 @@ from unittest.mock import AsyncMock, patch, MagicMock
from litellm.caching.caching_handler import (
LLMCachingHandler,
CachingHandlerResponse,
_is_chat_completion_cached_dict,
_should_defer_streaming_cache_hit_callbacks,
)
from litellm.caching.caching import LiteLLMCacheType
@ -40,6 +41,7 @@ from litellm.types.utils import (
from litellm.types.llms.openai import ResponsesAPIResponse
from datetime import timedelta, datetime
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
from litellm._logging import verbose_logger
import logging
@ -1072,6 +1074,70 @@ def test_convert_cached_streaming_responses_result_to_iterator():
)
def test_is_chat_completion_cached_dict():
assert _is_chat_completion_cached_dict(
{"id": "chatcmpl-abc", "object": "chat.completion", "choices": []}
)
assert _is_chat_completion_cached_dict(
{"id": "other", "object": "chat.completion.chunk", "choices": []}
)
assert not _is_chat_completion_cached_dict(
{"id": "resp_abc", "object": "response", "output": []}
)
def test_convert_cached_aresponses_bridge_chat_completion_stream():
"""
openai/responses chat-completions bridge caches ModelResponse JSON on aresponses
cache keys; replay must not call ResponsesAPIResponse(**chatcmpl_dict).
"""
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-5.4",
messages=[],
function_id=str(uuid.uuid4()),
stream=True,
start_time=datetime.now(),
)
cached_result = {
"id": "chatcmpl-bridge-cache-test",
"object": "chat.completion",
"created": int(time.time()),
"model": "gpt-5.4",
"choices": [
{
"index": 0,
"message": {"role": "assistant", "content": "Hi!"},
"finish_reason": "stop",
}
],
"usage": {
"prompt_tokens": 7,
"completion_tokens": 11,
"total_tokens": 18,
},
}
result = caching_handler._convert_cached_result_to_model_response(
cached_result=cached_result,
call_type=CallTypes.aresponses.value,
kwargs={
"model": "gpt-5.4",
"stream": True,
"messages": [{"role": "user", "content": "hi"}],
},
logging_obj=logging_obj,
model="gpt-5.4",
args=(),
)
assert isinstance(result, CustomStreamWrapper)
def test_convert_cached_streaming_reasoning_result_to_iterator():
caching_handler = LLMCachingHandler(
original_function=responses, request_kwargs={}, start_time=datetime.now()

View file

@ -19,14 +19,17 @@ sys.path.insert(
) # Adds the parent directory to the system path
import litellm
from tests._vcr_conftest_common import ( # noqa: E402
from tests._vcr_conftest_common import ( # noqa: E402,F401
VerboseReporterState,
_pin_multipart_boundary,
apply_vcr_auto_marker_to_items,
emit_cassette_cache_session_banner,
emit_vcr_classification_summary,
emit_vcr_diagnostic_log,
install_live_call_probe,
record_vcr_outcome,
register_persister_if_enabled,
reset_vcr_diag_dir,
vcr_config_dict,
)
@ -79,6 +82,7 @@ def _vcr_outcome_gate(request, vcr):
def pytest_configure(config):
_verbose_state.remember_pluginmanager(config)
reset_vcr_diag_dir()
def pytest_runtest_logreport(report):
@ -229,3 +233,4 @@ def pytest_collection_modifyitems(config, items):
def pytest_terminal_summary(terminalreporter, exitstatus, config):
emit_cassette_cache_session_banner(terminalreporter)
emit_vcr_classification_summary(terminalreporter)
emit_vcr_diagnostic_log(terminalreporter)

View file

@ -12,14 +12,17 @@ import pytest
sys.path.insert(0, os.path.abspath("../.."))
from tests._vcr_conftest_common import ( # noqa: E402
from tests._vcr_conftest_common import ( # noqa: E402,F401
VerboseReporterState,
_pin_multipart_boundary,
apply_vcr_auto_marker_to_items,
emit_cassette_cache_session_banner,
emit_vcr_classification_summary,
emit_vcr_diagnostic_log,
install_live_call_probe,
record_vcr_outcome,
register_persister_if_enabled,
reset_vcr_diag_dir,
vcr_config_dict,
)
@ -51,6 +54,7 @@ def _vcr_outcome_gate(request, vcr):
def pytest_configure(config):
_verbose_state.remember_pluginmanager(config)
reset_vcr_diag_dir()
def pytest_runtest_logreport(report):
@ -64,3 +68,4 @@ def pytest_collection_modifyitems(config, items):
def pytest_terminal_summary(terminalreporter, exitstatus, config):
emit_cassette_cache_session_banner(terminalreporter)
emit_vcr_classification_summary(terminalreporter)
emit_vcr_diagnostic_log(terminalreporter)

View file

@ -610,22 +610,18 @@ def extract_user_budget_metrics(metrics_text: str, user_id: str) -> Dict[str, fl
# Escape user_id for regex pattern matching
escaped_user_id = re.escape(user_id)
# Get remaining budget
remaining_pattern = (
f'litellm_remaining_user_budget_metric{{user="{escaped_user_id}"}} ([0-9.]+)'
)
# Get remaining budget (user_email and user_alias may also be present as labels)
remaining_pattern = rf'litellm_remaining_user_budget_metric{{[^}}]*user="{escaped_user_id}"[^}}]*}} ([0-9.]+)'
remaining_match = re.search(remaining_pattern, metrics_text)
metrics["remaining"] = float(remaining_match.group(1)) if remaining_match else None
# Get total budget
total_pattern = (
f'litellm_user_max_budget_metric{{user="{escaped_user_id}"}} ([0-9.]+)'
)
total_pattern = rf'litellm_user_max_budget_metric{{[^}}]*user="{escaped_user_id}"[^}}]*}} ([0-9.]+)'
total_match = re.search(total_pattern, metrics_text)
metrics["total"] = float(total_match.group(1)) if total_match else None
# Get remaining hours
hours_pattern = f'litellm_user_budget_remaining_hours_metric{{user="{escaped_user_id}"}} ([0-9.]+)'
hours_pattern = rf'litellm_user_budget_remaining_hours_metric{{[^}}]*user="{escaped_user_id}"[^}}]*}} ([0-9.]+)'
hours_match = re.search(hours_pattern, metrics_text)
metrics["remaining_hours"] = float(hours_match.group(1)) if hours_match else None

View file

@ -5,14 +5,17 @@ import pytest
sys.path.insert(0, os.path.abspath("../.."))
from tests._vcr_conftest_common import ( # noqa: E402
from tests._vcr_conftest_common import ( # noqa: E402,F401
VerboseReporterState,
_pin_multipart_boundary,
apply_vcr_auto_marker_to_items,
emit_cassette_cache_session_banner,
emit_vcr_classification_summary,
emit_vcr_diagnostic_log,
install_live_call_probe,
record_vcr_outcome,
register_persister_if_enabled,
reset_vcr_diag_dir,
vcr_config_dict,
)
@ -56,6 +59,7 @@ def _vcr_outcome_gate(request, vcr):
def pytest_configure(config):
_verbose_state.remember_pluginmanager(config)
reset_vcr_diag_dir()
def pytest_runtest_logreport(report):
@ -71,3 +75,4 @@ def pytest_collection_modifyitems(config, items):
def pytest_terminal_summary(terminalreporter, exitstatus, config):
emit_cassette_cache_session_banner(terminalreporter)
emit_vcr_classification_summary(terminalreporter)
emit_vcr_diagnostic_log(terminalreporter)

View file

@ -12,14 +12,17 @@ sys.path.insert(
) # Adds the parent directory to the system path
import litellm # noqa: E402,F401
from tests._vcr_conftest_common import ( # noqa: E402
from tests._vcr_conftest_common import ( # noqa: E402,F401
VerboseReporterState,
_pin_multipart_boundary,
apply_vcr_auto_marker_to_items,
emit_cassette_cache_session_banner,
emit_vcr_classification_summary,
emit_vcr_diagnostic_log,
install_live_call_probe,
record_vcr_outcome,
register_persister_if_enabled,
reset_vcr_diag_dir,
vcr_config_dict,
)
@ -97,6 +100,7 @@ def _vcr_outcome_gate(request, vcr):
def pytest_configure(config):
_verbose_state.remember_pluginmanager(config)
reset_vcr_diag_dir()
def pytest_runtest_logreport(report):
@ -123,3 +127,4 @@ def pytest_collection_modifyitems(config, items):
def pytest_terminal_summary(terminalreporter, exitstatus, config):
emit_cassette_cache_session_banner(terminalreporter)
emit_vcr_classification_summary(terminalreporter)
emit_vcr_diagnostic_log(terminalreporter)

View file

@ -13,14 +13,17 @@ import pytest
sys.path.insert(0, os.path.abspath("../.."))
from tests._vcr_conftest_common import ( # noqa: E402
from tests._vcr_conftest_common import ( # noqa: E402,F401
VerboseReporterState,
_pin_multipart_boundary,
apply_vcr_auto_marker_to_items,
emit_cassette_cache_session_banner,
emit_vcr_classification_summary,
emit_vcr_diagnostic_log,
install_live_call_probe,
record_vcr_outcome,
register_persister_if_enabled,
reset_vcr_diag_dir,
vcr_config_dict,
)
@ -52,6 +55,7 @@ def _vcr_outcome_gate(request, vcr):
def pytest_configure(config):
_verbose_state.remember_pluginmanager(config)
reset_vcr_diag_dir()
def pytest_runtest_logreport(report):
@ -65,3 +69,4 @@ def pytest_collection_modifyitems(config, items):
def pytest_terminal_summary(terminalreporter, exitstatus, config):
emit_cassette_cache_session_banner(terminalreporter)
emit_vcr_classification_summary(terminalreporter)
emit_vcr_diagnostic_log(terminalreporter)

View file

@ -9,12 +9,155 @@ import pytest
import asyncio
import aiohttp
import os
import re
import dotenv
from collections import Counter
from dotenv import load_dotenv
import pytest
load_dotenv()
# A *leak* is sustained, monotonic growth of one callback TYPE across the whole
# sampling window. A one-time bump that then plateaus is benign pollution from
# other tests sharing this proxy (this suite runs `pytest -n 4` against a single
# proxy container, so other workers legitimately add team/key-scoped callbacks
# while this test sleeps). We therefore sample N times and only flag a type
# whose normalized count never decreases, grows in >=2 distinct intervals, and
# nets >= LEAK_MIN_NET_GROWTH overall.
NUM_SAMPLES = 4
SAMPLE_INTERVAL_SECONDS = 20
LEAK_MIN_NET_GROWTH = 5
LEAK_MIN_GROWING_INTERVALS = 2
# A routing-strategy switch / alerting config is a *known, bounded, one-time*
# registration (CCI diagnostic 2026-05-16: total 85->95 on the first interval
# after switching to latency-based-routing, then flat at 95 for 2.5 min under
# load). We absorb that step by settling before the baseline sample, so only
# growth *after* the deliberate perturbation can count as a leak.
SETTLE_SECONDS = 30
# Strip instance-identity noise so N leaked instances of one class collapse to
# one rising counter instead of N opaque, unrelated-looking strings.
_ADDR_RE = re.compile(r" at 0x[0-9a-fA-F]+")
_OBJ_RE = re.compile(r"<([\w.]+) object")
def _normalize_callback(cb_str: str) -> str:
"""Reduce a callback's str() to a stable type key (drops 0x… addresses)."""
s = _ADDR_RE.sub("", cb_str)
m = _OBJ_RE.search(s)
if m:
return m.group(1).split(".")[-1]
# bound methods: "<bound method Cls.m of <... at 0x..>>" -> "Cls.m"
bm = re.search(r"bound method ([\w.]+)", s)
if bm:
return bm.group(1)
return s.strip()
def _summarize(all_litellm_callbacks) -> Counter:
return Counter(_normalize_callback(str(c)) for c in all_litellm_callbacks)
def _detect_leaks(samples):
"""
samples: list[Counter] taken in time order.
Returns {callback_type: [counts across samples]} for types that grew
monotonically (never decreased), in >=LEAK_MIN_GROWING_INTERVALS intervals,
and netted >=LEAK_MIN_NET_GROWTH overall i.e. a real leak, not a one-shot
step from a parallel test.
"""
leaks = {}
all_types = set().union(*[set(s) for s in samples]) if samples else set()
for t in all_types:
series = [s.get(t, 0) for s in samples]
deltas = [b - a for a, b in zip(series, series[1:])]
net = series[-1] - series[0]
non_decreasing = all(d >= 0 for d in deltas)
growing_intervals = sum(1 for d in deltas if d > 0)
if (
non_decreasing
and net >= LEAK_MIN_NET_GROWTH
and growing_intervals >= LEAK_MIN_GROWING_INTERVALS
):
leaks[t] = series
return leaks
def _terminal_suspects(samples):
"""
Types whose net growth clears the threshold monotonically but is confined
to the *final* interval `growing_intervals == 1` with that one growing
interval being the last. `_detect_leaks`' `>= 2` guard silently passes
these, so a real leak that accumulates entirely in the last sampled window
is indistinguishable from a one-time terminal step *without one more
sample*. Returns the set of such types so the caller can re-confirm.
"""
suspects = set()
all_types = set().union(*[set(s) for s in samples]) if samples else set()
for t in all_types:
series = [s.get(t, 0) for s in samples]
deltas = [b - a for a, b in zip(series, series[1:])]
if not deltas:
continue
net = series[-1] - series[0]
non_decreasing = all(d >= 0 for d in deltas)
growing = [i for i, d in enumerate(deltas) if d > 0]
if (
non_decreasing
and net >= LEAK_MIN_NET_GROWTH
and growing == [len(deltas) - 1]
):
suspects.add(t)
return suspects
async def _detect_leaks_confirmed(session, samples):
"""
`_detect_leaks`, plus a single confirmation sample when growth is confined
to the final interval (see `_terminal_suspects`). A genuine ongoing leak
keeps climbing -> now grows in >= 2 intervals -> flagged; a one-time
terminal registration plateaus -> still 1 growing interval -> ignored.
Returns `(leaks, samples)` (samples may have one extra entry appended).
"""
leaks = _detect_leaks(samples)
if not leaks and _terminal_suspects(samples):
await asyncio.sleep(SAMPLE_INTERVAL_SECONDS)
_, _, all_cb = await get_active_callbacks(session=session)
samples = samples + [_summarize(all_cb)]
leaks = _detect_leaks(samples)
return leaks, samples
def _format_report(samples, leaks) -> str:
lines = ["Callback count per type across samples (time order):"]
all_types = sorted(set().union(*[set(s) for s in samples]))
for t in all_types:
series = [s.get(t, 0) for s in samples]
marker = " <-- LEAK" if t in leaks else ""
lines.append(f" {t}: {series}{marker}")
totals = [sum(s.values()) for s in samples]
lines.append(f"TOTAL callbacks per sample: {totals}")
if leaks:
lines.append(
"Leaking callback types (sustained monotonic growth): "
+ ", ".join(sorted(leaks))
)
return "\n".join(lines)
async def _sample_callbacks(session, num_samples, interval):
"""Take `num_samples` callback snapshots `interval`s apart."""
samples = []
alerts = []
for i in range(num_samples):
if i > 0:
await asyncio.sleep(interval)
num_cb, num_alert, all_cb = await get_active_callbacks(session=session)
samples.append(_summarize(all_cb))
alerts.append(num_alert)
return samples, alerts
async def config_update(session, routing_strategy=None):
url = "http://0.0.0.0:4000/config/update"
@ -97,105 +240,65 @@ async def get_current_routing_strategy(session):
@pytest.mark.asyncio
@pytest.mark.order1
@pytest.mark.flaky(reruns=2, reruns_delay=5)
async def test_check_num_callbacks():
"""
Test 1: num callbacks should NOT increase over time
-> check current callbacks
-> sleep for 30 seconds
-> check current callbacks
-> sleep for 30 seconds
-> check current callbacks
PROD invariant: no callback TYPE should grow without bound over time.
This suite runs `pytest -n 4` against one shared proxy, so the raw count is
noisy other workers legitimately add team/key-scoped callbacks that then
plateau. We settle first, then sample several times, and only fail on
*sustained, monotonic* per-type growth (a genuine leak), naming the type.
"""
from litellm._uuid import uuid
async with aiohttp.ClientSession() as session:
await asyncio.sleep(30)
num_callbacks_1, _, all_litellm_callbacks_1 = await get_active_callbacks(
session=session
)
assert num_callbacks_1 > 0
await asyncio.sleep(30)
# Absorb proxy warmup / in-flight parallel registration before baseline.
await asyncio.sleep(SETTLE_SECONDS)
num_callbacks_2, _, all_litellm_callbacks_2 = await get_active_callbacks(
session=session
samples, _ = await _sample_callbacks(
session, NUM_SAMPLES, SAMPLE_INTERVAL_SECONDS
)
print("all_litellm_callbacks_1", all_litellm_callbacks_1)
assert sum(samples[0].values()) > 0, "expected some callbacks registered"
print(
"diff in callbacks=",
set(all_litellm_callbacks_1) - set(all_litellm_callbacks_2),
)
assert abs(num_callbacks_1 - num_callbacks_2) <= 4
await asyncio.sleep(30)
num_callbacks_3, _, all_litellm_callbacks_3 = await get_active_callbacks(
session=session
)
print(
"diff in callbacks = all_litellm_callbacks3 - all_litellm_callbacks2 ",
set(all_litellm_callbacks_3) - set(all_litellm_callbacks_2),
)
assert abs(num_callbacks_3 - num_callbacks_2) <= 4
leaks, samples = await _detect_leaks_confirmed(session, samples)
report = _format_report(samples, leaks)
print(report)
assert not leaks, f"Callback leak detected.\n{report}"
@pytest.mark.asyncio
@pytest.mark.order2
@pytest.mark.flaky(reruns=2, reruns_delay=5)
async def test_check_num_callbacks_on_lowest_latency():
"""
Test 1: num callbacks should NOT increase over time
-> Update to lowest latency
-> check current callbacks
-> sleep for 30s
-> check current callbacks
-> sleep for 30s
-> check current callbacks
-> update back to original routing-strategy
Same PROD invariant as test_check_num_callbacks, but after switching the
router to latency-based-routing. That switch is a *known, bounded* one-time
registration (it adds the latency strategy handler + Slack alerting); we
settle past it before baselining so only post-switch growth counts as a
leak. Also asserts the alerting count is stable.
"""
from litellm._uuid import uuid
async with aiohttp.ClientSession() as session:
await asyncio.sleep(30)
original_routing_strategy = await get_current_routing_strategy(session=session)
await config_update(session=session, routing_strategy="latency-based-routing")
await asyncio.sleep(30)
try:
# Absorb the deliberate one-time config/update registration step.
await asyncio.sleep(SETTLE_SECONDS)
num_callbacks_1, num_alerts_1, all_litellm_callbacks_1 = (
await get_active_callbacks(session=session)
)
samples, alerts = await _sample_callbacks(
session, NUM_SAMPLES, SAMPLE_INTERVAL_SECONDS
)
await asyncio.sleep(30)
num_callbacks_2, num_alerts_2, all_litellm_callbacks_2 = (
await get_active_callbacks(session=session)
)
print(
"diff in callbacks all_litellm_callbacks_2 - all_litellm_callbacks_1 =",
set(all_litellm_callbacks_2) - set(all_litellm_callbacks_1),
)
assert abs(num_callbacks_1 - num_callbacks_2) <= 4
await asyncio.sleep(30)
num_callbacks_3, num_alerts_3, all_litellm_callbacks_3 = (
await get_active_callbacks(session=session)
)
print(
"diff in callbacks all_litellm_callbacks_3 - all_litellm_callbacks_2 =",
set(all_litellm_callbacks_3) - set(all_litellm_callbacks_2),
)
assert abs(num_callbacks_2 - num_callbacks_3) <= 4
assert num_alerts_1 == num_alerts_2 == num_alerts_3
await config_update(session=session, routing_strategy=original_routing_strategy)
leaks, samples = await _detect_leaks_confirmed(session, samples)
report = _format_report(samples, leaks)
print(report)
assert not leaks, f"Callback leak detected.\n{report}"
assert (
len(set(alerts)) == 1
), f"alerting count changed across samples: {alerts}"
finally:
await config_update(
session=session, routing_strategy=original_routing_strategy
)

View file

@ -232,3 +232,207 @@ def test_combine_usage_handles_none_details():
combined = llm_caching_handler.combine_usage(usage_a, usage_c)
assert combined.prompt_tokens_details is not None
assert combined.prompt_tokens_details.image_count == 1
def test_is_chat_completion_cached_dict():
from litellm.caching.caching_handler import _is_chat_completion_cached_dict
assert _is_chat_completion_cached_dict(
{"id": "chatcmpl-abc", "object": "chat.completion", "choices": []}
)
assert _is_chat_completion_cached_dict(
{"id": "other", "object": "chat.completion.chunk", "choices": []}
)
assert _is_chat_completion_cached_dict(
{"id": "no-object", "choices": [{"index": 0}]}
)
assert not _is_chat_completion_cached_dict(
{"id": "resp_abc", "object": "response", "output": []}
)
def _build_logging_obj(call_type: str, stream: bool):
import uuid as _uuid
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
return LiteLLMLogging(
litellm_call_id=str(datetime.now()),
call_type=call_type,
model="gpt-5.4",
messages=[],
function_id=str(_uuid.uuid4()),
stream=stream,
start_time=datetime.now(),
)
def test_convert_cached_aresponses_bridge_chat_completion_stream():
"""openai/responses chat-completions bridge: streaming cache hit replays as chat stream."""
from litellm import aresponses
from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
from litellm.types.utils import CallTypes
caching_handler = LLMCachingHandler(
original_function=aresponses, request_kwargs={}, start_time=datetime.now()
)
cached_result = {
"id": "chatcmpl-bridge-cache-test",
"object": "chat.completion",
"created": int(time.time()),
"model": "gpt-5.4",
"choices": [
{
"index": 0,
"message": {"role": "assistant", "content": "Hi!"},
"finish_reason": "stop",
}
],
"usage": {"prompt_tokens": 7, "completion_tokens": 11, "total_tokens": 18},
}
result = caching_handler._convert_cached_result_to_model_response(
cached_result=cached_result,
call_type=CallTypes.aresponses.value,
kwargs={
"model": "gpt-5.4",
"stream": True,
"messages": [{"role": "user", "content": "hi"}],
},
logging_obj=_build_logging_obj(CallTypes.aresponses.value, stream=True),
model="gpt-5.4",
args=(),
)
assert isinstance(result, CustomStreamWrapper)
def test_convert_cached_responses_bridge_chat_completion_nonstream():
"""openai/responses chat-completions bridge: non-streaming cache hit replays as ModelResponse."""
from litellm import responses
from litellm.types.utils import CallTypes, ModelResponse
caching_handler = LLMCachingHandler(
original_function=responses, request_kwargs={}, start_time=datetime.now()
)
cached_result = {
"id": "chatcmpl-bridge-nonstream",
"object": "chat.completion",
"created": int(time.time()),
"model": "gpt-5.4",
"choices": [
{
"index": 0,
"message": {"role": "assistant", "content": "Hi!"},
"finish_reason": "stop",
}
],
"usage": {"prompt_tokens": 7, "completion_tokens": 11, "total_tokens": 18},
}
result = caching_handler._convert_cached_result_to_model_response(
cached_result=cached_result,
call_type=CallTypes.responses.value,
kwargs={
"model": "gpt-5.4",
"stream": False,
"messages": [{"role": "user", "content": "hi"}],
},
logging_obj=_build_logging_obj(CallTypes.responses.value, stream=False),
model="gpt-5.4",
args=(),
)
assert isinstance(result, ModelResponse)
assert result.choices[0].message.content == "Hi!"
def test_convert_cached_responses_legacy_nonstream_path():
"""Genuine ResponsesAPIResponse dict (no chatcmpl/choices) falls through legacy path."""
from litellm import responses
from litellm.types.llms.openai import ResponsesAPIResponse
from litellm.types.utils import CallTypes
caching_handler = LLMCachingHandler(
original_function=responses, request_kwargs={}, start_time=datetime.now()
)
cached_result = {
"id": "resp_legacy_nonstream",
"created_at": int(time.time()),
"status": "completed",
"model": "gpt-4o",
"object": "response",
"output": [
{
"type": "message",
"id": "msg_legacy",
"status": "completed",
"role": "assistant",
"content": [
{
"type": "output_text",
"text": "legacy response",
"annotations": [],
}
],
}
],
}
result = caching_handler._convert_cached_result_to_model_response(
cached_result=cached_result,
call_type=CallTypes.responses.value,
kwargs={"model": "gpt-4o", "input": "hi", "stream": False},
logging_obj=_build_logging_obj(CallTypes.responses.value, stream=False),
model="gpt-4o",
args=(),
)
assert isinstance(result, ResponsesAPIResponse)
assert result.id == "resp_legacy_nonstream"
def test_convert_cached_responses_legacy_stream_path():
"""Genuine ResponsesAPIResponse dict (no chatcmpl/choices) on stream falls through legacy path."""
from litellm import responses
from litellm.responses.streaming_iterator import (
CachedResponsesAPIStreamingIterator,
)
from litellm.types.utils import CallTypes
caching_handler = LLMCachingHandler(
original_function=responses, request_kwargs={}, start_time=datetime.now()
)
cached_result = {
"id": "resp_legacy_stream",
"created_at": int(time.time()),
"status": "completed",
"model": "gpt-4o",
"object": "response",
"output": [
{
"type": "message",
"id": "msg_legacy_stream",
"status": "completed",
"role": "assistant",
"content": [
{
"type": "output_text",
"text": "legacy stream",
"annotations": [],
}
],
}
],
}
result = caching_handler._convert_cached_result_to_model_response(
cached_result=cached_result,
call_type=CallTypes.responses.value,
kwargs={"model": "gpt-4o", "input": "hi", "stream": True},
logging_obj=_build_logging_obj(CallTypes.responses.value, stream=True),
model="gpt-4o",
args=(),
)
assert isinstance(result, CachedResponsesAPIStreamingIterator)

View file

@ -0,0 +1,150 @@
import os
import sys
from datetime import datetime
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
sys.path.insert(0, os.path.abspath("../../.."))
from litellm.completion_extras.litellm_responses_transformation.handler import (
ResponsesToCompletionBridgeHandler,
)
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
from litellm.types.utils import ModelResponse
def test_is_preformatted_cached_chat_stream_true():
stream = MagicMock(spec=CustomStreamWrapper)
stream.custom_llm_provider = "cached_response"
assert (
ResponsesToCompletionBridgeHandler._is_preformatted_cached_chat_stream(stream)
is True
)
def test_is_preformatted_cached_chat_stream_false_wrong_provider():
stream = MagicMock(spec=CustomStreamWrapper)
stream.custom_llm_provider = "openai"
assert (
ResponsesToCompletionBridgeHandler._is_preformatted_cached_chat_stream(stream)
is False
)
def test_is_preformatted_cached_chat_stream_false_wrong_type():
assert (
ResponsesToCompletionBridgeHandler._is_preformatted_cached_chat_stream(
{"object": "chat.completion.chunk"}
)
is False
)
def _bridge_kwargs(stream: bool):
logging_obj = LiteLLMLogging(
litellm_call_id="test-call",
call_type="completion",
model="gpt-5.4",
messages=[{"role": "user", "content": "hi"}],
function_id="fn-id",
stream=stream,
start_time=datetime.now(),
)
return {
"model": "gpt-5.4",
"custom_llm_provider": "openai",
"messages": [{"role": "user", "content": "hi"}],
"optional_params": {"stream": stream},
"litellm_params": {},
"headers": {},
"model_response": ModelResponse(),
"logging_obj": logging_obj,
}
def test_completion_returns_cached_model_response_directly():
"""Non-streaming bridge cache hit: responses() returns a ModelResponse -> bridge returns it as-is."""
cached = ModelResponse(id="chatcmpl-cached-nonstream", model="gpt-5.4")
bridge = ResponsesToCompletionBridgeHandler()
with (
patch.object(
bridge.transformation_handler,
"transform_request",
return_value={"model": "gpt-5.4", "input": "hi"},
),
patch("litellm.responses", return_value=cached),
):
result = bridge.completion(**_bridge_kwargs(stream=False))
assert result is cached
@pytest.mark.asyncio
async def test_acompletion_returns_cached_model_response_directly():
cached = ModelResponse(id="chatcmpl-cached-nonstream-async", model="gpt-5.4")
bridge = ResponsesToCompletionBridgeHandler()
with (
patch.object(
bridge.transformation_handler,
"transform_request",
return_value={"model": "gpt-5.4", "input": "hi"},
),
patch("litellm.aresponses", new=AsyncMock(return_value=cached)),
):
result = await bridge.acompletion(**_bridge_kwargs(stream=False))
assert result is cached
def test_completion_skips_rewrapping_preformatted_cached_chat_stream():
"""Streaming bridge cache hit returning CustomStreamWrapper(cached_response) -> bridge skips re-wrapping."""
stream = MagicMock(spec=CustomStreamWrapper)
stream.custom_llm_provider = "cached_response"
bridge = ResponsesToCompletionBridgeHandler()
with (
patch.object(
bridge.transformation_handler,
"transform_request",
return_value={"model": "gpt-5.4", "input": "hi"},
),
patch("litellm.responses", return_value=stream),
patch.object(
bridge,
"_apply_post_stream_processing",
side_effect=lambda s, *a, **kw: s,
) as post,
):
result = bridge.completion(**_bridge_kwargs(stream=True))
post.assert_called_once()
assert result is stream
@pytest.mark.asyncio
async def test_acompletion_skips_rewrapping_preformatted_cached_chat_stream():
stream = MagicMock(spec=CustomStreamWrapper)
stream.custom_llm_provider = "cached_response"
bridge = ResponsesToCompletionBridgeHandler()
with (
patch.object(
bridge.transformation_handler,
"transform_request",
return_value={"model": "gpt-5.4", "input": "hi"},
),
patch("litellm.aresponses", new=AsyncMock(return_value=stream)),
patch.object(
bridge,
"_apply_post_stream_processing",
side_effect=lambda s, *a, **kw: s,
) as post,
):
result = await bridge.acompletion(**_bridge_kwargs(stream=True))
post.assert_called_once()
assert result is stream

View file

@ -230,3 +230,30 @@ def test_transform_request_drops_user_metadata_with_additional_drop_params():
assert "metadata" not in result
assert result["litellm_metadata"]["internal_key"] == "secret"
def test_translate_responses_chunk_passthrough_chat_completion_chunk():
from litellm.completion_extras.litellm_responses_transformation.transformation import (
OpenAiResponsesToChatCompletionStreamIterator,
)
chat_chunk = {
"id": "chatcmpl-cache-passthrough",
"object": "chat.completion.chunk",
"created": 1779104834,
"model": "gpt-5.4",
"choices": [
{
"index": 0,
"delta": {"role": "assistant", "content": "Hi! How can I help?"},
"finish_reason": None,
}
],
}
result = OpenAiResponsesToChatCompletionStreamIterator.translate_responses_chunk_to_openai_stream(
chat_chunk
)
assert result.choices[0].delta.content == "Hi! How can I help?"
assert result.choices[0].finish_reason is None

View file

@ -460,6 +460,104 @@ async def test_assemble_user_object_does_not_override_metadata_max_budget(
), "max_budget from metadata must not be replaced by the DB value"
async def test_assemble_user_object_populates_user_email_and_alias_from_db(
prometheus_logger,
):
db_user = MagicMock()
db_user.max_budget = None
db_user.budget_reset_at = None
db_user.user_email = "alice@example.com"
db_user.user_alias = "Alice"
with patch("litellm.proxy.auth.auth_checks.get_user_object") as mock_get_user:
mock_get_user.return_value = db_user
user_object = await prometheus_logger._assemble_user_object(
user_id="user-abc-123",
spend=10.0,
max_budget=None,
response_cost=0.5,
)
assert user_object.user_email == "alice@example.com"
assert user_object.user_alias == "Alice"
def test_set_user_budget_metrics_default_no_email_alias_labels(
prometheus_logger,
):
"""By default (flag off), only user label is emitted."""
import litellm
from litellm.proxy._types import LiteLLM_UserTable
litellm.prometheus_user_budget_label_include_email_alias = False
user = LiteLLM_UserTable(
user_id="user-abc-123",
user_email="alice@example.com",
user_alias="Alice",
spend=25.0,
max_budget=100.0,
budget_reset_at=datetime(2026, 3, 1, tzinfo=timezone.utc),
)
prometheus_logger.litellm_remaining_user_budget_metric = MagicMock()
prometheus_logger.litellm_user_max_budget_metric = MagicMock()
prometheus_logger.litellm_user_budget_remaining_hours_metric = MagicMock()
prometheus_logger._set_user_budget_metrics(user)
prometheus_logger.litellm_remaining_user_budget_metric.labels.assert_called_once_with(
user="user-abc-123",
)
def test_set_user_budget_metrics_includes_user_email_and_alias_labels_when_opted_in(
prometheus_logger,
):
"""When prometheus_user_budget_label_include_email_alias=True, email+alias labels appear."""
import litellm
from litellm.proxy._types import LiteLLM_UserTable
litellm.prometheus_user_budget_label_include_email_alias = True
user = LiteLLM_UserTable(
user_id="user-abc-123",
user_email="alice@example.com",
user_alias="Alice",
spend=25.0,
max_budget=100.0,
budget_reset_at=datetime(2026, 3, 1, tzinfo=timezone.utc),
)
prometheus_logger.litellm_remaining_user_budget_metric = MagicMock()
prometheus_logger.litellm_user_max_budget_metric = MagicMock()
prometheus_logger.litellm_user_budget_remaining_hours_metric = MagicMock()
try:
prometheus_logger._set_user_budget_metrics(user)
prometheus_logger.litellm_remaining_user_budget_metric.labels.assert_called_once_with(
user="user-abc-123",
user_email="alice@example.com",
user_alias="Alice",
)
prometheus_logger.litellm_remaining_user_budget_metric.labels().set.assert_called_once_with(
75.0
)
prometheus_logger.litellm_user_max_budget_metric.labels.assert_called_once_with(
user="user-abc-123",
user_email="alice@example.com",
user_alias="Alice",
)
prometheus_logger.litellm_user_budget_remaining_hours_metric.labels.assert_called_once_with(
user="user-abc-123",
user_email="alice@example.com",
user_alias="Alice",
)
finally:
litellm.prometheus_user_budget_label_include_email_alias = False
async def test_set_user_budget_metrics_after_api_request_no_inf_when_metadata_budget_none(
prometheus_logger,
):

View file

@ -0,0 +1,119 @@
"""
Test that BedrockBatchesConfig._get_openai_compatible_batch_metadata
sanitizes non-string metadata values injected by proxy guardrail hooks.
The OpenAI Batch Pydantic model requires metadata: Dict[str, str].
Proxy hooks (Model Armor, OpenAI Moderations, queue time tracking) inject
dicts, floats, and other non-string values that cause a ValidationError
when constructing LiteLLMBatch. This test suite verifies the sanitization
layer prevents that.
"""
import os
import sys
import pytest
sys.path.insert(0, os.path.abspath("../../../../.."))
from litellm.llms.bedrock.batches.transformation import BedrockBatchesConfig
class TestGetOpenaiCompatibleBatchMetadata:
"""Tests for _get_openai_compatible_batch_metadata."""
def test_string_values_pass_through_unchanged(self):
metadata = {"user_key": "user_value", "run_id": "abc123"}
result = BedrockBatchesConfig._get_openai_compatible_batch_metadata(metadata)
assert result == {"user_key": "user_value", "run_id": "abc123"}
def test_dict_values_serialized_to_json_string(self):
metadata = {
"_model_armor_response": {
"sanitizationResult": {"filterMatchState": "MATCH_FOUND"}
}
}
result = BedrockBatchesConfig._get_openai_compatible_batch_metadata(metadata)
assert "_model_armor_response" in result
assert isinstance(result["_model_armor_response"], str)
assert "MATCH_FOUND" in result["_model_armor_response"]
def test_float_values_serialized_to_string(self):
metadata = {"queue_time_seconds": 0.5}
result = BedrockBatchesConfig._get_openai_compatible_batch_metadata(metadata)
assert result == {"queue_time_seconds": "0.5"}
def test_none_values_excluded(self):
metadata = {"key": "value", "empty": None}
result = BedrockBatchesConfig._get_openai_compatible_batch_metadata(metadata)
assert "empty" not in result
assert result == {"key": "value"}
def test_standard_logging_guardrail_information_excluded(self):
metadata = {
"standard_logging_guardrail_information": {"some": "logging_data"},
"user_key": "keep_me",
}
result = BedrockBatchesConfig._get_openai_compatible_batch_metadata(metadata)
assert "standard_logging_guardrail_information" not in result
assert result == {"user_key": "keep_me"}
def test_non_dict_input_returns_empty_dict(self):
assert BedrockBatchesConfig._get_openai_compatible_batch_metadata(None) == {}
assert BedrockBatchesConfig._get_openai_compatible_batch_metadata("string") == {}
assert BedrockBatchesConfig._get_openai_compatible_batch_metadata(123) == {}
def test_empty_dict_returns_empty_dict(self):
assert BedrockBatchesConfig._get_openai_compatible_batch_metadata({}) == {}
def test_mixed_metadata_from_guardrails(self):
"""Simulate real metadata contaminated by proxy guardrails."""
metadata = {
"_model_armor_response": {"sanitizationResult": {"key": "val"}},
"_model_armor_status": "success",
"_openai_moderation_response": {"id": "mod-123", "flagged": False},
"queue_time_seconds": 1.23,
"headers": {"Authorization": "Bearer sk-xxx"},
"standard_logging_guardrail_information": {"internal": True},
"user_metadata_key": "user_value",
"none_field": None,
}
result = BedrockBatchesConfig._get_openai_compatible_batch_metadata(metadata)
# All values must be strings
for key, value in result.items():
assert isinstance(value, str), f"metadata[{key!r}] is {type(value)}, not str"
# Excluded keys
assert "standard_logging_guardrail_information" not in result
assert "none_field" not in result
# Preserved keys
assert result["_model_armor_status"] == "success"
assert result["user_metadata_key"] == "user_value"
def test_result_compatible_with_litellm_batch(self):
"""Verify sanitized metadata can construct a LiteLLMBatch without error."""
import time
from litellm.types.utils import LiteLLMBatch
metadata = {
"_model_armor_response": {"blocked": True},
"queue_time_seconds": 0.05,
"user_key": "value",
}
sanitized = BedrockBatchesConfig._get_openai_compatible_batch_metadata(metadata)
# This would raise ValidationError before the fix
batch = LiteLLMBatch(
id="arn:aws:bedrock:us-east-1:123:model-invocation-job/test",
object="batch",
endpoint="/v1/chat/completions",
input_file_id="file-123",
completion_window="24h",
status="validating",
created_at=int(time.time()),
metadata=sanitized,
)
assert batch.metadata == sanitized

View file

@ -957,3 +957,50 @@ def test_titan_image_embedding_cost_uses_per_image_rate():
assert response.usage is not None
assert response.usage.prompt_tokens_details is not None
assert response.usage.prompt_tokens_details.image_count == 1
@pytest.mark.parametrize(
"encoding_format,expected_embedding_types",
[
("float", ["float"]),
("base64", ["base64"]),
(["float", "int8"], ["float", "int8"]),
],
)
def test_bedrock_cohere_embedding_types_wrapped_as_list(
encoding_format, expected_embedding_types
):
"""
Bedrock Cohere expects `embedding_types` as a JSON array, not a raw string.
Regression test for: Bedrock returns
Malformed input request: #/embedding_types: expected type: JSONArray, found: String
when `encoding_format` is passed as a string.
"""
litellm.set_verbose = True
client = HTTPHandler()
model = "bedrock/cohere.embed-multilingual-v3"
with patch.object(client, "post") as mock_post:
mock_response = Mock()
mock_response.status_code = 200
mock_response.text = json.dumps(cohere_embedding_response)
mock_response.json = lambda: json.loads(mock_response.text)
mock_post.return_value = mock_response
response = litellm.embedding(
model=model,
input=test_input,
encoding_format=encoding_format,
client=client,
aws_region_name="us-east-1",
aws_bedrock_runtime_endpoint="https://bedrock-runtime.us-east-1.amazonaws.com",
api_key="test-bearer-token-12345",
)
assert isinstance(response, litellm.EmbeddingResponse)
request_body = json.loads(mock_post.call_args.kwargs.get("data", "{}"))
assert "embedding_types" in request_body
assert request_body["embedding_types"] == expected_embedding_types
assert isinstance(request_body["embedding_types"], list)

View file

@ -0,0 +1,141 @@
"""
Unit tests for DashScope embedding transformation.
"""
import json
import os
import sys
from unittest.mock import MagicMock
import httpx
import pytest
sys.path.insert(0, os.path.abspath("../../../../.."))
from litellm.llms.dashscope.common_utils import DashScopeError
from litellm.llms.dashscope.embed.transformation import (
DEFAULT_API_BASE,
DashScopeEmbeddingConfig,
)
from litellm.types.utils import EmbeddingResponse
def test_validate_environment_and_url():
config = DashScopeEmbeddingConfig()
headers = config.validate_environment(
headers={},
model="text-embedding-v4",
messages=[],
optional_params={},
litellm_params={},
api_key="sk-test",
)
assert headers["Authorization"] == "Bearer sk-test"
url = config.get_complete_url(
api_base=None,
api_key="sk-test",
model="text-embedding-v4",
optional_params={},
litellm_params={},
)
assert url == f"{DEFAULT_API_BASE}/embeddings"
def test_transform_embedding_request():
config = DashScopeEmbeddingConfig()
data = config.transform_embedding_request(
model="text-embedding-v4",
input=["风急天高猿啸哀"],
optional_params={"dimensions": 1024, "encoding_format": "float"},
headers={},
)
assert data == {
"model": "text-embedding-v4",
"input": ["风急天高猿啸哀"],
"dimensions": 1024,
"encoding_format": "float",
}
def test_transform_embedding_response_success():
config = DashScopeEmbeddingConfig()
payload = {
"data": [
{"embedding": [0.1, 0.2], "index": 0, "object": "embedding"},
],
"model": "text-embedding-v4",
"object": "list",
"usage": {"prompt_tokens": 5, "total_tokens": 5},
"id": "73591b79-xxxx",
}
raw = httpx.Response(
status_code=200,
content=json.dumps(payload).encode("utf-8"),
request=httpx.Request("POST", "https://example.com"),
)
result = config.transform_embedding_response(
model="text-embedding-v4",
raw_response=raw,
model_response=EmbeddingResponse(),
logging_obj=MagicMock(),
api_key="sk-x",
request_data={"input": ["a"]},
optional_params={},
litellm_params={},
)
assert result.model == "text-embedding-v4"
assert len(result.data) == 1
assert result.usage.prompt_tokens == 5
def test_transform_embedding_request_user_param():
config = DashScopeEmbeddingConfig()
data = config.transform_embedding_request(
model="text-embedding-v4",
input=["hello"],
optional_params={"user": "user-123"},
headers={},
)
assert data["user"] == "user-123"
def test_map_openai_params_drops_unsupported_with_drop_params():
config = DashScopeEmbeddingConfig()
result = config.map_openai_params(
non_default_params={"dimensions": 512, "unknown_param": "value"},
optional_params={},
model="text-embedding-v4",
drop_params=True,
)
assert result == {"dimensions": 512}
assert "unknown_param" not in result
def test_transform_embedding_response_error():
config = DashScopeEmbeddingConfig()
payload = {
"error": {
"message": "Incorrect API key provided.",
"type": "invalid_request_error",
"code": "invalid_api_key",
}
}
raw = httpx.Response(
status_code=401,
content=json.dumps(payload).encode("utf-8"),
request=httpx.Request("POST", "https://example.com"),
)
with pytest.raises(DashScopeError) as exc:
config.transform_embedding_response(
model="text-embedding-v4",
raw_response=raw,
model_response=EmbeddingResponse(),
logging_obj=MagicMock(),
api_key="sk-bad",
request_data={"input": ["a"]},
optional_params={},
litellm_params={},
)
assert exc.value.status_code == 401
assert "Incorrect API key" in exc.value.message

View file

@ -0,0 +1,328 @@
"""
Unit tests for DashScope rerank transformation.
"""
import json
import os
import sys
from unittest.mock import MagicMock
import httpx
import pytest
sys.path.insert(0, os.path.abspath("../../../../.."))
from litellm.llms.dashscope.common_utils import DashScopeError
from litellm.llms.dashscope.rerank.transformation import (
DEFAULT_RERANK_URL,
DashScopeRerankConfig,
)
from litellm.types.rerank import RerankResponse
class TestDashScopeRerankURL:
def setup_method(self):
self.config = DashScopeRerankConfig()
def test_default_url(self):
url = self.config.get_complete_url(api_base=None, model="qwen3-rerank")
assert url == DEFAULT_RERANK_URL
def test_explicit_v1_base_appends_reranks(self):
url = self.config.get_complete_url(
api_base="https://dashscope.aliyuncs.com/compatible-mode/v1",
model="qwen3-rerank",
)
assert url == "https://dashscope.aliyuncs.com/compatible-mode/v1/reranks"
def test_intl_v1_base_appends_reranks(self):
url = self.config.get_complete_url(
api_base="https://dashscope-intl.aliyuncs.com/compatible-mode/v1",
model="qwen3-rerank",
)
assert url == "https://dashscope-intl.aliyuncs.com/compatible-mode/v1/reranks"
def test_already_complete_url_passthrough(self):
full = "https://dashscope.aliyuncs.com/compatible-api/v1/reranks"
assert self.config.get_complete_url(api_base=full, model="qwen3-rerank") == full
def test_trailing_slash_stripped(self):
full = "https://dashscope.aliyuncs.com/compatible-api/v1/reranks/"
assert self.config.get_complete_url(
api_base=full, model="qwen3-rerank"
) == full.rstrip("/")
def test_custom_v1_base_appends_reranks(self):
url = self.config.get_complete_url(
api_base="https://my-proxy.example.com/v1", model="qwen3-rerank"
)
assert url == "https://my-proxy.example.com/v1/reranks"
class TestDashScopeRerankRequest:
def setup_method(self):
self.config = DashScopeRerankConfig()
def test_validate_environment_with_explicit_key(self):
headers = self.config.validate_environment(
headers={}, model="qwen3-rerank", api_key="sk-test"
)
assert headers["Authorization"] == "Bearer sk-test"
assert headers["content-type"] == "application/json"
def test_validate_environment_missing_key(self, monkeypatch):
monkeypatch.delenv("DASHSCOPE_API_KEY", raising=False)
with pytest.raises(ValueError, match="DASHSCOPE_API_KEY"):
self.config.validate_environment(
headers={}, model="qwen3-rerank", api_key=None
)
def test_validate_environment_falls_back_to_env(self, monkeypatch):
monkeypatch.setenv("DASHSCOPE_API_KEY", "env-key")
headers = self.config.validate_environment(
headers={}, model="qwen3-rerank", api_key=None
)
assert headers["Authorization"] == "Bearer env-key"
def test_supported_params(self):
assert self.config.get_supported_cohere_rerank_params("qwen3-rerank") == [
"query",
"documents",
"top_n",
"return_documents",
]
def test_map_params_drops_unsupported(self):
# qwen3-rerank accepts query/documents/top_n/return_documents.
# rank_fields and max_*_per_doc are silently dropped.
params = self.config.map_cohere_rerank_params(
non_default_params={},
model="qwen3-rerank",
drop_params=False,
query="什么是文本排序模型",
documents=["d1", "d2"],
top_n=2,
rank_fields=["title"],
return_documents=True,
max_chunks_per_doc=5,
max_tokens_per_doc=100,
)
assert params == {
"query": "什么是文本排序模型",
"documents": ["d1", "d2"],
"top_n": 2,
"return_documents": True,
}
def test_transform_request_full(self):
body = self.config.transform_rerank_request(
model="qwen3-rerank",
optional_rerank_params={
"query": "如何制作美味的苹果派?",
"documents": ["a", "b"],
"top_n": 5,
"return_documents": True,
},
headers={},
)
assert body == {
"model": "qwen3-rerank",
"query": "如何制作美味的苹果派?",
"documents": ["a", "b"],
"top_n": 5,
"return_documents": True,
}
def test_transform_request_omits_unset_optional(self):
body = self.config.transform_rerank_request(
model="qwen3-rerank",
optional_rerank_params={"query": "q", "documents": ["a"]},
headers={},
)
assert "top_n" not in body
assert "return_documents" not in body
def test_transform_request_requires_query(self):
with pytest.raises(ValueError, match="query"):
self.config.transform_rerank_request(
model="qwen3-rerank",
optional_rerank_params={"documents": ["a"]},
headers={},
)
def test_transform_request_requires_documents(self):
with pytest.raises(ValueError, match="documents"):
self.config.transform_rerank_request(
model="qwen3-rerank",
optional_rerank_params={"query": "q"},
headers={},
)
class TestDashScopeRerankResponse:
def setup_method(self):
self.config = DashScopeRerankConfig()
self.logging = MagicMock()
def _resp(self, body, status_code=200):
return httpx.Response(
status_code=status_code,
content=json.dumps(body).encode(),
request=httpx.Request("POST", "https://example.com"),
)
def test_success_response(self):
body = {
"object": "list",
"results": [
{"index": 0, "relevance_score": 0.93},
{"index": 2, "relevance_score": 0.34},
],
"model": "qwen3-rerank",
"id": "85ba5752",
"usage": {"total_tokens": 79},
}
out = self.config.transform_rerank_response(
model="qwen3-rerank",
raw_response=self._resp(body),
model_response=RerankResponse(),
logging_obj=self.logging,
api_key="sk",
request_data={"query": "q"},
)
assert out.id == "85ba5752"
assert out.results == [
{"index": 0, "relevance_score": 0.93},
{"index": 2, "relevance_score": 0.34},
]
assert out.meta == {
"billed_units": {"total_tokens": 79},
"tokens": {"input_tokens": 79},
}
def test_response_with_return_documents_real_payload(self):
# Verbatim sample from a real qwen3-rerank call with return_documents=true.
body = {
"object": "list",
"results": [
{
"document": {
"text": "苹果派的制作步骤包括准备面团、切苹果、调制馅料、组装和烘烤。"
},
"index": 1,
"relevance_score": 0.8304247466067356,
},
{
"document": {
"text": "制作苹果派时,预先煮软苹果可以缩短烘烤时间。"
},
"index": 3,
"relevance_score": 0.7142660211908354,
},
],
"model": "qwen3-rerank",
"id": "e191b077-97c4-9929-b121-c2fbd2c7b0af",
"usage": {"total_tokens": 192},
}
out = self.config.transform_rerank_response(
model="qwen3-rerank",
raw_response=self._resp(body),
model_response=RerankResponse(),
logging_obj=self.logging,
request_data={"query": "如何制作美味的苹果派?"},
)
assert out.id == "e191b077-97c4-9929-b121-c2fbd2c7b0af"
assert out.results == [
{
"index": 1,
"relevance_score": 0.8304247466067356,
"document": {
"text": "苹果派的制作步骤包括准备面团、切苹果、调制馅料、组装和烘烤。"
},
},
{
"index": 3,
"relevance_score": 0.7142660211908354,
"document": {"text": "制作苹果派时,预先煮软苹果可以缩短烘烤时间。"},
},
]
assert out.meta == {
"billed_units": {"total_tokens": 192},
"tokens": {"input_tokens": 192},
}
def test_response_string_document_normalized(self):
# Defensive path: if a future API revision returns a bare string,
# normalize to {"text": ...} so downstream code stays consistent.
body = {
"results": [{"index": 0, "relevance_score": 0.9, "document": "hello"}],
"model": "qwen3-rerank",
"usage": {"total_tokens": 5},
}
out = self.config.transform_rerank_response(
model="qwen3-rerank",
raw_response=self._resp(body),
model_response=RerankResponse(),
logging_obj=self.logging,
)
assert out.results[0]["document"] == {"text": "hello"}
def test_missing_id_generates_uuid(self):
body = {"results": [{"index": 0, "relevance_score": 0.5}], "usage": {}}
out = self.config.transform_rerank_response(
model="qwen3-rerank",
raw_response=self._resp(body),
model_response=RerankResponse(),
logging_obj=self.logging,
)
assert out.id is not None and len(out.id) > 0
def test_error_envelope_raises(self):
body = {
"code": "InvalidApiKey",
"message": "Invalid API-key provided.",
"request_id": "fb53",
}
with pytest.raises(DashScopeError) as exc_info:
self.config.transform_rerank_response(
model="qwen3-rerank",
raw_response=self._resp(body, status_code=401),
model_response=RerankResponse(),
logging_obj=self.logging,
)
assert "Invalid API-key provided." in str(exc_info.value)
def test_non_json_response_raises(self):
bad = httpx.Response(
status_code=500,
content=b"<html>bad gateway</html>",
request=httpx.Request("POST", "https://example.com"),
)
with pytest.raises(DashScopeError):
self.config.transform_rerank_response(
model="qwen3-rerank",
raw_response=bad,
model_response=RerankResponse(),
logging_obj=self.logging,
)
def test_get_error_class(self):
err = self.config.get_error_class(
error_message="boom", status_code=500, headers={}
)
assert isinstance(err, DashScopeError)
assert err.status_code == 500
class TestProviderConfigManagerDispatch:
def test_dashscope_returns_rerank_config(self):
import litellm
from litellm.utils import ProviderConfigManager
cfg = ProviderConfigManager.get_provider_rerank_config(
model="qwen3-rerank",
provider=litellm.LlmProviders.DASHSCOPE,
api_base=None,
present_version_params=[],
)
assert isinstance(cfg, DashScopeRerankConfig)

View file

@ -0,0 +1,189 @@
import litellm
from litellm.llms.anthropic.experimental_pass_through.messages.transformation import (
AnthropicMessagesConfig,
)
from litellm.llms.deepseek.messages.transformation import (
DeepSeekAnthropicMessagesConfig,
)
from litellm.types.router import GenericLiteLLMParams
from litellm.utils import ProviderConfigManager
def test_deepseek_provider_uses_anthropic_messages_config():
config = ProviderConfigManager.get_provider_anthropic_messages_config(
model="deepseek-v4-pro",
provider=litellm.LlmProviders.DEEPSEEK,
)
assert isinstance(config, DeepSeekAnthropicMessagesConfig)
assert config.custom_llm_provider == "deepseek"
def test_deepseek_anthropic_messages_config_defaults():
config = DeepSeekAnthropicMessagesConfig()
assert config.custom_llm_provider == "deepseek"
assert config.get_api_base() == "https://api.deepseek.com/anthropic"
def test_anthropic_provider_keeps_default_config_for_deepseek_named_model():
config = ProviderConfigManager.get_provider_anthropic_messages_config(
model="deepseek-v4-pro",
provider=litellm.LlmProviders.ANTHROPIC,
)
assert isinstance(config, AnthropicMessagesConfig)
assert not isinstance(config, DeepSeekAnthropicMessagesConfig)
def test_deepseek_anthropic_messages_url_defaults_to_anthropic_endpoint():
config = DeepSeekAnthropicMessagesConfig()
assert (
config.get_complete_url(
api_base=None,
api_key=None,
model="deepseek-v4-pro",
optional_params={},
litellm_params={},
)
== "https://api.deepseek.com/anthropic/v1/messages"
)
assert (
config.get_complete_url(
api_base="https://api.deepseek.com/anthropic/v1",
api_key=None,
model="deepseek-v4-pro",
optional_params={},
litellm_params={},
)
== "https://api.deepseek.com/anthropic/v1/messages"
)
assert (
config.get_complete_url(
api_base="https://api.deepseek.com/anthropic",
api_key=None,
model="deepseek-v4-pro",
optional_params={},
litellm_params={},
)
== "https://api.deepseek.com/anthropic/v1/messages"
)
assert (
config.get_complete_url(
api_base="https://api.deepseek.com",
api_key=None,
model="deepseek-v4-pro",
optional_params={},
litellm_params={},
)
== "https://api.deepseek.com/anthropic/v1/messages"
)
assert (
config.get_complete_url(
api_base="https://api.deepseek.com/v1",
api_key=None,
model="deepseek-v4-pro",
optional_params={},
litellm_params={},
)
== "https://api.deepseek.com/anthropic/v1/messages"
)
assert (
config.get_complete_url(
api_base="https://api.deepseek.com/v1/messages",
api_key=None,
model="deepseek-v4-pro",
optional_params={},
litellm_params={},
)
== "https://api.deepseek.com/anthropic/v1/messages"
)
def test_deepseek_anthropic_messages_headers_use_deepseek_key():
config = DeepSeekAnthropicMessagesConfig()
headers, api_base = config.validate_anthropic_messages_environment(
headers={},
model="deepseek-v4-pro",
messages=[],
optional_params={},
litellm_params={},
api_key="sk-deepseek",
api_base="https://example.test/anthropic",
)
assert api_base == "https://example.test/anthropic"
assert headers["x-api-key"] == "sk-deepseek"
assert headers["anthropic-version"] == "2023-06-01"
assert headers["content-type"] == "application/json"
def test_deepseek_anthropic_messages_preserves_thinking_and_sanitizes_custom_tools():
config = DeepSeekAnthropicMessagesConfig()
messages = [
{
"role": "user",
"content": "Use the tool.",
},
{
"role": "assistant",
"content": [
{
"type": "thinking",
"thinking": "I should call the tool.",
"signature": "sig",
},
{
"type": "tool_use",
"id": "toolu_123",
"name": "get_weather",
"input": {"city": "Sao Paulo"},
},
],
},
{
"role": "user",
"content": [
{
"type": "tool_result",
"tool_use_id": "toolu_123",
"content": "Sunny",
}
],
},
]
request = config.transform_anthropic_messages_request(
model="deepseek-v4-pro",
messages=messages,
anthropic_messages_optional_request_params={
"max_tokens": 100,
"thinking": {"type": "enabled", "budget_tokens": 1024},
"tools": [
{
"type": "custom",
"name": "get_weather",
"description": "Get weather",
"input_schema": {"type": "object"},
},
{
"type": "web_search_20260209",
"name": "web_search",
"max_uses": 1,
},
],
},
litellm_params=GenericLiteLLMParams(),
headers={},
)
assert request["messages"] == messages
assert request["thinking"] == {"type": "enabled", "budget_tokens": 1024}
assert request["tools"][0] == {
"name": "get_weather",
"description": "Get weather",
"input_schema": {"type": "object"},
}
assert request["tools"][1]["type"] == "web_search_20260209"

View file

@ -0,0 +1,433 @@
"""
Tests for handling malformed or invalid 'file' content blocks (missing or null
`file` sub-field, HTTP file_id URLs for Google AI Studio).
Regression tests for:
- litellm/llms/vertex_ai/gemini/transformation.py
- litellm/llms/gemini/chat/transformation.py
- litellm/litellm_core_utils/prompt_templates/common_utils.py
(migrate_file_to_image_url raises on missing `file`; file-id helpers skip non-OpenAI shapes)
- litellm/litellm_core_utils/prompt_templates/factory.py (Bedrock + Anthropic)
- litellm/llms/openai/chat/gpt_transformation.py
"""
import asyncio
import copy
from typing import List, cast
import pytest
import litellm
from litellm.litellm_core_utils.prompt_templates.common_utils import (
get_file_ids_from_messages,
migrate_file_to_image_url,
update_messages_with_model_file_ids,
)
from litellm.litellm_core_utils.prompt_templates.factory import (
BedrockConverseMessagesProcessor,
anthropic_process_openai_file_message,
)
from litellm.llms.gemini.chat.transformation import GoogleAIStudioGeminiConfig
from litellm.llms.openai.chat.gpt_transformation import OpenAIGPTConfig
from litellm.llms.vertex_ai.gemini.transformation import (
_gemini_convert_messages_with_history,
)
from litellm.types.llms.openai import (
AllMessageValues,
ChatCompletionFileObject,
OpenAIMessageContentListBlock,
)
_MALFORMED_MESSAGES_RAW = [
{
"role": "user",
"content": [
{"type": "text", "text": "hello"},
{"type": "file"}, # Missing required "file" sub-field
],
}
]
_WELL_FORMED_MESSAGES_RAW = [
{
"role": "user",
"content": [
{"type": "text", "text": "hello"},
{
"type": "file",
"file": {"file_id": "file-abc123", "format": "pdf"},
},
],
}
]
MALFORMED_FILE_OBJECT: ChatCompletionFileObject = cast(
ChatCompletionFileObject, {"type": "file"}
)
EXPLICIT_NULL_FILE_OBJECT: ChatCompletionFileObject = cast(
ChatCompletionFileObject,
{"type": "file", "file": None},
)
def _malformed() -> List[AllMessageValues]:
return copy.deepcopy(cast(List[AllMessageValues], _MALFORMED_MESSAGES_RAW))
def _well_formed() -> List[AllMessageValues]:
return copy.deepcopy(cast(List[AllMessageValues], _WELL_FORMED_MESSAGES_RAW))
def _explicit_null_file_in_content() -> List[AllMessageValues]:
return copy.deepcopy(
cast(
List[AllMessageValues],
[
{
"role": "user",
"content": [
{"type": "text", "text": "hello"},
{"type": "file", "file": None},
],
}
],
)
)
# ---------------------------------------------------------------------------
# vertex_ai/gemini/transformation.py
# ---------------------------------------------------------------------------
def test_gemini_convert_messages_malformed_file_raises_bad_request():
"""_gemini_convert_messages_with_history should raise BadRequestError (not KeyError)
when a content block has type='file' but no 'file' sub-field."""
with pytest.raises(litellm.BadRequestError, match="missing the required 'file' field"):
_gemini_convert_messages_with_history(
messages=_malformed(),
model="gemini-2.0-flash",
)
def test_gemini_convert_messages_explicit_null_file_field_raises_bad_request():
"""Explicit JSON null for `file` must be rejected like a missing `file` key."""
with pytest.raises(litellm.BadRequestError, match="missing the required 'file' field"):
_gemini_convert_messages_with_history(
messages=_explicit_null_file_in_content(),
model="gemini-2.0-flash",
)
# ---------------------------------------------------------------------------
# gemini/chat/transformation.py - GoogleAIStudioGeminiConfig
# ---------------------------------------------------------------------------
def test_google_ai_studio_transform_messages_malformed_file_raises_bad_request():
"""GoogleAIStudioGeminiConfig._transform_messages should raise BadRequestError
when a content block has type='file' but no 'file' sub-field."""
config = GoogleAIStudioGeminiConfig()
with pytest.raises(litellm.BadRequestError, match="missing the required 'file' field"):
config._transform_messages(messages=_malformed(), model="gemini-2.0-flash")
def test_google_ai_studio_transform_messages_explicit_null_file_field_raises_bad_request():
"""Explicit JSON null for `file` must be rejected like a missing `file` key."""
config = GoogleAIStudioGeminiConfig()
with pytest.raises(litellm.BadRequestError, match="missing the required 'file' field"):
config._transform_messages(
messages=_explicit_null_file_in_content(), model="gemini-2.0-flash"
)
def test_google_ai_studio_transform_messages_http_file_id_converts_to_base64(monkeypatch):
"""Google AI Studio rejects raw HTTP(S) file URLs; _transform_messages should
fetch and replace them with base64 `file_data` before conversion."""
# Data URL shape so downstream Gemini media parsing accepts the inlined bytes
# (mirrors real `convert_url_to_base64` output from `_process_image_response`).
fake_file_data = "data:application/pdf;base64,aGVsbG8="
def _fake_convert_url_to_base64(url: str) -> str:
assert url == "https://example.com/doc.pdf"
return fake_file_data
monkeypatch.setattr(
"litellm.llms.gemini.chat.transformation.convert_url_to_base64",
_fake_convert_url_to_base64,
)
messages = cast(
List[AllMessageValues],
[
{
"role": "user",
"content": [
{"type": "text", "text": "hello"},
{
"type": "file",
"file": {
"file_id": "https://example.com/doc.pdf",
"format": "pdf",
},
},
],
}
],
)
config = GoogleAIStudioGeminiConfig()
config._transform_messages(messages=messages, model="gemini-2.0-flash")
content = messages[0].get("content")
assert isinstance(content, list)
file_block = next(c for c in content if isinstance(c, dict) and c.get("type") == "file")
file_field = file_block.get("file")
assert isinstance(file_field, dict)
assert file_field.get("file_data") == fake_file_data
assert "file_id" not in file_field
def test_google_ai_studio_transform_messages_http_file_id_convert_failure_leaves_file_unchanged(
monkeypatch,
):
"""If convert_url_to_base64 fails, the Studio prep step must not mutate the block
(see try/except in GoogleAIStudioGeminiConfig._transform_messages)."""
https_id = "https://example.com/missing.pdf"
def _raise(_url: str) -> str:
raise litellm.ImageFetchError("simulated fetch failure")
monkeypatch.setattr(
"litellm.llms.gemini.chat.transformation.convert_url_to_base64",
_raise,
)
messages = cast(
List[AllMessageValues],
[
{
"role": "user",
"content": [
{"type": "text", "text": "hello"},
{
"type": "file",
"file": {
"file_id": https_id,
"format": "application/pdf",
},
},
],
}
],
)
config = GoogleAIStudioGeminiConfig()
config._transform_messages(messages=messages, model="gemini-2.0-flash")
content = messages[0].get("content")
assert isinstance(content, list)
file_block = next(c for c in content if isinstance(c, dict) and c.get("type") == "file")
file_field = file_block.get("file")
assert isinstance(file_field, dict)
assert file_field.get("file_id") == https_id
assert file_field.get("format") == "application/pdf"
assert "file_data" not in file_field
# ---------------------------------------------------------------------------
# common_utils.py - update_messages_with_model_file_ids
# ---------------------------------------------------------------------------
def test_update_messages_with_model_file_ids_malformed_skips_non_openai_file_block():
"""Non-OpenAI file blocks (e.g. missing nested `file` dict) are skipped so callers
relying on LangChain v1 / provider-native shapes are not rejected here."""
messages = _malformed()
result = update_messages_with_model_file_ids(
messages=messages,
model_id="some-model",
model_file_id_mapping={},
)
assert result == messages
content = result[0].get("content")
assert isinstance(content, list)
file_block = next(c for c in content if isinstance(c, dict) and c.get("type") == "file")
assert "file" not in file_block
def test_update_messages_with_model_file_ids_well_formed_updates():
"""update_messages_with_model_file_ids should update file_id for well-formed blocks."""
mapping = {"file-abc123": {"some-model": "provider-file-xyz"}}
result = update_messages_with_model_file_ids(
messages=_well_formed(),
model_id="some-model",
model_file_id_mapping=mapping,
)
content = result[0].get("content")
assert isinstance(content, list)
file_block = next(c for c in content if c.get("type") == "file")
assert file_block.get("file", {}).get("file_id") == "provider-file-xyz"
# ---------------------------------------------------------------------------
# common_utils.py - get_file_ids_from_messages
# ---------------------------------------------------------------------------
def test_get_file_ids_from_messages_malformed_skips_non_openai_file_block():
"""Blocks with type='file' but no OpenAI `file` sub-dict yield no extracted ids."""
assert get_file_ids_from_messages(messages=_malformed()) == []
def test_get_file_ids_from_messages_well_formed_returns_ids():
"""get_file_ids_from_messages should extract file_id from well-formed blocks."""
messages: List[AllMessageValues] = cast(
List[AllMessageValues],
[
{
"role": "user",
"content": [
{"type": "text", "text": "hello"},
{"type": "file", "file": {"file_id": "file-abc123", "format": "pdf"}},
],
}
],
)
result = get_file_ids_from_messages(messages=messages)
assert result == ["file-abc123"]
# ---------------------------------------------------------------------------
# factory.py - BedrockConverseMessagesProcessor (sync + async)
# ---------------------------------------------------------------------------
def test_bedrock_process_file_message_malformed_raises_bad_request():
"""_process_file_message should raise BadRequestError (not KeyError)
when the file object is missing the 'file' sub-field."""
with pytest.raises(litellm.BadRequestError, match="missing the required 'file' field"):
BedrockConverseMessagesProcessor._process_file_message(MALFORMED_FILE_OBJECT)
def test_bedrock_process_file_message_explicit_null_file_field_raises_bad_request():
with pytest.raises(litellm.BadRequestError, match="missing the required 'file' field"):
BedrockConverseMessagesProcessor._process_file_message(EXPLICIT_NULL_FILE_OBJECT)
def test_bedrock_async_process_file_message_malformed_raises_bad_request():
"""_async_process_file_message should raise BadRequestError (not KeyError)
when the file object is missing the 'file' sub-field."""
async def _run() -> None:
with pytest.raises(
litellm.BadRequestError, match="missing the required 'file' field"
):
await BedrockConverseMessagesProcessor._async_process_file_message(
MALFORMED_FILE_OBJECT
)
asyncio.run(_run())
def test_bedrock_async_process_file_message_explicit_null_file_field_raises_bad_request():
async def _run() -> None:
with pytest.raises(
litellm.BadRequestError, match="missing the required 'file' field"
):
await BedrockConverseMessagesProcessor._async_process_file_message(
EXPLICIT_NULL_FILE_OBJECT
)
asyncio.run(_run())
# ---------------------------------------------------------------------------
# openai/chat/gpt_transformation.py
# ---------------------------------------------------------------------------
def test_openai_apply_common_transform_malformed_file_raises_bad_request():
"""_apply_common_transform_content_item should raise BadRequestError (not KeyError)
when a content block has type='file' but no 'file' sub-field."""
config = OpenAIGPTConfig()
malformed_block: OpenAIMessageContentListBlock = cast(
OpenAIMessageContentListBlock, {"type": "file"}
)
with pytest.raises(litellm.BadRequestError, match="missing the required 'file' field"):
config._apply_common_transform_content_item(malformed_block)
def test_openai_apply_common_transform_explicit_null_file_field_raises_bad_request():
config = OpenAIGPTConfig()
explicit_null_block: OpenAIMessageContentListBlock = cast(
OpenAIMessageContentListBlock,
{"type": "file", "file": None},
)
with pytest.raises(litellm.BadRequestError, match="missing the required 'file' field"):
config._apply_common_transform_content_item(explicit_null_block)
def test_openai_apply_common_transform_well_formed_file_does_not_raise():
"""_apply_common_transform_content_item should not raise for well-formed file blocks."""
config = OpenAIGPTConfig()
well_formed_block: OpenAIMessageContentListBlock = cast(
OpenAIMessageContentListBlock,
{"type": "file", "file": {"file_id": "file-abc123"}},
)
result = config._apply_common_transform_content_item(well_formed_block)
assert result.get("type") == "file"
file_field = cast(ChatCompletionFileObject, result).get("file", {})
assert file_field.get("file_id") == "file-abc123"
# ---------------------------------------------------------------------------
# factory.py - anthropic_process_openai_file_message
# ---------------------------------------------------------------------------
def test_anthropic_process_openai_file_message_malformed_raises_bad_request():
"""anthropic_process_openai_file_message should raise BadRequestError (not KeyError)
when the file object is missing the 'file' sub-field."""
with pytest.raises(litellm.BadRequestError, match="missing the required 'file' field"):
anthropic_process_openai_file_message(MALFORMED_FILE_OBJECT)
def test_anthropic_process_openai_file_message_explicit_null_file_field_raises_bad_request():
with pytest.raises(litellm.BadRequestError, match="missing the required 'file' field"):
anthropic_process_openai_file_message(EXPLICIT_NULL_FILE_OBJECT)
def test_anthropic_process_openai_file_message_well_formed_file_id_does_not_raise():
"""anthropic_process_openai_file_message should not raise for a well-formed file_id block."""
well_formed: ChatCompletionFileObject = cast(
ChatCompletionFileObject,
{"type": "file", "file": {"file_id": "file-abc123"}},
)
result = anthropic_process_openai_file_message(well_formed)
assert result.get("type") in ("document", "image", "container_upload")
# ---------------------------------------------------------------------------
# common_utils.py - migrate_file_to_image_url
# ---------------------------------------------------------------------------
def test_migrate_file_to_image_url_malformed_raises_bad_request():
"""migrate_file_to_image_url should raise BadRequestError (not KeyError)
when the file object is missing the 'file' sub-field."""
with pytest.raises(litellm.BadRequestError, match="missing the required 'file' field"):
migrate_file_to_image_url(MALFORMED_FILE_OBJECT)
def test_migrate_file_to_image_url_explicit_null_file_field_raises_bad_request():
with pytest.raises(litellm.BadRequestError, match="missing the required 'file' field"):
migrate_file_to_image_url(EXPLICIT_NULL_FILE_OBJECT)
def test_migrate_file_to_image_url_well_formed_returns_image_url():
"""migrate_file_to_image_url should return an image_url block for a well-formed file."""
well_formed: ChatCompletionFileObject = cast(
ChatCompletionFileObject,
{"type": "file", "file": {"file_id": "file-abc123", "format": "png"}},
)
result = migrate_file_to_image_url(well_formed)
assert result.get("type") == "image_url"
image_url = result.get("image_url", {})
assert isinstance(image_url, dict)
assert image_url.get("url") == "file-abc123"

View file

@ -0,0 +1,52 @@
import pytest
from typing import List, cast
import litellm
from litellm.llms.vertex_ai.gemini.transformation import (
_gemini_convert_messages_with_history,
)
from litellm.types.llms.openai import AllMessageValues
def test_missing_image_url_field_raises_bad_request_error():
"""When element type is 'image_url' but 'image_url' field is missing, a BadRequestError is raised."""
messages = cast(
List[AllMessageValues],
[{"role": "user", "content": [{"type": "image_url"}]}],
)
with pytest.raises(litellm.BadRequestError) as exc_info:
_gemini_convert_messages_with_history(messages, model="gemini-1.5-pro")
assert "'image_url' field is missing" in str(exc_info.value)
def test_missing_url_inside_image_url_dict_raises_bad_request_error():
"""When image_url is a dict but 'url' key is absent, a BadRequestError is raised."""
messages = cast(
List[AllMessageValues],
[{"role": "user", "content": [{"type": "image_url", "image_url": {"detail": "high"}}]}],
)
with pytest.raises(litellm.BadRequestError) as exc_info:
_gemini_convert_messages_with_history(messages, model="gemini-1.5-pro")
assert "'url' field is missing inside" in str(exc_info.value)
def test_explicit_null_image_url_raises_bad_request_error():
"""When image_url key is present but explicitly null, a BadRequestError is raised."""
messages = cast(
List[AllMessageValues],
[{"role": "user", "content": [{"type": "image_url", "image_url": None}]}],
)
with pytest.raises(litellm.BadRequestError) as exc_info:
_gemini_convert_messages_with_history(messages, model="gemini-1.5-pro")
assert "'image_url' field is missing" in str(exc_info.value)
def test_empty_dict_image_url_raises_bad_request_error():
"""When image_url is an empty dict (no url), a BadRequestError is raised."""
messages = cast(
List[AllMessageValues],
[{"role": "user", "content": [{"type": "image_url", "image_url": {}}]}],
)
with pytest.raises(litellm.BadRequestError) as exc_info:
_gemini_convert_messages_with_history(messages, model="gemini-1.5-pro")
assert "'url' field is missing inside" in str(exc_info.value)

View file

@ -4290,3 +4290,444 @@ def test_transform_response_does_not_leak_body_on_parse_failure():
msg = str(exc_info.value)
assert "secret content" not in msg
assert "Error converting to valid response block" in msg
def test_chunk_parser_raises_on_429_error_chunk():
"""Test chunk_parser raises VertexAIError on 429 RESOURCE_EXHAUSTED error chunk"""
from unittest.mock import Mock
from litellm.llms.vertex_ai.common_utils import VertexAIError
from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
ModelResponseIterator,
)
error_chunk = {
"error": {
"code": 429,
"message": "Resource exhausted. Please try again later. Please refer to https://cloud.google.com/vertex-ai/generative-ai/docs/error-code-429 for more details.",
"status": "RESOURCE_EXHAUSTED",
}
}
logging_obj = Mock()
logging_obj.optional_params = {}
streaming_obj = ModelResponseIterator(
streaming_response=iter([]),
sync_stream=True,
logging_obj=logging_obj,
)
with pytest.raises(VertexAIError) as exc_info:
streaming_obj.chunk_parser(error_chunk)
assert exc_info.value.status_code == 429
assert "RESOURCE_EXHAUSTED" in exc_info.value.message
assert "Resource exhausted" in exc_info.value.message
def test_chunk_parser_raises_on_500_error_chunk():
"""Test chunk_parser raises VertexAIError on 500 INTERNAL error chunk"""
from unittest.mock import Mock
from litellm.llms.vertex_ai.common_utils import VertexAIError
from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
ModelResponseIterator,
)
error_chunk = {
"error": {
"code": 500,
"message": "Internal error encountered.",
"status": "INTERNAL",
}
}
logging_obj = Mock()
logging_obj.optional_params = {}
streaming_obj = ModelResponseIterator(
streaming_response=iter([]),
sync_stream=True,
logging_obj=logging_obj,
)
with pytest.raises(VertexAIError) as exc_info:
streaming_obj.chunk_parser(error_chunk)
assert exc_info.value.status_code == 500
assert "INTERNAL" in exc_info.value.message
def test_chunk_parser_raises_on_error_chunk_with_minimal_fields():
"""Test chunk_parser handles error chunks with missing optional fields"""
from unittest.mock import Mock
from litellm.llms.vertex_ai.common_utils import VertexAIError
from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
ModelResponseIterator,
)
error_chunk = {
"error": {
"code": 429,
"message": "Resource exhausted.",
}
}
logging_obj = Mock()
logging_obj.optional_params = {}
streaming_obj = ModelResponseIterator(
streaming_response=iter([]),
sync_stream=True,
logging_obj=logging_obj,
)
with pytest.raises(VertexAIError) as exc_info:
streaming_obj.chunk_parser(error_chunk)
assert exc_info.value.status_code == 429
def test_chunk_parser_normal_chunk_unaffected_by_error_check():
"""Test that normal streaming chunks still work correctly after error check addition"""
from unittest.mock import Mock
from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
ModelResponseIterator,
)
normal_chunk = {
"candidates": [
{
"content": {
"role": "model",
"parts": [{"text": "Hello"}],
},
"index": 0,
}
],
"usageMetadata": {
"promptTokenCount": 5,
"candidatesTokenCount": 1,
"totalTokenCount": 6,
},
}
logging_obj = Mock()
logging_obj.optional_params = {}
streaming_obj = ModelResponseIterator(
streaming_response=iter([]),
sync_stream=True,
logging_obj=logging_obj,
)
result = streaming_obj.chunk_parser(normal_chunk)
assert result is not None
assert len(result.choices) > 0
assert result.choices[0].delta.content == "Hello"
def test_chunk_parser_raises_on_non_dict_error():
"""Test chunk_parser raises VertexAIError when chunk['error'] is not a dict"""
from unittest.mock import Mock
from litellm.llms.vertex_ai.common_utils import VertexAIError
from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
ModelResponseIterator,
)
error_chunk = {"error": "something went wrong"}
logging_obj = Mock()
logging_obj.optional_params = {}
streaming_obj = ModelResponseIterator(
streaming_response=iter([]),
sync_stream=True,
logging_obj=logging_obj,
)
with pytest.raises(VertexAIError) as exc_info:
streaming_obj.chunk_parser(error_chunk)
assert exc_info.value.status_code == 500
assert "Unexpected error format" in exc_info.value.message
def test_chunk_parser_raises_on_string_error_code():
"""Test chunk_parser correctly converts string error code to int"""
from unittest.mock import Mock
from litellm.llms.vertex_ai.common_utils import VertexAIError
from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
ModelResponseIterator,
)
# code field is a string "429" rather than an int
error_chunk = {
"error": {
"code": "429",
"message": "Resource exhausted.",
"status": "RESOURCE_EXHAUSTED",
}
}
logging_obj = Mock()
logging_obj.optional_params = {}
streaming_obj = ModelResponseIterator(
streaming_response=iter([]),
sync_stream=True,
logging_obj=logging_obj,
)
with pytest.raises(VertexAIError) as exc_info:
streaming_obj.chunk_parser(error_chunk)
assert exc_info.value.status_code == 429
assert isinstance(exc_info.value.status_code, int)
def test_chunk_parser_error_chunk_explicit_null_code_uses_500():
"""JSON null for code must not call int(None); status defaults to 500."""
from unittest.mock import Mock
from litellm.llms.vertex_ai.common_utils import VertexAIError
from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
ModelResponseIterator,
)
error_chunk = {
"error": {
"code": None,
"message": "Something went wrong.",
"status": "UNKNOWN",
}
}
logging_obj = Mock()
logging_obj.optional_params = {}
streaming_obj = ModelResponseIterator(
streaming_response=iter([]),
sync_stream=True,
logging_obj=logging_obj,
)
with pytest.raises(VertexAIError) as exc_info:
streaming_obj.chunk_parser(error_chunk)
assert exc_info.value.status_code == 500
assert "Something went wrong" in exc_info.value.message
def test_chunk_parser_error_chunk_non_numeric_code_defaults_to_500():
"""Non-numeric code must not become ValueError -> RuntimeError in __next__."""
from unittest.mock import Mock
from litellm.llms.vertex_ai.common_utils import VertexAIError
from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
ModelResponseIterator,
)
error_chunk = {
"error": {
"code": "NOT_A_NUMBER",
"message": "Malformed.",
"status": "INVALID",
}
}
logging_obj = Mock()
logging_obj.optional_params = {}
streaming_obj = ModelResponseIterator(
streaming_response=iter([]),
sync_stream=True,
logging_obj=logging_obj,
)
with pytest.raises(VertexAIError) as exc_info:
streaming_obj.chunk_parser(error_chunk)
assert exc_info.value.status_code == 500
assert "Malformed" in exc_info.value.message
def test_chunk_parser_error_chunk_empty_dict_defaults_to_500():
"""Empty error object {} uses default code 500 and default message/status strings."""
from unittest.mock import Mock
from litellm.llms.vertex_ai.common_utils import VertexAIError
from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
ModelResponseIterator,
)
error_chunk = {"error": {}}
logging_obj = Mock()
logging_obj.optional_params = {}
streaming_obj = ModelResponseIterator(
streaming_response=iter([]),
sync_stream=True,
logging_obj=logging_obj,
)
with pytest.raises(VertexAIError) as exc_info:
streaming_obj.chunk_parser(error_chunk)
assert exc_info.value.status_code == 500
assert "UNKNOWN" in exc_info.value.message
assert "Unknown error" in exc_info.value.message
def test_chunk_parser_error_chunk_non_dict_int_value():
"""Non-dict error payloads (e.g. bare JSON number) must raise with status 500, not TypeError."""
from unittest.mock import Mock
from litellm.llms.vertex_ai.common_utils import VertexAIError
from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
ModelResponseIterator,
)
error_chunk = {"error": 503}
logging_obj = Mock()
logging_obj.optional_params = {}
streaming_obj = ModelResponseIterator(
streaming_response=iter([]),
sync_stream=True,
logging_obj=logging_obj,
)
with pytest.raises(VertexAIError) as exc_info:
streaming_obj.chunk_parser(error_chunk)
assert exc_info.value.status_code == 500
assert "Unexpected error format" in exc_info.value.message
assert "503" in exc_info.value.message
def test_chunk_parser_error_chunk_non_dict_null_value():
"""JSON null for error must hit the non-dict branch (same as int/string)."""
from unittest.mock import Mock
from litellm.llms.vertex_ai.common_utils import VertexAIError
from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
ModelResponseIterator,
)
error_chunk = {"error": None}
logging_obj = Mock()
logging_obj.optional_params = {}
streaming_obj = ModelResponseIterator(
streaming_response=iter([]),
sync_stream=True,
logging_obj=logging_obj,
)
with pytest.raises(VertexAIError) as exc_info:
streaming_obj.chunk_parser(error_chunk)
assert exc_info.value.status_code == 500
assert "Unexpected error format" in exc_info.value.message
def test_mid_stream_429_error_raises_during_iteration():
"""
Simulate a full streaming scenario: normal thinking chunks arrive first,
then a 429 RESOURCE_EXHAUSTED error chunk arrives mid-stream.
Verify that ModelResponseIterator raises VertexAIError during iteration.
"""
import json
from unittest.mock import Mock
from litellm.llms.vertex_ai.common_utils import VertexAIError
from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
ModelResponseIterator,
)
# Simulate Vertex AI SSE stream: normal chunks followed by a 429 error chunk
normal_chunk_1 = json.dumps(
{
"candidates": [
{
"content": {
"role": "model",
"parts": [{"text": "Let me think about this...", "thought": True}],
},
"index": 0,
}
],
"usageMetadata": {
"promptTokenCount": 10,
"candidatesTokenCount": 5,
"totalTokenCount": 15,
},
"modelVersion": "gemini-3.1-flash-image-preview",
}
)
normal_chunk_2 = json.dumps(
{
"candidates": [
{
"content": {
"role": "model",
"parts": [{"text": "I'll generate the image now.", "thought": True}],
},
"index": 0,
}
],
"usageMetadata": {
"promptTokenCount": 10,
"candidatesTokenCount": 12,
"totalTokenCount": 22,
},
}
)
error_chunk = json.dumps(
{
"error": {
"code": 429,
"message": "Resource exhausted. Please try again later. Please refer to https://cloud.google.com/vertex-ai/generative-ai/docs/error-code-429 for more details.",
"status": "RESOURCE_EXHAUSTED",
}
}
)
# Build a mock SSE stream (lines returned by iter_lines)
sse_lines = iter([normal_chunk_1, normal_chunk_2, error_chunk])
logging_obj = Mock()
logging_obj.optional_params = {}
streaming_obj = ModelResponseIterator(
streaming_response=sse_lines,
sync_stream=True,
logging_obj=logging_obj,
)
# Iterate the stream: first chunks should succeed, then 429 error should be raised
results = []
with pytest.raises(VertexAIError) as exc_info:
for chunk in streaming_obj:
if chunk is not None:
results.append(chunk)
# Verify: received normal chunks before the error
assert (
len(results) >= 1
), "Should have received at least 1 normal chunk before the error"
# Verify: 429 error is properly raised
assert exc_info.value.status_code == 429
assert "RESOURCE_EXHAUSTED" in str(exc_info.value.message)

View file

@ -1219,6 +1219,32 @@ def test_process_gemini_media():
mime_type="image/jpeg", file_uri="gs://bucket/image"
)
# Test gs url without extension using mime_type from image_url object
image_message = [
{
"role": "user",
"content": [
{
"type": "image_url",
"image_url": {
"url": "gs://bucket/image-without-extension",
"mime_type": "image/png",
},
}
],
}
]
from litellm.llms.vertex_ai.gemini.transformation import (
_gemini_convert_messages_with_history,
)
converted = _gemini_convert_messages_with_history(
messages=image_message, model="gemini-2.5-flash"
)
assert converted[0]["parts"][0]["file_data"] == FileDataType(
mime_type="image/png", file_uri="gs://bucket/image-without-extension"
)
# Test HTTPS JPG URL
https_result = _process_gemini_media("https://example.com/image.jpg")
print("https_result JPG", https_result)
@ -1256,6 +1282,7 @@ def test_process_gemini_media():
assert base64_result["inline_data"]["data"] == "/9j/4AAQSkZJRg..."
def test_get_image_mime_type_from_url():
"""Test the _get_image_mime_type_from_url function for different image URLs"""
from litellm.llms.vertex_ai.gemini.transformation import (

View file

@ -0,0 +1,466 @@
"""Vertex Gemini: extensionless gs:// MIME + GCS metadata tests.
Split from test_vertex.py to satisfy CI per-file size limits.
"""
import asyncio
import os
import sys
import time
from dotenv import load_dotenv
load_dotenv()
import pytest
import litellm
from unittest.mock import MagicMock, patch
sys.path.insert(0, os.path.abspath("../.."))
from litellm.llms.vertex_ai.gemini.transformation import _process_gemini_media
def test_process_gemini_media_gcs_explicit_format_octet_stream_and_alias():
"""Explicit format bypasses registry; image/jpg alias still applies."""
from litellm.types.llms.vertex_ai import FileDataType
r1 = _process_gemini_media(
"gs://bucket/object-no-ext",
format="application/octet-stream",
)
assert r1["file_data"] == FileDataType(
mime_type="application/octet-stream",
file_uri="gs://bucket/object-no-ext",
)
r2 = _process_gemini_media("gs://bucket/object-no-ext", format="image/jpg")
assert r2["file_data"] == FileDataType(
mime_type="image/jpeg",
file_uri="gs://bucket/object-no-ext",
)
def test_process_gemini_media_gcs_without_extension_errors_and_metadata_mock():
with patch(
"litellm.llms.vertex_ai.gemini.transformation._get_gcs_object_content_type",
return_value=None,
):
with pytest.raises(litellm.BadRequestError) as exc:
_process_gemini_media("gs://bucket/image-without-extension")
assert "Unable to determine mime type for gs URI" in str(exc.value)
from litellm.types.llms.vertex_ai import FileDataType
with patch(
"litellm.llms.vertex_ai.gemini.transformation._get_gcs_object_content_type",
return_value="image/jpeg",
) as m:
r = _process_gemini_media("gs://bucket/image-without-extension")
assert r["file_data"] == FileDataType(
mime_type="image/jpeg", file_uri="gs://bucket/image-without-extension"
)
m.assert_called()
with patch(
"litellm.llms.vertex_ai.gemini.transformation._get_gcs_object_content_type",
return_value="image/jpg",
):
r_alias = _process_gemini_media("gs://bucket/image-without-extension")
assert r_alias["file_data"]["mime_type"] == "image/jpeg"
def test_process_gemini_media_rejects_gcs_metadata_mime_not_supported_by_gemini():
"""Non-empty GCS contentType that fails _normalize_and_validate_gemini_mime_type."""
with patch(
"litellm.llms.vertex_ai.gemini.transformation._get_gcs_object_content_type",
return_value="application/x-litellm-unit-test-unknown-mime",
):
with pytest.raises(
litellm.BadRequestError,
match="File type not supported by gemini",
):
_process_gemini_media("gs://bucket/object-without-extension")
def test_file_block_uses_mime_type_alias_for_extensionless_gcs():
from litellm.llms.vertex_ai.gemini.transformation import (
_gemini_convert_messages_with_history,
)
from litellm.types.llms.vertex_ai import FileDataType
messages = [
{
"role": "user",
"content": [
{
"type": "file",
"file": {
"file_id": "gs://bucket/no-extension-object",
"mime_type": "application/pdf",
},
}
],
}
]
converted = _gemini_convert_messages_with_history(
messages=messages, model="gemini-2.5-flash"
)
assert converted[0]["parts"][0]["file_data"] == FileDataType(
mime_type="application/pdf", file_uri="gs://bucket/no-extension-object"
)
@pytest.mark.parametrize(
"bucket,expected",
[
(("a." * 110) + "aa", True),
("ab", False),
("a" * 64, False),
("ab..cd", False),
("1.2.3.4", False),
("192.168.0.1", False),
("Bucket-Upper", False),
("bucket@name", False),
("bucket name", False),
("-mybucket", False),
("mybucket-", False),
(".mybucket", False),
("mybucket.", False),
],
)
def test_is_valid_gcs_bucket_name_matrix(bucket, expected):
from litellm.llms.vertex_ai.gemini.transformation import _is_valid_gcs_bucket_name
assert _is_valid_gcs_bucket_name(bucket) is expected
def test_get_gcs_object_content_type_explicit_vertex_success_and_token_failure():
from litellm.llms.vertex_ai.gemini import transformation as gt
mock_v = MagicMock()
mock_v.get_access_token.return_value = ("test-token", "test-project")
resp = MagicMock()
resp.is_error = False
resp.status_code = 200
resp.json.return_value = {"contentType": "image/png"}
http = MagicMock()
http.get.return_value = resp
with (
patch.object(gt, "_GCS_METADATA_VERTEX_BASE", mock_v),
patch(
"litellm.llms.vertex_ai.gemini.transformation._get_gcs_metadata_http_handler",
return_value=http,
),
):
assert (
gt._get_gcs_object_content_type(
image_url="gs://my-bucket/path/to/image-without-extension",
vertex_project="project-123",
vertex_credentials="credential-json",
)
== "image/png"
)
mock_v.get_access_token.assert_called_once_with(
credentials="credential-json",
project_id="project-123",
)
mock_v2 = MagicMock()
mock_v2.get_access_token.side_effect = Exception("token failure")
with patch.object(gt, "_GCS_METADATA_VERTEX_BASE", mock_v2):
with pytest.raises(
litellm.BadRequestError,
match="Unable to fetch GCS metadata with provided Vertex credentials/project",
):
gt._get_gcs_object_content_type(
image_url="gs://my-bucket/path/to/image-without-extension",
vertex_project="project-123",
vertex_credentials="credential-json",
)
def test_get_gcs_object_content_type_http_error_explicit_vs_anonymous():
from litellm.llms.vertex_ai.gemini import transformation as gt
mock_v = MagicMock()
mock_v.get_access_token.return_value = ("t", "p")
err_resp = MagicMock()
err_resp.is_error = True
err_resp.status_code = 403
err_resp.text = '{"error":{"message":"Permission denied"}}'
http = MagicMock()
http.get.return_value = err_resp
with (
patch.object(gt, "_GCS_METADATA_VERTEX_BASE", mock_v),
patch(
"litellm.llms.vertex_ai.gemini.transformation._get_gcs_metadata_http_handler",
return_value=http,
),
):
with pytest.raises(litellm.BadRequestError, match="HTTP 403") as ei:
gt._get_gcs_object_content_type(
image_url="gs://my-bucket/path/to/obj",
vertex_project="project-123",
vertex_credentials="credential-json",
)
assert "Permission denied" in str(ei.value)
mock_v2 = MagicMock()
anon_err = MagicMock()
anon_err.is_error = True
anon_err.status_code = 403
anon_err.text = "Forbidden"
http2 = MagicMock()
http2.get.return_value = anon_err
with (
patch.object(gt, "_GCS_METADATA_VERTEX_BASE", mock_v2),
patch(
"litellm.llms.vertex_ai.gemini.transformation._get_gcs_metadata_http_handler",
return_value=http2,
),
):
assert (
gt._get_gcs_object_content_type(image_url="gs://public-bucket/public-object")
is None
)
mock_v2.get_access_token.assert_not_called()
def test_get_gcs_object_content_type_anonymous_success_no_auth_header():
from litellm.llms.vertex_ai.gemini import transformation as gt
mock_v = MagicMock()
ok = MagicMock()
ok.is_error = False
ok.status_code = 200
ok.json.return_value = {"contentType": "image/jpeg"}
http = MagicMock()
http.get.return_value = ok
with (
patch.object(gt, "_GCS_METADATA_VERTEX_BASE", mock_v),
patch(
"litellm.llms.vertex_ai.gemini.transformation._get_gcs_metadata_http_handler",
return_value=http,
),
):
assert (
gt._get_gcs_object_content_type(image_url="gs://public-bucket/public-object")
== "image/jpeg"
)
mock_v.get_access_token.assert_not_called()
hdrs = http.get.call_args.kwargs.get("headers")
assert hdrs is None or "Authorization" not in hdrs
def test_async_transform_request_body_offloads_extensionless_gs_not_plain_text():
from litellm.llms.vertex_ai.gemini import transformation as gemini_transformation
messages = [
{
"role": "user",
"content": [
{
"type": "image_url",
"image_url": {"url": "gs://bucket/image-without-extension"},
}
],
}
]
def slow_http_get(*args, **kwargs):
time.sleep(0.5)
response = MagicMock()
response.is_error = False
response.status_code = 200
response.raise_for_status.return_value = None
response.json.return_value = {"contentType": "image/png"}
return response
async def fake_check_and_create_cache(self, **kwargs):
return kwargs["messages"], kwargs["optional_params"], None
mock_v = MagicMock()
mock_v.get_access_token.return_value = ("token", "project")
mock_http = MagicMock()
mock_http.get.side_effect = slow_http_get
async def run_scenario() -> float:
async def concurrent_sleep() -> float:
start = time.monotonic()
await asyncio.sleep(0.05)
return time.monotonic() - start
task = asyncio.create_task(
gemini_transformation.async_transform_request_body(
gemini_api_key=None,
messages=messages,
api_base=None,
model="gemini-2.5-flash",
client=None,
timeout=None,
extra_headers=None,
optional_params={},
logging_obj=MagicMock(),
custom_llm_provider="vertex_ai",
litellm_params={},
vertex_project=None,
vertex_location=None,
vertex_auth_header=None,
)
)
elapsed = await concurrent_sleep()
await task
return elapsed
with (
patch.object(gemini_transformation, "_GCS_METADATA_VERTEX_BASE", mock_v),
patch(
"litellm.llms.vertex_ai.gemini.transformation._get_gcs_metadata_http_handler",
return_value=mock_http,
),
patch(
"litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching."
"ContextCachingEndpoints.async_check_and_create_cache",
new=fake_check_and_create_cache,
),
):
sleep_elapsed = asyncio.run(run_scenario())
assert sleep_elapsed < 0.4, (
f"Event loop blocked for {sleep_elapsed:.3f}s; "
"async_transform_request_body did not offload sync GCS metadata"
)
async def fake_cache2(self, **kwargs):
return kwargs["messages"], kwargs["optional_params"], None
async def run_plain():
with patch(
"litellm.llms.vertex_ai.gemini.transformation.asyncify",
side_effect=AssertionError("asyncify must not run without extensionless gs://"),
):
return await gemini_transformation.async_transform_request_body(
gemini_api_key=None,
messages=[{"role": "user", "content": "hello"}],
api_base=None,
model="gemini-2.5-flash",
client=None,
timeout=None,
extra_headers=None,
optional_params={},
logging_obj=MagicMock(),
custom_llm_provider="vertex_ai",
litellm_params={},
vertex_project=None,
vertex_location=None,
vertex_auth_header=None,
)
with patch(
"litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching."
"ContextCachingEndpoints.async_check_and_create_cache",
new=fake_cache2,
):
body = asyncio.run(run_plain())
assert body is not None and "contents" in body
@pytest.mark.parametrize(
"messages,expected",
[
([{"role": "user", "content": "hello"}], False),
(
[
{
"role": "user",
"content": [
{
"type": "image_url",
"image_url": {"url": "gs://bucket/image-without-extension"},
}
],
}
],
True,
),
(
[
{
"role": "user",
"content": [
{
"type": "image_url",
"image_url": {"url": "gs://bucket/image.png"},
}
],
}
],
False,
),
(
[
{
"role": "user",
"content": [
{
"type": "image_url",
"image_url": {
"url": "gs://bucket/image-without-extension",
"mime_type": "image/png",
},
}
],
}
],
False,
),
(
[
{
"role": "assistant",
"content": [],
"images": [
{"image_url": {"url": "gs://bucket/gen-without-extension"}},
],
}
],
True,
),
(
[
{
"role": "assistant",
"content": [],
"images": [{"image_url": {"url": "gs://bucket/gen.png"}}],
}
],
False,
),
(
[
{
"role": "assistant",
"content": [],
"images": [
{
"image_url": {
"url": "gs://bucket/gen-no-ext",
"mime_type": "image/png",
},
}
],
}
],
False,
),
],
)
def test_openai_messages_may_need_sync_gcs_metadata_fetch_matrix(messages, expected):
from litellm.llms.vertex_ai.gemini.transformation import (
_openai_messages_may_need_sync_gcs_metadata_fetch,
)
assert _openai_messages_may_need_sync_gcs_metadata_fetch(messages) is expected

View file

@ -0,0 +1,121 @@
"""
Regression tests for #28084:
`VertexAIPartnerModels.count_tokens` (for Claude / Mistral / Llama on Vertex)
used to gate on `import vertexai` even though the actual count-tokens path goes
through `VertexAIPartnerModelsTokenCounter.handle_count_tokens_request`, which
talks to the publisher's `:rawPredict` endpoint over plain httpx and never
touches the Gemini SDK. The unused gate broke `/v1/messages/count_tokens` for
any LiteLLM install that did not pull in `google-cloud-aiplatform` (which is
not in the default `proxy` / `proxy-dev` extras).
These tests pin the absence of that gate by:
1. simulating `vertexai` being unimportable and verifying the partner-model
path does not raise the historical "vertexai import failed" error before
reaching the network/auth layer, and
2. asserting that import of the partner-model count-tokens handler module by
itself does not pull `vertexai` into `sys.modules`.
"""
import sys
import pytest
from litellm.llms.vertex_ai.vertex_ai_partner_models.count_tokens.handler import (
VertexAIPartnerModelsTokenCounter,
)
from litellm.llms.vertex_ai.vertex_ai_partner_models.main import VertexAIPartnerModels
@pytest.mark.asyncio
async def test_count_tokens_does_not_require_vertexai_sdk(monkeypatch):
"""Even when `import vertexai` would fail, count_tokens must not raise the
historical "vertexai import failed" gate. The downstream handler talks to
`:rawPredict` over httpx with an access token no Gemini SDK needed."""
# Simulate `vertexai` being unimportable, regardless of what is actually on
# the test environment's sys.path.
monkeypatch.setitem(sys.modules, "vertexai", None)
monkeypatch.setitem(sys.modules, "vertexai.preview", None)
captured = {}
async def fake_ensure_access_token(
self, credentials, project_id, custom_llm_provider
):
return "fake-token", "fake-project"
def fake_build_endpoint(self, model, project_id, vertex_location, api_base=None):
captured["model_to_endpoint"] = model
return "https://fake-endpoint"
monkeypatch.setattr(
VertexAIPartnerModelsTokenCounter,
"_ensure_access_token_async",
fake_ensure_access_token,
)
monkeypatch.setattr(
VertexAIPartnerModelsTokenCounter,
"_build_count_tokens_endpoint",
fake_build_endpoint,
)
class FakeResponse:
status_code = 200
def json(self):
return {"input_tokens": 9}
class FakeClient:
async def post(self, url, headers=None, json=None, **kwargs):
captured["url"] = url
captured["headers"] = headers
captured["json"] = json
return FakeResponse()
import litellm.llms.vertex_ai.vertex_ai_partner_models.count_tokens.handler as handler_mod
monkeypatch.setattr(
handler_mod, "get_async_httpx_client", lambda **kwargs: FakeClient()
)
result = await VertexAIPartnerModels().count_tokens(
model="claude-sonnet-4-6",
messages=[{"role": "user", "content": "hello"}],
litellm_params={"vertex_location": "us-east5"},
vertex_project="test-project",
vertex_location="us-east5",
vertex_credentials=None,
)
# We should reach the publisher endpoint and parse its response, not raise
# the vertexai-import gate.
assert result == {
"input_tokens": 9,
"tokenizer_used": "vertex_ai_partner_models",
}
assert captured["headers"] == {"Authorization": "Bearer fake-token"}
assert captured["model_to_endpoint"] == "claude-sonnet-4-6"
def test_handler_module_does_not_import_vertexai_sdk():
"""Importing the partner-model count-tokens handler must not load the
Gemini SDK into sys.modules. Operators who only need Claude-on-Vertex
token counting should not pay for `google-cloud-aiplatform`."""
# Force-evict any prior load so this assertion measures what THIS module
# pulls in, not what an unrelated earlier test did.
for mod in list(sys.modules):
if mod == "vertexai" or mod.startswith("vertexai."):
sys.modules.pop(mod, None)
# Re-import the handler module to verify it stays SDK-free.
import importlib
import litellm.llms.vertex_ai.vertex_ai_partner_models.count_tokens.handler as handler_mod
importlib.reload(handler_mod)
leaked = [m for m in sys.modules if m == "vertexai" or m.startswith("vertexai.")]
assert leaked == [], f"unexpected vertexai SDK imports: {leaked}"

View file

@ -553,6 +553,7 @@ class TestGuardrailActions:
# Verify the exception has the clean error message (no wrapper)
assert str(exc_info.value) == "Content contains harmful instructions"
assert exc_info.value.guardrail_name == "generic_guardrail_api"
assert exc_info.value.status_code == 400
@pytest.mark.asyncio
async def test_action_intervened_modifies_content(

View file

@ -1446,3 +1446,175 @@ class TestGetTeamDeployments:
result = await _get_team_deployments(team_id, prisma_client)
assert len(result) == 1
assert result[0] is dep1
def _build_db_model_for_blocked_test():
from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo
return Deployment(
model_name="gpt-4o",
litellm_params=LiteLLM_Params(model="openai/gpt-4o"),
model_info=ModelInfo(id="dep-0"),
)
class TestUpdateDBModelBlocked:
"""`update_db_model` must thread `blocked` through to the Prisma payload only
when the caller explicitly set it PATCH semantics: an absent field means
"leave the stored value untouched"."""
def test_update_db_model_passes_blocked_true_to_db(self):
from litellm.proxy.management_endpoints.model_management_endpoints import (
update_db_model,
)
result = update_db_model(
db_model=_build_db_model_for_blocked_test(),
updated_patch=updateDeployment(blocked=True),
)
assert result["blocked"] is True
def test_update_db_model_passes_blocked_false_to_db(self):
from litellm.proxy.management_endpoints.model_management_endpoints import (
update_db_model,
)
result = update_db_model(
db_model=_build_db_model_for_blocked_test(),
updated_patch=updateDeployment(blocked=False),
)
assert result["blocked"] is False
def test_update_db_model_omits_blocked_when_patch_is_none(self):
from litellm.proxy.management_endpoints.model_management_endpoints import (
update_db_model,
)
result = update_db_model(
db_model=_build_db_model_for_blocked_test(),
updated_patch=updateDeployment(),
)
assert "blocked" not in result
class TestGetModelInfoWithIdBlocked:
"""`ProxyConfig.get_model_info_with_id` must propagate the DB-level `blocked`
column into the in-memory `model_info` dict so the router filter can read it."""
def test_get_model_info_with_id_propagates_blocked_true(self):
from litellm.proxy.proxy_server import ProxyConfig
model = MagicMock()
model.model_id = "dep-1"
model.model_info = {}
model.blocked = True
info = ProxyConfig().get_model_info_with_id(model=model, db_model=True)
assert info.id == "dep-1"
assert getattr(info, "blocked") is True
def test_get_model_info_with_id_defaults_blocked_to_false_when_missing(self):
from litellm.proxy.proxy_server import ProxyConfig
model = MagicMock(spec=["model_id", "model_info"])
model.model_id = "dep-2"
model.model_info = {}
info = ProxyConfig().get_model_info_with_id(model=model, db_model=True)
assert getattr(info, "blocked") is False
class TestPatchModelBlockedAuthGate:
"""Only proxy admins may flip `blocked` — team admins authorized for
team-scoped models via `can_user_make_model_call` must still be rejected
when they attempt to toggle the pause flag."""
@pytest.mark.asyncio
async def test_team_admin_cannot_toggle_blocked(self):
from litellm.proxy.management_endpoints.model_management_endpoints import (
patch_model,
)
non_admin = UserAPIKeyAuth(
user_id="team_admin",
user_role=LitellmUserRoles.INTERNAL_USER,
)
existing_row = MagicMock()
existing_row.litellm_params = {"model": "openai/gpt-4o-mini"}
existing_row.model_dump.return_value = {
"model_name": "gpt-4o-mini",
"litellm_params": existing_row.litellm_params,
"model_info": {"id": "m1"},
}
existing_row.model_dump_json.return_value = "{}"
mock_prisma = MagicMock()
mock_prisma.db.litellm_proxymodeltable.find_unique = AsyncMock(
return_value=existing_row
)
with (
patch("litellm.proxy.proxy_server.prisma_client", mock_prisma),
patch("litellm.proxy.proxy_server.llm_router", MagicMock()),
patch("litellm.proxy.proxy_server.store_model_in_db", True),
patch("litellm.proxy.proxy_server.premium_user", True),
patch(
"litellm.proxy.management_endpoints.model_management_endpoints.ModelManagementAuthChecks.can_user_make_model_call",
new=AsyncMock(return_value=None),
),
):
with pytest.raises(Exception) as exc_info:
await patch_model(
model_id="m1",
patch_data=updateDeployment(blocked=True),
user_api_key_dict=non_admin,
)
err = exc_info.value
assert getattr(err, "param", "") == "blocked"
assert "proxy admin" in getattr(err, "message", "").lower()
@pytest.mark.asyncio
async def test_proxy_admin_can_toggle_blocked(self):
from litellm.proxy.management_endpoints.model_management_endpoints import (
patch_model,
)
admin = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN)
existing_row = MagicMock()
existing_row.litellm_params = {"model": "openai/gpt-4o-mini"}
existing_row.model_dump.return_value = {
"model_name": "gpt-4o-mini",
"litellm_params": existing_row.litellm_params,
"model_info": {"id": "m1"},
}
existing_row.model_dump_json.return_value = "{}"
updated_row = MagicMock()
updated_row.model_dump_json.return_value = "{}"
mock_prisma = MagicMock()
mock_prisma.db.litellm_proxymodeltable.find_unique = AsyncMock(
return_value=existing_row
)
mock_prisma.db.litellm_proxymodeltable.update = AsyncMock(
return_value=updated_row
)
with (
patch("litellm.proxy.proxy_server.prisma_client", mock_prisma),
patch("litellm.proxy.proxy_server.llm_router", MagicMock()),
patch("litellm.proxy.proxy_server.store_model_in_db", True),
patch("litellm.proxy.proxy_server.premium_user", True),
patch(
"litellm.proxy.management_endpoints.model_management_endpoints.ModelManagementAuthChecks.can_user_make_model_call",
new=AsyncMock(return_value=None),
),
patch(
"litellm.proxy.management_endpoints.model_management_endpoints.clear_cache",
new=AsyncMock(return_value=None),
),
):
result = await patch_model(
model_id="m1",
patch_data=updateDeployment(blocked=True),
user_api_key_dict=admin,
)
assert result is updated_row
mock_prisma.db.litellm_proxymodeltable.update.assert_awaited_once()

View file

@ -7943,3 +7943,103 @@ async def test_team_member_me_returns_404_for_unknown_team(mock_db_client):
user_api_key_dict=caller_auth,
)
assert exc_info.value.status_code == 404
def _non_admin_auth():
return UserAPIKeyAuth(
user_id="u-team-admin", user_role=LitellmUserRoles.INTERNAL_USER
)
def test_check_passthrough_routes_caller_permission_team():
from litellm.proxy._types import NewTeamRequest
from litellm.proxy.management_endpoints.common_utils import (
_check_passthrough_routes_caller_permission,
)
admin = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN)
non_admin = _non_admin_auth()
_check_passthrough_routes_caller_permission(
NewTeamRequest(allowed_passthrough_routes=["/foo/*"]), admin, entity="team"
)
_check_passthrough_routes_caller_permission(
NewTeamRequest(), non_admin, entity="team"
)
_check_passthrough_routes_caller_permission(
NewTeamRequest(allowed_passthrough_routes=[]), non_admin, entity="team"
)
with pytest.raises(HTTPException) as exc:
_check_passthrough_routes_caller_permission(
NewTeamRequest(allowed_passthrough_routes=["/admin/*"]),
non_admin,
entity="team",
)
assert exc.value.status_code == 403
assert "allowed_passthrough_routes" in str(exc.value.detail)
assert "team" in str(exc.value.detail)
with pytest.raises(HTTPException) as exc:
_check_passthrough_routes_caller_permission(
NewTeamRequest(metadata={"allowed_passthrough_routes": ["/admin/*"]}),
non_admin,
entity="team",
)
assert exc.value.status_code == 403
assert "metadata.allowed_passthrough_routes" in str(exc.value.detail)
@pytest.mark.asyncio
async def test_new_team_blocks_non_admin_passthrough_routes(mock_db_client):
"""A non-proxy-admin cannot self-grant pass-through routes via /team/new."""
mock_db_client.db.litellm_teamtable.count = AsyncMock(return_value=0)
from fastapi import Request
from litellm.proxy._types import NewTeamRequest, ProxyException
from litellm.proxy.management_endpoints.team_endpoints import new_team
with patch(
"litellm.proxy.management_endpoints.team_endpoints._check_user_team_limits",
AsyncMock(return_value=None),
):
with pytest.raises(ProxyException) as exc:
await new_team(
data=NewTeamRequest(
team_alias="t", allowed_passthrough_routes=["/admin/*"]
),
http_request=MagicMock(spec=Request),
user_api_key_dict=_non_admin_auth(),
)
assert str(exc.value.code) == "403"
assert "allowed_passthrough_routes" in str(exc.value.message)
@pytest.mark.asyncio
async def test_update_team_blocks_non_admin_passthrough_routes(mock_db_client):
"""Even a team manager (non-proxy-admin) cannot set pass-through routes via
/team/update the gate runs after _verify_team_access."""
from fastapi import Request
from litellm.proxy._types import ProxyException, UpdateTeamRequest
from litellm.proxy.management_endpoints.team_endpoints import update_team
existing = MagicMock()
existing.model_dump.return_value = {"team_id": "t1"}
mock_db_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=existing)
with patch(
"litellm.proxy.management_endpoints.team_endpoints._verify_team_access",
AsyncMock(return_value=None),
):
with pytest.raises(ProxyException) as exc:
await update_team(
data=UpdateTeamRequest(
team_id="t1", allowed_passthrough_routes=["/admin/*"]
),
http_request=MagicMock(spec=Request),
user_api_key_dict=_non_admin_auth(),
)
assert str(exc.value.code) == "403"
assert "allowed_passthrough_routes" in str(exc.value.message)

View file

@ -2057,3 +2057,317 @@ def test_openrouter_gemini_3_1_flash_lite_preview_pricing():
assert model_info["output_cost_per_token"] == 1.5e-06
assert model_info["max_input_tokens"] == 1048576
assert model_info["max_output_tokens"] == 65536
def test_custom_pricing_applies_cache_read_input_cost():
"""
Bug 1 reproduction: custom_cost_per_token with cache_read_input_token_cost
should bill cached prompt tokens at the cache rate, not the full input rate.
"""
usage = Usage(
prompt_tokens=6074,
completion_tokens=285,
total_tokens=6359,
prompt_tokens_details=PromptTokensDetailsWrapper(
cached_tokens=3456,
audio_tokens=0,
),
)
response = ModelResponse(
id="test-id",
created=1234567890,
model="openai/gpt-5.4",
object="chat.completion",
choices=[],
usage=usage,
)
cost = litellm.completion_cost(
completion_response=response,
model="openai/gpt-5.4",
custom_llm_provider="openai",
custom_cost_per_token={
"input_cost_per_token": 0.0000025,
"output_cost_per_token": 0.000015,
"cache_read_input_token_cost": 0.00000025,
},
)
expected = (6074 - 3456) * 0.0000025 + 3456 * 0.00000025 + 285 * 0.000015
assert cost == pytest.approx(expected)
def test_custom_pricing_applies_cache_creation_input_cost_via_prompt_details():
"""
OpenAI-compatible providers report cache-write tokens under
prompt_tokens_details.cache_creation_tokens. The custom-pricing helper must
bill those at cache_creation_input_token_cost, not the full input rate.
"""
pt_details = PromptTokensDetailsWrapper(cached_tokens=1000, audio_tokens=0)
pt_details.cache_creation_tokens = 500
usage = Usage(
prompt_tokens=4000,
completion_tokens=100,
total_tokens=4100,
prompt_tokens_details=pt_details,
)
response = ModelResponse(
id="test-id",
created=1234567890,
model="openai/gpt-5.4",
object="chat.completion",
choices=[],
usage=usage,
)
cost = litellm.completion_cost(
completion_response=response,
model="openai/gpt-5.4",
custom_llm_provider="openai",
custom_cost_per_token={
"input_cost_per_token": 0.0000025,
"output_cost_per_token": 0.000015,
"cache_read_input_token_cost": 0.00000025,
"cache_creation_input_token_cost": 0.000003125,
},
)
expected = (
(4000 - 1000 - 500) * 0.0000025
+ 1000 * 0.00000025
+ 500 * 0.000003125
+ 100 * 0.000015
)
assert cost == pytest.approx(expected)
def test_custom_pricing_applies_cache_creation_input_cost_via_cache_write_tokens_alias():
"""
Some OpenAI-compatible providers (e.g. kimi-k2) emit cache-write tokens as
`cache_write_tokens` rather than `cache_creation_tokens`. The cost
calculator must mirror db_spend_update_writer and accept either name
otherwise daily aggregation counts the tokens but the per-request cost
bills them at the full input rate.
Drives `cost_per_token` directly with a SimpleNamespace usage stub so the
`cache_write_tokens` alias survives the call (Pydantic's Usage init
rebuilds prompt_tokens_details and drops dynamic attributes).
"""
from types import SimpleNamespace
from litellm.cost_calculator import cost_per_token
pt_details = SimpleNamespace(cached_tokens=1000, cache_write_tokens=500)
usage_stub = SimpleNamespace(
prompt_tokens=4000,
completion_tokens=100,
total_tokens=4100,
prompt_tokens_details=pt_details,
cache_read_input_tokens=None,
cache_creation_input_tokens=None,
)
prompt_cost, completion_cost = cost_per_token(
model="moonshotai/kimi-k2",
prompt_tokens=4000,
completion_tokens=100,
custom_llm_provider="openai",
usage_object=usage_stub,
custom_cost_per_token={
"input_cost_per_token": 0.0000025,
"output_cost_per_token": 0.000015,
"cache_read_input_token_cost": 0.00000025,
"cache_creation_input_token_cost": 0.000003125,
},
)
expected_prompt = (
(4000 - 1000 - 500) * 0.0000025 + 1000 * 0.00000025 + 500 * 0.000003125
)
expected_completion = 100 * 0.000015
assert prompt_cost == pytest.approx(expected_prompt)
assert completion_cost == pytest.approx(expected_completion)
# ---------------------------------------------------------------------------
# Bug 2 — db_spend_update_writer cache token extraction helpers.
# ---------------------------------------------------------------------------
def test_extract_cache_read_tokens_anthropic_top_level():
from litellm.proxy.db.db_spend_update_writer import _extract_cache_read_tokens
usage_obj = {
"prompt_tokens": 100,
"cache_read_input_tokens": 80,
"prompt_tokens_details": {"cached_tokens": 80},
}
# Anthropic top-level value should win over prompt_tokens_details fallback.
assert _extract_cache_read_tokens(usage_obj) == 80
def test_extract_cache_read_tokens_openai_compatible_fallback():
from litellm.proxy.db.db_spend_update_writer import _extract_cache_read_tokens
# Anthropic field absent — fall back to prompt_tokens_details.cached_tokens.
usage_obj = {
"prompt_tokens": 22583,
"prompt_tokens_details": {"cached_tokens": 22016},
}
assert _extract_cache_read_tokens(usage_obj) == 22016
def test_extract_cache_read_tokens_zero_when_missing():
from litellm.proxy.db.db_spend_update_writer import _extract_cache_read_tokens
assert _extract_cache_read_tokens({}) == 0
assert _extract_cache_read_tokens({"cache_read_input_tokens": None}) == 0
assert (
_extract_cache_read_tokens({"prompt_tokens_details": {"cached_tokens": None}})
== 0
)
def test_extract_cache_creation_tokens_anthropic_top_level():
from litellm.proxy.db.db_spend_update_writer import (
_extract_cache_creation_tokens,
)
usage_obj = {
"prompt_tokens": 100,
"cache_creation_input_tokens": 50,
"prompt_tokens_details": {"cache_write_tokens": 50},
}
# Anthropic top-level should short-circuit the fallback.
assert _extract_cache_creation_tokens(usage_obj) == 50
def test_extract_cache_creation_tokens_openai_cache_write_alias():
from litellm.proxy.db.db_spend_update_writer import (
_extract_cache_creation_tokens,
)
# kimi-k2 emits cache_write_tokens.
usage_obj = {
"prompt_tokens": 1000,
"prompt_tokens_details": {"cache_write_tokens": 200},
}
assert _extract_cache_creation_tokens(usage_obj) == 200
def test_extract_cache_creation_tokens_openai_cache_creation_alias():
from litellm.proxy.db.db_spend_update_writer import (
_extract_cache_creation_tokens,
)
# Other OpenAI-compatible providers emit cache_creation_tokens.
usage_obj = {
"prompt_tokens": 1000,
"prompt_tokens_details": {"cache_creation_tokens": 300},
}
assert _extract_cache_creation_tokens(usage_obj) == 300
def test_extract_cache_creation_tokens_zero_when_missing():
from litellm.proxy.db.db_spend_update_writer import (
_extract_cache_creation_tokens,
)
assert _extract_cache_creation_tokens({}) == 0
assert _extract_cache_creation_tokens({"cache_creation_input_tokens": None}) == 0
assert (
_extract_cache_creation_tokens(
{"prompt_tokens_details": {"cache_write_tokens": None}}
)
== 0
)
def test_custom_pricing_anthropic_style_cache_tokens_not_double_counted():
"""
Anthropic providers report cache tokens at the top level of Usage, and
`prompt_tokens` EXCLUDES them. The helper expects `prompt_tokens` to
include cache tokens, so cost_per_token must adjust before invoking it
otherwise regular_prompt_tokens goes negative and clamps to 0.
"""
usage = Usage(
prompt_tokens=2000,
completion_tokens=100,
total_tokens=2100,
cache_read_input_tokens=1500,
cache_creation_input_tokens=300,
)
response = ModelResponse(
id="test-id",
created=1234567890,
model="anthropic/claude-3-5-sonnet",
object="chat.completion",
choices=[],
usage=usage,
)
cost = litellm.completion_cost(
completion_response=response,
model="anthropic/claude-3-5-sonnet",
custom_llm_provider="anthropic",
custom_cost_per_token={
"input_cost_per_token": 0.000003,
"output_cost_per_token": 0.000015,
"cache_read_input_token_cost": 0.0000003,
"cache_creation_input_token_cost": 0.00000375,
},
)
# Anthropic prompt_tokens=2000 excludes cache. After normalization the
# helper sees 2000 + 1500 + 300 = 3800, of which 2000 are uncached.
expected = 2000 * 0.000003 + 1500 * 0.0000003 + 300 * 0.00000375 + 100 * 0.000015
assert cost == pytest.approx(expected)
def test_custom_pricing_without_cache_keys_preserves_legacy_behavior():
"""
Backward compatibility: when custom_cost_per_token omits both cache rates,
cached tokens must be billed at input_cost_per_token (matching the pre-fix
behavior) so existing callers see no change.
"""
usage = Usage(
prompt_tokens=1000,
completion_tokens=100,
total_tokens=1100,
prompt_tokens_details=PromptTokensDetailsWrapper(
cached_tokens=400,
audio_tokens=0,
),
)
response = ModelResponse(
id="test-id",
created=1234567890,
model="openai/gpt-5.4",
object="chat.completion",
choices=[],
usage=usage,
)
cost = litellm.completion_cost(
completion_response=response,
model="openai/gpt-5.4",
custom_llm_provider="openai",
custom_cost_per_token={
"input_cost_per_token": 0.0000025,
"output_cost_per_token": 0.000015,
},
)
# All 1000 prompt tokens billed at input rate, regardless of cached_tokens.
expected = 1000 * 0.0000025 + 100 * 0.000015
assert cost == pytest.approx(expected)

View file

@ -0,0 +1,66 @@
"""
Tests for guardrail exception status codes.
GuardrailRaisedException and BlockedPiiEntityError must carry
``status_code = 400`` so the proxy exception handler
(``getattr(e, "status_code", 500)``) returns HTTP 400 instead of 500
for intentional guardrail blocks.
"""
from litellm.exceptions import BlockedPiiEntityError, GuardrailRaisedException
class TestGuardrailRaisedExceptionStatusCode:
"""GuardrailRaisedException should default to status_code=400."""
def test_default_status_code(self):
exc = GuardrailRaisedException(
guardrail_name="test_guardrail",
message="blocked",
)
assert exc.status_code == 400
def test_custom_status_code(self):
exc = GuardrailRaisedException(
guardrail_name="test_guardrail",
message="rate limited",
status_code=429,
)
assert exc.status_code == 429
def test_getattr_fallback_resolves_to_400(self):
"""The proxy uses ``getattr(e, 'status_code', 500)`` — verify it
resolves to 400, not the 500 default."""
exc = GuardrailRaisedException(
guardrail_name="test_guardrail",
message="blocked",
)
assert getattr(exc, "status_code", 500) == 400
class TestBlockedPiiEntityErrorStatusCode:
"""BlockedPiiEntityError should default to status_code=400."""
def test_default_status_code(self):
exc = BlockedPiiEntityError(
entity_type="CREDIT_CARD",
guardrail_name="presidio",
)
assert exc.status_code == 400
def test_custom_status_code(self):
exc = BlockedPiiEntityError(
entity_type="SSN",
guardrail_name="presidio",
status_code=403,
)
assert exc.status_code == 403
def test_getattr_fallback_resolves_to_400(self):
"""The proxy uses ``getattr(e, 'status_code', 500)`` — verify it
resolves to 400, not the 500 default."""
exc = BlockedPiiEntityError(
entity_type="PHONE_NUMBER",
guardrail_name="presidio",
)
assert getattr(exc, "status_code", 500) == 400

View file

@ -1,5 +1,4 @@
import json
import os
from unittest.mock import MagicMock, patch
import pytest
@ -165,6 +164,13 @@ def test_max_connections_in_cluster_kwargs():
), "max_connections should be in available Redis cluster kwargs"
def test_socket_timeouts_in_cluster_kwargs():
"""Test that Redis cluster clients can receive socket timeout configuration"""
kwargs = _get_redis_cluster_kwargs()
assert "socket_timeout" in kwargs
assert "socket_connect_timeout" in kwargs
def test_get_redis_async_client_with_connection_pool():
"""Test that connection_pool parameter is properly passed to Redis client"""
# Create a mock connection pool

View file

@ -3697,3 +3697,173 @@ def test_try_early_resolve_deployments_for_model_not_in_names():
default_router.default_deployment["litellm_params"]["model"]
== "openai/will-be-overridden"
)
def _router_with_two_deployments(blocked_flags):
import litellm
model_list = []
for idx, blocked in enumerate(blocked_flags):
model_list.append(
{
"model_name": "gpt-4o",
"litellm_params": {"model": f"openai/gpt-4o-{idx}"},
"model_info": {"id": f"dep-{idx}", "blocked": blocked},
}
)
return litellm.Router(model_list=model_list)
def test_get_fully_blocked_model_names_marks_name_when_all_deployments_blocked():
router = _router_with_two_deployments([True, True])
assert router.get_fully_blocked_model_names() == {"gpt-4o"}
def test_get_fully_blocked_model_names_keeps_name_when_partial_blocked():
router = _router_with_two_deployments([True, False])
assert router.get_fully_blocked_model_names() == set()
def test_get_fully_blocked_model_names_treats_missing_key_as_unblocked():
import litellm
router = litellm.Router(
model_list=[
{
"model_name": "gpt-4o",
"litellm_params": {"model": "openai/gpt-4o"},
"model_info": {"id": "dep-0"},
}
]
)
assert router.get_fully_blocked_model_names() == set()
@pytest.mark.asyncio
async def test_async_get_healthy_deployments_skips_blocked_deployment():
router = _router_with_two_deployments([True, False])
healthy, all_dep = await router._async_get_healthy_deployments(
model="gpt-4o", parent_otel_span=None
)
healthy_ids = [d["model_info"]["id"] for d in healthy]
assert "dep-0" not in healthy_ids
assert "dep-1" in healthy_ids
assert len(all_dep) == 2
def test_get_healthy_deployments_sync_skips_blocked_deployment():
router = _router_with_two_deployments([False, True])
healthy, all_dep = router._get_healthy_deployments(
model="gpt-4o", parent_otel_span=None
)
healthy_ids = [d["model_info"]["id"] for d in healthy]
assert "dep-0" in healthy_ids
assert "dep-1" not in healthy_ids
assert len(all_dep) == 2
def test_filter_blocked_deployments_drops_blocked_keeps_unblocked():
router = _router_with_two_deployments([True, False])
filtered = router._filter_blocked_deployments(router.get_model_list() or [])
ids = [d["model_info"]["id"] for d in filtered]
assert ids == ["dep-1"]
@pytest.mark.asyncio
async def test_public_async_get_healthy_deployments_skips_blocked_on_primary_path():
router = _router_with_two_deployments([True, False])
deployments = await router.async_get_healthy_deployments(
model="gpt-4o", request_kwargs={}
)
assert isinstance(deployments, list)
ids = [d["model_info"]["id"] for d in deployments]
assert "dep-0" not in ids
assert "dep-1" in ids
def test_public_get_available_deployment_skips_blocked_on_primary_path():
router = _router_with_two_deployments([True, False])
deployment = router.get_available_deployment(model="gpt-4o", request_kwargs={})
assert deployment["model_info"]["id"] == "dep-1"
def test_get_available_deployment_raises_when_addressed_dict_is_blocked():
import litellm
router = _router_with_two_deployments([True, True])
with pytest.raises(litellm.ServiceUnavailableError):
router.get_available_deployment(model="dep-0", request_kwargs={})
def _router_with_two_pass_through_deployments(blocked_flags):
import litellm
model_list = []
for idx, blocked in enumerate(blocked_flags):
model_list.append(
{
"model_name": "gpt-4o",
"litellm_params": {
"model": f"openai/gpt-4o-{idx}",
"api_key": "sk-fake-for-tests",
"use_in_pass_through": True,
},
"model_info": {"id": f"pt-{idx}", "blocked": blocked},
}
)
return litellm.Router(model_list=model_list)
def test_get_available_deployment_for_pass_through_skips_blocked():
router = _router_with_two_pass_through_deployments([True, False])
deployment = router.get_available_deployment_for_pass_through(
model="gpt-4o", request_kwargs={}
)
assert deployment["model_info"]["id"] == "pt-1"
def test_get_available_deployment_for_pass_through_raises_when_dict_blocked():
import litellm
router = _router_with_two_pass_through_deployments([True, True])
with pytest.raises(litellm.ServiceUnavailableError):
router.get_available_deployment_for_pass_through(
model="pt-0", request_kwargs={}
)
def test_get_deployment_credentials_returns_none_for_blocked_deployment():
router = _router_with_two_deployments([True, False])
assert router.get_deployment_credentials(model_id="dep-0") is None
assert router.get_deployment_credentials(model_id="dep-1") is not None
def test_get_deployment_credentials_with_provider_returns_none_for_blocked_deployment():
router = _router_with_two_deployments([True, False])
assert router.get_deployment_credentials_with_provider(model_id="dep-0") is None
assert router.get_deployment_credentials_with_provider(model_id="dep-1") is not None
def test_is_deployment_blocked_static_helper_reflects_blocked_flag():
"""
Exercises Router._is_deployment_blocked so router_code_coverage.py (AST call graph)
marks the helper as covered by router-named tests.
"""
import types
import litellm
router = _router_with_two_deployments([True, False])
blocked_dep = router.get_deployment("dep-0")
unblocked_dep = router.get_deployment("dep-1")
assert blocked_dep is not None and unblocked_dep is not None
assert litellm.Router._is_deployment_blocked(blocked_dep) is True
assert litellm.Router._is_deployment_blocked(unblocked_dep) is False
# No model_info on deployment object → treated as not blocked
assert litellm.Router._is_deployment_blocked(object()) is False
missing_blocked = types.SimpleNamespace()
assert litellm.Router._is_deployment_blocked(types.SimpleNamespace(model_info=missing_blocked)) is False
assert litellm.Router._is_deployment_blocked(
types.SimpleNamespace(model_info=types.SimpleNamespace(blocked=True))
) is True

View file

@ -12,14 +12,17 @@ sys.path.insert(
) # Adds the parent directory to the system path
import litellm # noqa: E402,F401
from tests._vcr_conftest_common import ( # noqa: E402
from tests._vcr_conftest_common import ( # noqa: E402,F401
VerboseReporterState,
_pin_multipart_boundary,
apply_vcr_auto_marker_to_items,
emit_cassette_cache_session_banner,
emit_vcr_classification_summary,
emit_vcr_diagnostic_log,
install_live_call_probe,
record_vcr_outcome,
register_persister_if_enabled,
reset_vcr_diag_dir,
vcr_config_dict,
)
@ -84,6 +87,7 @@ def _vcr_outcome_gate(request, vcr):
def pytest_configure(config):
_verbose_state.remember_pluginmanager(config)
reset_vcr_diag_dir()
def pytest_runtest_logreport(report):
@ -110,3 +114,4 @@ def pytest_collection_modifyitems(config, items):
def pytest_terminal_summary(terminalreporter, exitstatus, config):
emit_cassette_cache_session_banner(terminalreporter)
emit_vcr_classification_summary(terminalreporter)
emit_vcr_diagnostic_log(terminalreporter)

View file

@ -654,12 +654,7 @@ const UsagePage: React.FC<UsagePageProps> = ({ teams, organizations }) => {
<Card>
<Title>Input Tokens</Title>
<Text className="text-2xl font-bold mt-2 text-blue-600">
{Math.max(
0,
(userSpendData.metadata?.total_prompt_tokens || 0) -
(userSpendData.metadata?.total_cache_read_input_tokens || 0) -
(userSpendData.metadata?.total_cache_creation_input_tokens || 0)
).toLocaleString()}
{(userSpendData.metadata?.total_prompt_tokens || 0).toLocaleString()}
</Text>
</Card>
<Card>

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