diff --git a/README.md b/README.md index 32b0160dbaa..68aaa09ec98 100644 --- a/README.md +++ b/README.md @@ -292,6 +292,7 @@ curl -X POST 'http://0.0.0.0:4000/v1/chat/completions' \ | [Clarifai (`clarifai`)](https://docs.litellm.ai/docs/providers/clarifai) | ✅ | ✅ | ✅ | | | | | | | | | [Cloudflare AI Workers (`cloudflare`)](https://docs.litellm.ai/docs/providers/cloudflare_workers) | ✅ | ✅ | ✅ | | | | | | | | | [Codestral (`codestral`)](https://docs.litellm.ai/docs/providers/codestral) | ✅ | ✅ | ✅ | | | | | | | | +| [Cognition (`cognition`)](https://docs.litellm.ai/docs/providers/cognition) | ✅ | ✅ | ✅ | | | | | | | | | [Cohere (`cohere`)](https://docs.litellm.ai/docs/providers/cohere) | ✅ | ✅ | ✅ | ✅ | | | | | | ✅ | | [Cohere Chat (`cohere_chat`)](https://docs.litellm.ai/docs/providers/cohere) | ✅ | ✅ | ✅ | | | | | | | | | [CometAPI (`cometapi`)](https://docs.litellm.ai/docs/providers/cometapi) | ✅ | ✅ | ✅ | ✅ | | | | | | | diff --git a/litellm/__init__.py b/litellm/__init__.py index 00f67ea0ff5..e95b553c5d4 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -453,6 +453,7 @@ max_end_user_budget_id: Optional[str] = None # backwards compatibility — arbitrary client-supplied identifiers still # pass through unchanged. validate_end_user_id_in_db: bool = False +block_requests_for_models_without_pricing: bool = False disable_end_user_cost_tracking: Optional[bool] = None disable_end_user_cost_tracking_prometheus_only: Optional[bool] = None enable_end_user_cost_tracking_prometheus_only: Optional[bool] = None diff --git a/litellm/caching/_embedding_router.py b/litellm/caching/_embedding_router.py index 8dfcddf158a..cec25634bb8 100644 --- a/litellm/caching/_embedding_router.py +++ b/litellm/caching/_embedding_router.py @@ -16,6 +16,7 @@ from collections.abc import Sequence from typing import TYPE_CHECKING, Any, Final import litellm +from litellm.constants import SEMANTIC_CACHE_EMBEDDING_TIMEOUT_SECONDS if TYPE_CHECKING: from litellm.router import Router @@ -60,6 +61,13 @@ def resolve_embedding_max_input_tokens( return deployment_max_input_tokens +def resolve_embedding_timeout(configured_timeout: float | None) -> float: + """Explicit cache setting first, else the short semantic-cache default.""" + if configured_timeout is not None: + return configured_timeout + return SEMANTIC_CACHE_EMBEDDING_TIMEOUT_SECONDS + + def truncate_embedding_input(prompt: str, embedding_model: str, max_input_tokens: int | None) -> str: """Keep only the first ``max_input_tokens`` tokens of ``prompt`` for the embedding call.""" if max_input_tokens is None: diff --git a/litellm/caching/caching.py b/litellm/caching/caching.py index 6b68ae98111..cefe6aae9ed 100644 --- a/litellm/caching/caching.py +++ b/litellm/caching/caching.py @@ -98,6 +98,7 @@ class Cache: qdrant_semantic_cache_embedding_model: str = "text-embedding-ada-002", qdrant_semantic_cache_vector_size: int | None = None, semantic_cache_embedding_max_input_tokens: int | None = None, + semantic_cache_embedding_timeout: float | None = None, # GCP IAM authentication parameters gcp_service_account: str | None = None, gcp_ssl_ca_certs: str | None = None, @@ -124,6 +125,7 @@ class Cache: qdrant_collection_name (str, optional): The name for your qdrant collection. Required if type is "qdrant-semantic". similarity_threshold (float, optional): The similarity threshold for semantic-caching, Required if type is "redis-semantic" or "qdrant-semantic". semantic_cache_embedding_max_input_tokens (int, optional): Truncate prompts to this many tokens before embedding them for semantic caching. Defaults to the embedding deployment's configured max_input_tokens. + semantic_cache_embedding_timeout (float, optional): Seconds a semantic-cache lookup may spend embedding the prompt before it gives up and lets the request continue to the LLM. Defaults to SEMANTIC_CACHE_EMBEDDING_TIMEOUT_SECONDS. # Disk Cache Args disk_cache_dir (str, optional): The directory for the disk cache. Defaults to None. @@ -195,6 +197,7 @@ class Cache: embedding_model=redis_semantic_cache_embedding_model, index_name=redis_semantic_cache_index_name, embedding_max_input_tokens=semantic_cache_embedding_max_input_tokens, + embedding_timeout=semantic_cache_embedding_timeout, **kwargs, ) elif type == LiteLLMCacheType.VALKEY_SEMANTIC: @@ -211,6 +214,7 @@ class Cache: index_name=valkey_semantic_cache_index_name, startup_nodes=redis_startup_nodes, embedding_max_input_tokens=semantic_cache_embedding_max_input_tokens, + embedding_timeout=semantic_cache_embedding_timeout, **kwargs, ) elif type == LiteLLMCacheType.QDRANT_SEMANTIC: @@ -223,6 +227,7 @@ class Cache: embedding_model=qdrant_semantic_cache_embedding_model, vector_size=qdrant_semantic_cache_vector_size, embedding_max_input_tokens=semantic_cache_embedding_max_input_tokens, + embedding_timeout=semantic_cache_embedding_timeout, ) elif type == LiteLLMCacheType.LOCAL: self.cache = InMemoryCache() diff --git a/litellm/caching/qdrant_semantic_cache.py b/litellm/caching/qdrant_semantic_cache.py index 8270c655d82..4898700c403 100644 --- a/litellm/caching/qdrant_semantic_cache.py +++ b/litellm/caching/qdrant_semantic_cache.py @@ -16,7 +16,11 @@ from typing import TYPE_CHECKING, Any, Final, cast import litellm from litellm._logging import print_verbose -from litellm.constants import QDRANT_SCALAR_QUANTILE, QDRANT_VECTOR_SIZE +from litellm.constants import ( + QDRANT_SCALAR_QUANTILE, + QDRANT_VECTOR_SIZE, + SEMANTIC_CACHE_EMBEDDING_TIMEOUT_SECONDS, +) from litellm.litellm_core_utils.prompt_templates.common_utils import ( get_str_from_messages, ) @@ -26,6 +30,7 @@ from ._embedding_router import ( build_router_embedding_metadata, resolve_embedding_max_input_tokens, resolve_embedding_router, + resolve_embedding_timeout, truncate_embedding_input, ) from .base_cache import BaseCache @@ -37,6 +42,7 @@ if TYPE_CHECKING: class QdrantSemanticCache(BaseCache): CACHE_KEY_FIELD_NAME = "litellm_cache_key" embedding_max_input_tokens: int | None = None + embedding_timeout: float = SEMANTIC_CACHE_EMBEDDING_TIMEOUT_SECONDS def __init__( self, @@ -49,6 +55,7 @@ class QdrantSemanticCache(BaseCache): host_type=None, vector_size=None, embedding_max_input_tokens: int | None = None, + embedding_timeout: float | None = None, ): from litellm.llms.custom_httpx.http_handler import ( _get_httpx_client, @@ -68,6 +75,7 @@ class QdrantSemanticCache(BaseCache): self.similarity_threshold = similarity_threshold self.embedding_model = embedding_model self.embedding_max_input_tokens = embedding_max_input_tokens + self.embedding_timeout = resolve_embedding_timeout(embedding_timeout) self.vector_size = vector_size if vector_size is not None else QDRANT_VECTOR_SIZE headers = {} @@ -222,11 +230,15 @@ class QdrantSemanticCache(BaseCache): input=embedding_input, cache={"no-store": True, "no-cache": True}, metadata=build_router_embedding_metadata(metadata), + timeout=self.embedding_timeout, + num_retries=0, ) return litellm.embedding( model=self.embedding_model, input=embedding_input, cache={"no-store": True, "no-cache": True}, + timeout=self.embedding_timeout, + num_retries=0, ) async def _get_async_embedding(self, prompt: str, metadata: dict[str, Any] | None = None) -> EmbeddingResponse: @@ -238,19 +250,25 @@ class QdrantSemanticCache(BaseCache): router: Final = resolve_embedding_router(self.embedding_model, llm_router, llm_model_list) embedding_input: Final = self._embedding_input(prompt, router) - if router is not None: - return await router.aembedding( + embedding_call: Final = ( + router.aembedding( model=self.embedding_model, input=embedding_input, cache={"no-store": True, "no-cache": True}, metadata=build_router_embedding_metadata(metadata), + timeout=self.embedding_timeout, + num_retries=0, + ) + if router is not None + else litellm.aembedding( + model=self.embedding_model, + input=embedding_input, + cache={"no-store": True, "no-cache": True}, + timeout=self.embedding_timeout, + num_retries=0, ) - - return await litellm.aembedding( - model=self.embedding_model, - input=embedding_input, - cache={"no-store": True, "no-cache": True}, ) + return await asyncio.wait_for(embedding_call, self.embedding_timeout) def set_cache(self, key, value, **kwargs): print_verbose(f"qdrant semantic-cache set_cache, kwargs: {kwargs}") diff --git a/litellm/caching/redis_semantic_cache.py b/litellm/caching/redis_semantic_cache.py index d91260f4d9c..f5264e28124 100644 --- a/litellm/caching/redis_semantic_cache.py +++ b/litellm/caching/redis_semantic_cache.py @@ -18,6 +18,7 @@ from typing import TYPE_CHECKING, Any, Final, cast import litellm from litellm._logging import print_verbose, verbose_logger +from litellm.constants import SEMANTIC_CACHE_EMBEDDING_TIMEOUT_SECONDS from litellm.litellm_core_utils.prompt_templates.common_utils import ( get_str_from_messages, ) @@ -27,6 +28,7 @@ from ._embedding_router import ( build_router_embedding_metadata, resolve_embedding_max_input_tokens, resolve_embedding_router, + resolve_embedding_timeout, truncate_embedding_input, ) from .base_cache import BaseCache @@ -47,6 +49,7 @@ class RedisSemanticCache(BaseCache): DEFAULT_REDIS_INDEX_NAME: str = "litellm_semantic_cache_index" CACHE_KEY_FIELD_NAME: str = "litellm_cache_key" embedding_max_input_tokens: int | None = None + embedding_timeout: float = SEMANTIC_CACHE_EMBEDDING_TIMEOUT_SECONDS def __init__( self, @@ -58,6 +61,7 @@ class RedisSemanticCache(BaseCache): embedding_model: str = "text-embedding-ada-002", index_name: str | None = None, embedding_max_input_tokens: int | None = None, + embedding_timeout: float | None = None, **kwargs: object, ): """ @@ -74,6 +78,8 @@ class RedisSemanticCache(BaseCache): index_name: Name for the Redis index embedding_max_input_tokens: Truncate prompts to this many tokens before embedding; defaults to the Router deployment's configured max_input_tokens + embedding_timeout: Seconds a cache lookup may spend embedding the prompt before it + gives up and lets the request continue to the LLM ttl: Default time-to-live for cache entries in seconds **kwargs: Additional arguments passed to the Redis client @@ -99,6 +105,7 @@ class RedisSemanticCache(BaseCache): self.distance_threshold = 1 - similarity_threshold self.embedding_model = embedding_model self.embedding_max_input_tokens = embedding_max_input_tokens + self.embedding_timeout = resolve_embedding_timeout(embedding_timeout) # Set up Redis connection if redis_url is None: @@ -349,6 +356,8 @@ class RedisSemanticCache(BaseCache): input=embedding_input, cache={"no-store": True, "no-cache": True}, metadata=build_router_embedding_metadata(metadata), + timeout=self.embedding_timeout, + num_retries=0, ), ) else: @@ -358,6 +367,8 @@ class RedisSemanticCache(BaseCache): model=self.embedding_model, input=embedding_input, cache={"no-store": True, "no-cache": True}, + timeout=self.embedding_timeout, + num_retries=0, ), ) return embedding_response["data"][0]["embedding"] @@ -512,20 +523,26 @@ class RedisSemanticCache(BaseCache): router: Final = resolve_embedding_router(self.embedding_model, llm_router, llm_model_list) embedding_input: Final = self._embedding_input(prompt, router) + embedding_call: Final = ( + router.aembedding( + model=self.embedding_model, + input=embedding_input, + cache={"no-store": True, "no-cache": True}, + metadata=build_router_embedding_metadata(metadata), + timeout=self.embedding_timeout, + num_retries=0, + ) + if router is not None + else litellm.aembedding( + model=self.embedding_model, + input=embedding_input, + cache={"no-store": True, "no-cache": True}, + timeout=self.embedding_timeout, + num_retries=0, + ) + ) try: - if router is not None: - embedding_response = await router.aembedding( - model=self.embedding_model, - input=embedding_input, - cache={"no-store": True, "no-cache": True}, - metadata=build_router_embedding_metadata(metadata), - ) - else: - embedding_response = await litellm.aembedding( - model=self.embedding_model, - input=embedding_input, - cache={"no-store": True, "no-cache": True}, - ) + embedding_response: Final = await asyncio.wait_for(embedding_call, self.embedding_timeout) return embedding_response["data"][0]["embedding"] except Exception as e: print_verbose(f"Error generating async embedding: {e}") diff --git a/litellm/caching/valkey_semantic_cache.py b/litellm/caching/valkey_semantic_cache.py index 737d212a89d..c66f6873383 100644 --- a/litellm/caching/valkey_semantic_cache.py +++ b/litellm/caching/valkey_semantic_cache.py @@ -30,6 +30,7 @@ from litellm._logging import print_verbose from litellm._uuid import uuid from litellm.llms.valkey.common_utils import build_valkey_url, pack_vector +from ._embedding_router import resolve_embedding_timeout from .redis_semantic_cache import RedisSemanticCache @@ -62,6 +63,7 @@ class ValkeySemanticCache(RedisSemanticCache): sync_client: Redis | None = None, async_client: AsyncRedis | None = None, embedding_max_input_tokens: int | None = None, + embedding_timeout: float | None = None, **kwargs: Any, ): if similarity_threshold is None: @@ -80,6 +82,7 @@ class ValkeySemanticCache(RedisSemanticCache): self.similarity_threshold = similarity_threshold self.embedding_model = embedding_model self.embedding_max_input_tokens = embedding_max_input_tokens + self.embedding_timeout = resolve_embedding_timeout(embedding_timeout) self.index_name = index_name or self.DEFAULT_VALKEY_INDEX_NAME self.key_prefix = f"{self.index_name}:" self._index_dim: int | None = None diff --git a/litellm/completion_extras/litellm_responses_transformation/transformation.py b/litellm/completion_extras/litellm_responses_transformation/transformation.py index 5f3e9ac753c..6103b1bf484 100644 --- a/litellm/completion_extras/litellm_responses_transformation/transformation.py +++ b/litellm/completion_extras/litellm_responses_transformation/transformation.py @@ -113,6 +113,58 @@ def _build_reasoning_item( } +def _reasoning_item_from_output_item(item: object) -> _BuiltReasoningItem | None: + from openai.types.responses import ResponseReasoningItem + + if isinstance(item, ResponseReasoningItem): + return _build_reasoning_item( + item_id=item.id, + encrypted_content=getattr(item, "encrypted_content", None), + summary_raw=item.summary, + ) + if isinstance(item, dict) and item.get("type") == "reasoning": + return _build_reasoning_item( + item_id=item.get("id", ""), + encrypted_content=item.get("encrypted_content"), + summary_raw=item.get("summary"), + ) + return None + + +def _reasoning_items_from_output_items(output_items: Sequence[object]) -> tuple[_BuiltReasoningItem, ...]: + return tuple( + reasoning_item + for reasoning_item in (_reasoning_item_from_output_item(item) for item in output_items) + if reasoning_item is not None + ) + + +def _as_chat_reasoning_items( + reasoning_items: Sequence[_BuiltReasoningItem], +) -> list[ChatCompletionReasoningItem] | None: + if not reasoning_items: + return None + # cast-ok: _BuiltReasoningItem is the structural shape ChatCompletionReasoningItem + # describes, and TypedDict invariance is what stops the two from unifying here. + return cast(list[ChatCompletionReasoningItem], list(reasoning_items)) + + +def _map_incomplete_reason_to_finish_reason(incomplete_reason: str | None) -> Literal["length", "content_filter"]: + if incomplete_reason == "content_filter": + return "content_filter" + return "length" + + +def _incomplete_reason_from_response_payload(response_payload: object) -> str | None: + if not isinstance(response_payload, Mapping): + return None + incomplete_details: Final = response_payload.get("incomplete_details") + if not isinstance(incomplete_details, Mapping): + return None + reason: Final = incomplete_details.get("reason") + return reason if isinstance(reason, str) else None + + class _ChatToolCallDict(ChatCompletionToolCallChunk, total=False): provider_specific_fields: Mapping[str, object] @@ -657,6 +709,27 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): return choices + @staticmethod + def _build_empty_incomplete_choice( + output_items: Sequence[object], + finish_reason: Literal["length", "content_filter"], + ) -> "Choices": + from litellm.types.utils import Choices, Message + + reasoning_items: Final = _reasoning_items_from_output_items(output_items) + reasoning_content: Final = " ".join( + summary_block["text"] + for reasoning_item in reasoning_items + for summary_block in reasoning_item["summary"] + if summary_block.get("text") + ) + message: Final = Message( + content="", + reasoning_content=reasoning_content if reasoning_content else None, + reasoning_items=_as_chat_reasoning_items(reasoning_items), + ) + return Choices(message=message, finish_reason=finish_reason, index=0) + @classmethod def _extract_output_from_completed_event(cls, parsed_chunk: Mapping[str, object]) -> list[dict[str, object]] | None: response_payload: Final = parsed_chunk.get("response") @@ -763,11 +836,22 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): handle_raw_dict_callback=self._handle_raw_dict_response_item, ) - if len(choices) == 0: - if raw_response.incomplete_details is not None and raw_response.incomplete_details.reason is not None: - raise ValueError(f"{model} unable to complete request: {raw_response.incomplete_details.reason}") + response_is_incomplete: Final = raw_response.status == "incomplete" or ( + raw_response.incomplete_details is not None and raw_response.incomplete_details.reason is not None + ) + + if len(choices) == 0 and not response_is_incomplete: + raise ValueError(f"Unknown items in responses API response: {output_items}") + + if response_is_incomplete: + incomplete_finish_reason: Final = _map_incomplete_reason_to_finish_reason( + raw_response.incomplete_details.reason if raw_response.incomplete_details is not None else None + ) + if len(choices) == 0: + choices.append(self._build_empty_incomplete_choice(output_items, incomplete_finish_reason)) else: - raise ValueError(f"Unknown items in responses API response: {output_items}") + for choice in choices: + choice.finish_reason = incomplete_finish_reason setattr(model_response, "choices", choices) @@ -1392,12 +1476,7 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator): ) ] ) - elif event_type == "response.completed": - # Response is fully complete - now we can signal is_finished=True - # This ensures we don't prematurely end the stream before tool_calls arrive - - # Check if response contains function_call items in output - # to determine correct finish_reason + elif event_type in ("response.completed", "response.incomplete"): response_data: Final = parsed_chunk.get("response", {}) output_items: Final = response_data.get("output", []) if response_data else [] @@ -1407,25 +1486,14 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator): if isinstance(item, dict) ) - finish_reason: Final = "tool_calls" if has_function_calls else "stop" + finish_reason: Final = ( + _map_incomplete_reason_to_finish_reason(_incomplete_reason_from_response_payload(response_data)) + if event_type == "response.incomplete" + else ("tool_calls" if has_function_calls else "stop") + ) - # Extract reasoning items with encrypted_content for round-tripping - completed_reasoning_items: list[_BuiltReasoningItem] | None = None - for item in output_items: - if not isinstance(item, dict) or item.get("type") != "reasoning": - continue - if completed_reasoning_items is None: - completed_reasoning_items = [] - completed_reasoning_items.append( - _build_reasoning_item( - item_id=item.get("id", ""), - encrypted_content=item.get("encrypted_content"), - summary_raw=item.get("summary"), - ) - ) - completed_reasoning_items_typed: Final = cast( - list[ChatCompletionReasoningItem] | None, - completed_reasoning_items, + terminal_reasoning_items_typed: Final = _as_chat_reasoning_items( + _reasoning_items_from_output_items(output_items) ) usage = None @@ -1439,7 +1507,7 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator): index=0, delta=Delta( content="", - reasoning_items=completed_reasoning_items_typed, + reasoning_items=terminal_reasoning_items_typed, ), finish_reason=finish_reason, ) diff --git a/litellm/constants.py b/litellm/constants.py index b9c92609361..928f7246a4a 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -436,6 +436,9 @@ DEFAULT_REQUEST_TIMEOUT_SECONDS: Final[float] = 6000.0 # deadline and connect handshake (see ``http_handler`` cached handler paths). COMPLETION_HTTP_FALLBACK_SECONDS: Final[float] = 600.0 HTTP_HANDLER_CONNECT_TIMEOUT_SECONDS: Final[float] = 5.0 +SEMANTIC_CACHE_EMBEDDING_TIMEOUT_SECONDS: Final[float] = float( + os.getenv("SEMANTIC_CACHE_EMBEDDING_TIMEOUT_SECONDS", "5.0") +) request_timeout: float = float(os.getenv("REQUEST_TIMEOUT", str(int(DEFAULT_REQUEST_TIMEOUT_SECONDS)))) request_timeout_explicitly_set: bool = "REQUEST_TIMEOUT" in os.environ DEFAULT_A2A_AGENT_TIMEOUT: Final[float] = float(os.getenv("DEFAULT_A2A_AGENT_TIMEOUT", 6000)) # 10 minutes @@ -779,6 +782,7 @@ openai_compatible_endpoints: Final[list] = [ "https://api.libertai.io/v1", "https://pinstripes.io/v1", "https://api.meta.ai/v1", + "https://api.cognition.ai/v1", ] @@ -846,6 +850,7 @@ openai_compatible_providers: Final[list] = [ "pinstripes", # Pinstripes - JSON-configured provider "darkbloom", "meta", # Meta Model API (Muse Spark) - JSON-configured provider + "cognition", ] openai_text_completion_compatible_providers: Final[list] = [ # providers that support `/v1/completions` "together_ai", @@ -1537,6 +1542,11 @@ DEFAULT_CRON_JOB_LOCK_TTL_SECONDS: Final = int(os.getenv("DEFAULT_CRON_JOB_LOCK_ PROXY_BUDGET_RESCHEDULER_MIN_TIME: Final = int(os.getenv("PROXY_BUDGET_RESCHEDULER_MIN_TIME", 597)) RESET_BUDGET_JOB_BATCH_SIZE: Final = max(1, int(os.getenv("RESET_BUDGET_JOB_BATCH_SIZE", "500"))) RESET_BUDGET_JOB_MAX_CHUNKS_PER_RUN: Final = max(1, int(os.getenv("RESET_BUDGET_JOB_MAX_CHUNKS_PER_RUN", "100"))) +RESET_BUDGET_JOB_NAME: Final = "reset_budget_job" +# Comfortably longer than one PROXY_BUDGET_RESCHEDULER_MIN_TIME tick, so a healthy +# leader keeps the lease across its own run, and a crashed one strands the sweep for +# at most a single tick. +RESET_BUDGET_JOB_LOCK_TTL_SECONDS: Final[int] = 900 PROXY_BATCH_POLLING_INTERVAL: Final = int(os.getenv("PROXY_BATCH_POLLING_INTERVAL", 3600)) MAX_OBJECTS_PER_POLL_CYCLE: Final = max(1, int(os.getenv("MAX_OBJECTS_PER_POLL_CYCLE", 50))) MANAGED_OBJECT_STALENESS_CUTOFF_DAYS: Final = max(1, int(os.getenv("MANAGED_OBJECT_STALENESS_CUTOFF_DAYS", 7))) @@ -1601,6 +1611,7 @@ LITELLM_SETTINGS_SAFE_DB_OVERRIDES: Final = [ "public_model_groups_links", "cost_discount_config", "cost_margin_config", + "block_requests_for_models_without_pricing", "budget_exceeded_throttle_percentage", # Every field editable from the Admin UI (proxy_server._GENERAL_SETTINGS_UI_LITELLM_FIELDS) # must be listed here so a DB write from one worker overrides the live litellm attribute on diff --git a/litellm/litellm_core_utils/get_llm_provider_logic.py b/litellm/litellm_core_utils/get_llm_provider_logic.py index dbb40913e14..e674fc37673 100644 --- a/litellm/litellm_core_utils/get_llm_provider_logic.py +++ b/litellm/litellm_core_utils/get_llm_provider_logic.py @@ -349,6 +349,9 @@ def get_llm_provider( elif endpoint == "https://api.meta.ai/v1": custom_llm_provider = "meta" dynamic_api_key = get_secret_str("META_API_KEY") + elif (json_provider := JSONProviderRegistry.get_by_base_url(endpoint)) is not None: + custom_llm_provider = json_provider.slug + dynamic_api_key = api_key if api_key is not None else get_secret_str(json_provider.api_key_env) if api_base is not None and not isinstance(api_base, str): raise Exception(f"api base needs to be a string. api_base={api_base}") diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index 0793fe20b21..0a52e1d283e 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -1371,6 +1371,7 @@ class CostCalculatorUtils: return fal_ai_image_cost_calculator( model=model, image_response=completion_response, + optional_params=optional_params, ) elif custom_llm_provider == litellm.LlmProviders.RUNWAYML.value: from litellm.llms.runwayml.cost_calculator import ( diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index 0ed15c43ccf..b676077ab0e 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -1200,13 +1200,14 @@ def _encode_tool_call_id_with_signature(tool_call_id: str, thought_signature: st return tool_call_id -def _get_thought_signature_from_tool(tool: dict, model: str | None = None) -> str | None: +def _get_thought_signature_from_tool(tool: dict) -> str | None: """Extract thought signature from tool call's provider_specific_fields. If not provided try to extract thought signature from tool call id Checks both tool.provider_specific_fields and tool.function.provider_specific_fields. - If no signature is found and model is gemini-3, returns a dummy signature. + Returns None when the tool call carries no signature; callers decide whether a + placeholder signature is needed. """ # First check tool's provider_specific_fields provider_fields: Final = tool.get("provider_specific_fields") or {} @@ -1236,13 +1237,6 @@ def _get_thought_signature_from_tool(tool: dict, model: str | None = None) -> st if len(parts) == 2: _, signature = parts return signature - # If no signature found and model is gemini-3, return dummy signature - from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( - VertexGeminiConfig, - ) - - if model and VertexGeminiConfig._is_gemini_3_or_newer(model): - return _get_dummy_thought_signature() return None @@ -1251,10 +1245,14 @@ def _get_dummy_thought_signature() -> str: This is used when transferring conversation history from older models (like gemini-2.5-flash) to gemini-3, which requires thought_signature - for strict validation. + for strict validation. Google documents it as a last resort that "will + negatively impact model performance", so callers must only fall back to it + when no real signature is available. + + See: + https://ai.google.dev/gemini-api/docs/thought-signatures#faqs + https://docs.cloud.google.com/gemini-enterprise-agent-platform/models/thinking/thought-signatures """ - # Return a base64-encoded dummy signature string - # Below dummy signature is recommended by google - https://ai.google.dev/gemini-api/docs/thought-signatures#faqs dummy_data: Final = b"skip_thought_signature_validator" return base64.b64encode(dummy_data).decode("utf-8") @@ -1312,8 +1310,10 @@ def convert_to_gemini_tool_call_invoke( VertexGeminiConfig, ) + needs_dummy_signature: Final = model is not None and VertexGeminiConfig._is_gemini_3_or_newer(model) + if tool_calls is not None: - for idx, tool in enumerate(tool_calls): + for tool in tool_calls: if "function" in tool: gemini_function_call: VertexFunctionCall | None = _gemini_tool_call_invoke_helper( function_call_params=tool["function"], @@ -1321,7 +1321,13 @@ def convert_to_gemini_tool_call_invoke( ) if gemini_function_call is not None: part_dict: VertexPartType = {"function_call": gemini_function_call} - thought_signature = _get_thought_signature_from_tool(dict(tool), model=model) + thought_signature = _get_thought_signature_from_tool(dict(tool)) + # Gemini signs only the first functionCall part of a parallel batch, so scope the + # placeholder fallback to that part instead of fabricating one per sibling call: + # https://docs.cloud.google.com/gemini-enterprise-agent-platform/models/thinking/thought-signatures#parallel_function_calling_example + is_first_function_call = len(_parts_list) == 0 + if not thought_signature and is_first_function_call and needs_dummy_signature: + thought_signature = _get_dummy_thought_signature() if thought_signature: part_dict["thoughtSignature"] = thought_signature @@ -1344,7 +1350,7 @@ def convert_to_gemini_tool_call_invoke( thought_signature = provider_fields.get("thought_signature") # If no signature found and model is gemini-3, use dummy signature - if not thought_signature and model and VertexGeminiConfig._is_gemini_3_or_newer(model): + if not thought_signature and needs_dummy_signature: thought_signature = _get_dummy_thought_signature() if thought_signature: diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 485091bccd0..f6340426c1b 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -191,7 +191,7 @@ class CustomStreamWrapper: custom_llm_provider: str | None = None, stream_options=None, make_call: Callable | None = None, - _response_headers: dict | None = None, + _response_headers: dict | httpx.Headers | None = None, ): self.model = model self.make_call = make_call @@ -2315,10 +2315,18 @@ class CustomStreamWrapper: if self.logging_obj is None or not self.chunks: return try: - partial_response: Final = litellm.stream_chunk_builder(chunks=self.chunks) + partial_response: Final = litellm.stream_chunk_builder( + chunks=self.chunks, + messages=self.messages if isinstance(self.messages, list) else None, + ) + if partial_response is None: + return usage: Final = cast(Usage | None, getattr(partial_response, "usage", None)) if usage is None: return + if self.model: + partial_response.model = self.model + backfill_missing_cache_usage_fields(usage) self.logging_obj.model_call_details["combined_usage_object"] = usage self.logging_obj.model_call_details["response_cost"] = ( self.logging_obj._response_cost_calculator(result=partial_response) or 0.0 @@ -2439,6 +2447,35 @@ class CustomStreamWrapper: return chunk +def _cache_token_count(details: PromptTokensDetailsWrapper | None, keys: tuple[str, ...]) -> int: + for key in keys: + value = getattr(details, key, None) + if isinstance(value, int) and not isinstance(value, bool) and value: + return value + return 0 + + +def backfill_missing_cache_usage_fields(usage: Usage) -> None: + """Give partial-stream usage the same cache fields a complete stream reports. + + Carries OpenAI-style ``prompt_tokens_details`` counts up to the Anthropic-style + top-level keys, defaulting to zero. It must carry the real count rather than a + flat zero: downstream readers treat these keys as authoritative once present and + skip their own normalization, so a zero here would overwrite a real cache read. + """ + details: Final = usage.prompt_tokens_details + if getattr(usage, "cache_read_input_tokens", None) is None: + usage.cache_read_input_tokens = _cache_token_count( # rebind-ok: in-place backfill is the contract + details, ("cached_tokens",) + ) + if getattr(usage, "cache_creation_input_tokens", None) is None: + usage.cache_creation_input_tokens = _cache_token_count( # rebind-ok: in-place backfill is the contract + details, ("cache_write_tokens", "cache_creation_tokens") + ) + if usage.prompt_tokens_details is None: + usage.prompt_tokens_details = PromptTokensDetailsWrapper(cached_tokens=0) # rebind-ok: backfill in place + + _TokenDetails = TypeVar("_TokenDetails", PromptTokensDetailsWrapper, CompletionTokensDetailsWrapper) diff --git a/litellm/llms/anthropic/chat/transformation.py b/litellm/llms/anthropic/chat/transformation.py index 414fd23381a..bca9b6bbec4 100644 --- a/litellm/llms/anthropic/chat/transformation.py +++ b/litellm/llms/anthropic/chat/transformation.py @@ -2216,7 +2216,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): def calculate_usage( self, - usage_object: dict, + usage_object: Mapping[str, Any], reasoning_content: str | None, completion_response: dict | None = None, speed: str | None = None, diff --git a/litellm/llms/bedrock/chat/converse_handler.py b/litellm/llms/bedrock/chat/converse_handler.py index 3a093e1f939..ca5f1298360 100644 --- a/litellm/llms/bedrock/chat/converse_handler.py +++ b/litellm/llms/bedrock/chat/converse_handler.py @@ -35,7 +35,7 @@ def make_sync_call( json_mode: bool | None = False, fake_stream: bool = False, stream_chunk_size: int | None = None, -): +) -> tuple[Any, httpx.Headers]: if client is None: client = _get_httpx_client() # Create a new client if none provided @@ -76,7 +76,7 @@ def make_sync_call( additional_args={"complete_input_dict": data}, ) - return completion_stream + return completion_stream, response.headers class BedrockConverseLLM(BaseAWSLLM): @@ -134,7 +134,7 @@ class BedrockConverseLLM(BaseAWSLLM): }, ) - completion_stream: Final = await make_call( + completion_stream, response_headers = await make_call( client=client, api_base=api_base, headers=dict(prepped.headers), @@ -151,6 +151,7 @@ class BedrockConverseLLM(BaseAWSLLM): model=model, custom_llm_provider="bedrock", logging_obj=logging_obj, + _response_headers=response_headers, ) return streaming_response @@ -232,7 +233,7 @@ class BedrockConverseLLM(BaseAWSLLM): except httpx.TimeoutException: raise BedrockError(status_code=408, message="Timeout error occurred.") - return litellm.AmazonConverseConfig()._transform_response( + transformed_response: Final = litellm.AmazonConverseConfig()._transform_response( model=model, response=response, model_response=model_response, @@ -244,6 +245,8 @@ class BedrockConverseLLM(BaseAWSLLM): optional_params=optional_params, encoding=encoding, ) + transformed_response.set_provider_response_headers(response.headers) + return transformed_response def completion( self, @@ -541,7 +544,7 @@ class BedrockConverseLLM(BaseAWSLLM): client = client if stream is not None and stream is True: - completion_stream: Final = make_sync_call( + completion_stream, response_headers = make_sync_call( client=(client if client is not None and isinstance(client, HTTPHandler) else None), api_base=proxy_endpoint_url, headers=prepped.headers, @@ -558,6 +561,7 @@ class BedrockConverseLLM(BaseAWSLLM): model=model, custom_llm_provider="bedrock", logging_obj=logging_obj, + _response_headers=response_headers, ) return streaming_response @@ -578,7 +582,7 @@ class BedrockConverseLLM(BaseAWSLLM): except httpx.TimeoutException: raise BedrockError(status_code=408, message="Timeout error occurred.") - return litellm.AmazonConverseConfig()._transform_response( + sync_transformed_response: Final = litellm.AmazonConverseConfig()._transform_response( model=model, response=response, model_response=model_response, @@ -590,3 +594,5 @@ class BedrockConverseLLM(BaseAWSLLM): optional_params=optional_params, encoding=encoding, ) + sync_transformed_response.set_provider_response_headers(response.headers) + return sync_transformed_response diff --git a/litellm/llms/bedrock/chat/invoke_handler.py b/litellm/llms/bedrock/chat/invoke_handler.py index 2a125e38a82..ce89c6c23e2 100644 --- a/litellm/llms/bedrock/chat/invoke_handler.py +++ b/litellm/llms/bedrock/chat/invoke_handler.py @@ -163,7 +163,7 @@ async def make_call( json_mode: bool | None = False, bedrock_invoke_provider: litellm.BEDROCK_INVOKE_PROVIDERS_LITERAL | None = None, stream_chunk_size: int | None = None, -): +) -> tuple[Any, httpx.Headers]: try: if client is None: client = get_async_httpx_client( @@ -225,7 +225,7 @@ async def make_call( additional_args={"complete_input_dict": data}, ) - return completion_stream + return completion_stream, response.headers except httpx.HTTPStatusError as err: error_code: Final = err.response.status_code raise BedrockError(status_code=error_code, message=err.response.text) @@ -248,7 +248,7 @@ def make_sync_call( json_mode: bool | None = False, bedrock_invoke_provider: litellm.BEDROCK_INVOKE_PROVIDERS_LITERAL | None = None, stream_chunk_size: int | None = None, -): +) -> tuple[Any, httpx.Headers]: try: if client is None: client = _get_httpx_client( @@ -309,7 +309,7 @@ def make_sync_call( additional_args={"complete_input_dict": data}, ) - return completion_stream + return completion_stream, response.headers except httpx.HTTPStatusError as err: error_code: Final = err.response.status_code raise BedrockError(status_code=error_code, message=err.response.text) diff --git a/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py index 76f91aa9115..333326a766b 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py @@ -1,7 +1,6 @@ import copy import json import time -from functools import partial from typing import TYPE_CHECKING, Any, Final, cast, get_args import httpx @@ -446,24 +445,24 @@ class AmazonInvokeConfig(BaseConfig, BaseAWSLLM): json_mode: bool | None = None, signed_json_body: bytes | None = None, ) -> CustomStreamWrapper: + completion_stream, response_headers = await make_call( + client=client, + api_base=api_base, + headers=headers, + data=json.dumps(data), + model=model, + messages=messages, + logging_obj=logging_obj, + fake_stream=True if "ai21" in api_base else False, + bedrock_invoke_provider=self.get_bedrock_invoke_provider(model), + json_mode=json_mode, + ) streaming_response: Final = CustomStreamWrapper( - completion_stream=None, - make_call=partial( - make_call, - client=client, - api_base=api_base, - headers=headers, - data=json.dumps(data), - model=model, - messages=messages, - logging_obj=logging_obj, - fake_stream=True if "ai21" in api_base else False, - bedrock_invoke_provider=self.get_bedrock_invoke_provider(model), - json_mode=json_mode, - ), + completion_stream=completion_stream, model=model, custom_llm_provider="bedrock", logging_obj=logging_obj, + _response_headers=response_headers, ) return streaming_response @@ -481,27 +480,28 @@ class AmazonInvokeConfig(BaseConfig, BaseAWSLLM): json_mode: bool | None = None, signed_json_body: bytes | None = None, ) -> CustomStreamWrapper: - if client is None or isinstance(client, AsyncHTTPHandler): - client = _get_httpx_client(params={}) + sync_client: Final = ( + _get_httpx_client(params={}) if client is None or isinstance(client, AsyncHTTPHandler) else client + ) + completion_stream, response_headers = make_sync_call( + client=sync_client, + api_base=api_base, + headers=headers, + data=json.dumps(data), + signed_json_body=signed_json_body, + model=model, + messages=messages, + logging_obj=logging_obj, + fake_stream=True if "ai21" in api_base else False, + bedrock_invoke_provider=self.get_bedrock_invoke_provider(model), + json_mode=json_mode, + ) streaming_response: Final = CustomStreamWrapper( - completion_stream=None, - make_call=partial( - make_sync_call, - client=client, - api_base=api_base, - headers=headers, - data=json.dumps(data), - signed_json_body=signed_json_body, - model=model, - messages=messages, - logging_obj=logging_obj, - fake_stream=True if "ai21" in api_base else False, - bedrock_invoke_provider=self.get_bedrock_invoke_provider(model), - json_mode=json_mode, - ), + completion_stream=completion_stream, model=model, custom_llm_provider="bedrock", logging_obj=logging_obj, + _response_headers=response_headers, ) return streaming_response diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 9a950d7f920..8c98c526da1 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -635,6 +635,7 @@ class BaseLLMHTTPHandler: model=model, custom_llm_provider=custom_llm_provider, logging_obj=logging_obj, + _response_headers=headers, ) if client is None or not isinstance(client, HTTPHandler): @@ -798,6 +799,7 @@ class BaseLLMHTTPHandler: model=model, custom_llm_provider=custom_llm_provider, logging_obj=logging_obj, + _response_headers=_response_headers, ) return streamwrapper diff --git a/litellm/llms/fal_ai/cost_calculator.py b/litellm/llms/fal_ai/cost_calculator.py index 8c5ad5a8c64..74848784c5b 100644 --- a/litellm/llms/fal_ai/cost_calculator.py +++ b/litellm/llms/fal_ai/cost_calculator.py @@ -1,25 +1,75 @@ -from typing import Any, Final +from collections.abc import Mapping +from types import MappingProxyType +from typing import Final import litellm from litellm.types.utils import ImageResponse +FAL_KEYED_PRICING_DEFAULT_QUALITY: Final[str] = "high" +FAL_TEXT_TO_IMAGE_DEFAULT_SIZE: Final[str] = "1024-x-768" +FAL_NAMED_IMAGE_SIZES: Final[Mapping[str, str]] = MappingProxyType( + { + "square_hd": "1024-x-1024", + "square": "512-x-512", + "portrait_4_3": "768-x-1024", + "portrait_16_9": "576-x-1024", + "landscape_4_3": "1024-x-768", + "landscape_16_9": "1024-x-576", + } +) + + +def _keyed_size(model: str, optional_params: Mapping[str, object]) -> str | None: + image_size: Final = optional_params.get("image_size") + if image_size is None: + return None if model.endswith("/edit") else FAL_TEXT_TO_IMAGE_DEFAULT_SIZE + if isinstance(image_size, Mapping): + width: Final = image_size.get("width") + height: Final = image_size.get("height") + if isinstance(width, int) and isinstance(height, int): + return f"{width}-x-{height}" + return None + if isinstance(image_size, str): + return FAL_NAMED_IMAGE_SIZES.get(image_size) + return None + + +def _keyed_cost_per_image(model: str, optional_params: Mapping[str, object] | None) -> float | None: + if optional_params is None: + return None + size: Final = _keyed_size(model=model, optional_params=optional_params) + if size is None: + return None + raw_quality: Final = optional_params.get("quality") + quality: Final = ( + raw_quality if isinstance(raw_quality, str) and raw_quality != "auto" else FAL_KEYED_PRICING_DEFAULT_QUALITY + ) + keyed_entry: Final = litellm.model_cost.get(f"fal_ai/{quality}/{size}/{model}") + if keyed_entry is None: + return None + keyed_cost: Final = keyed_entry.get("output_cost_per_image") + return float(keyed_cost) if isinstance(keyed_cost, (int, float)) else None + def cost_calculator( model: str, - image_response: Any, + image_response: object, + optional_params: Mapping[str, object] | None = None, ) -> float: """ fal.ai image generation cost calculator """ + if not isinstance(image_response, ImageResponse): + raise ValueError(f"image_response must be of type ImageResponse got type={type(image_response)}") + # the proxy cost path passes the provider-prefixed model name + model = model.removeprefix(f"{litellm.LlmProviders.FAL_AI.value}/") + num_images: Final[int] = len(image_response.data) if image_response.data else 0 + keyed_cost_per_image: Final = _keyed_cost_per_image(model=model, optional_params=optional_params) + if keyed_cost_per_image is not None: + return keyed_cost_per_image * num_images _model_info: Final = litellm.get_model_info( model=model, custom_llm_provider=litellm.LlmProviders.FAL_AI.value, ) output_cost_per_image: Final[float] = _model_info.get("output_cost_per_image") or 0.0 - num_images: int = 0 - if isinstance(image_response, ImageResponse): - if image_response.data: - num_images = len(image_response.data) - return output_cost_per_image * num_images - else: - raise ValueError(f"image_response must be of type ImageResponse got type={type(image_response)}") + return output_cost_per_image * num_images diff --git a/litellm/llms/openai_like/json_loader.py b/litellm/llms/openai_like/json_loader.py index 38f3866cfc3..5cdaff90d24 100644 --- a/litellm/llms/openai_like/json_loader.py +++ b/litellm/llms/openai_like/json_loader.py @@ -65,6 +65,11 @@ class JSONProviderRegistry: """Check if a provider is defined via JSON""" return slug in cls._providers + @classmethod + def get_by_base_url(cls, base_url: str) -> SimpleProviderConfig | None: + """Get a provider configuration by its default base url""" + return next((provider for provider in cls._providers.values() if provider.base_url == base_url), None) + @classmethod def supports_responses_api(cls, slug: str) -> bool: """Check if a JSON provider supports the Responses API""" diff --git a/litellm/llms/openai_like/providers.json b/litellm/llms/openai_like/providers.json index 164100d4194..5f57aaa78d8 100644 --- a/litellm/llms/openai_like/providers.json +++ b/litellm/llms/openai_like/providers.json @@ -175,6 +175,11 @@ "base_class": "openai_gpt", "supported_endpoints": ["/v1/chat/completions", "/v1/responses", "/v1/messages"] }, + "cognition": { + "base_url": "https://api.cognition.ai/v1", + "api_key_env": "COGNITION_API_KEY", + "api_base_env": "COGNITION_API_BASE" + }, "pinstripes": { "base_url": "https://pinstripes.io/v1", "api_key_env": "PINSTRIPES_API_KEY", diff --git a/litellm/llms/sagemaker/chat/transformation.py b/litellm/llms/sagemaker/chat/transformation.py index 99543e7add1..37ddd813d6f 100644 --- a/litellm/llms/sagemaker/chat/transformation.py +++ b/litellm/llms/sagemaker/chat/transformation.py @@ -54,7 +54,30 @@ class SagemakerChatConfig(OpenAIGPTConfig, BaseAWSLLM): api_key: str | None = None, api_base: str | None = None, ) -> dict: - return headers + inference_component_name: Final = optional_params.get("model_id") + if not isinstance(inference_component_name, str): + return headers + return {**headers, "X-Amzn-SageMaker-Inference-Component": inference_component_name} + + def transform_request( + self, + model: str, + messages: list[AllMessageValues], # mutable-ok: matches the base chat transform signature + optional_params: dict, # mutable-ok: matches the base chat transform signature + litellm_params: dict, # mutable-ok: matches the base chat transform signature + headers: dict, # mutable-ok: matches the base chat transform signature + ) -> dict: # mutable-ok: the handler sends this body straight to httpx + request: Final = super().transform_request( + model=model, + messages=messages, + optional_params=optional_params, + litellm_params=litellm_params, + headers=headers, + ) + served_model_name: Final = litellm_params.get("hf_model_name") + if not isinstance(served_model_name, str): + return request + return {**request, "model": served_model_name} def get_complete_url( self, diff --git a/litellm/llms/vertex_ai/gemini/transformation.py b/litellm/llms/vertex_ai/gemini/transformation.py index f2d318a9ffd..11c026010ee 100644 --- a/litellm/llms/vertex_ai/gemini/transformation.py +++ b/litellm/llms/vertex_ai/gemini/transformation.py @@ -645,10 +645,9 @@ def _collect_tool_call_thought_signatures( the text part as well would send two copies and double-bill the previous turn's reasoning tokens on gemini-3 and newer models. - Detection deliberately calls _get_thought_signature_from_tool without the - model argument: with a gemini-3 model that helper synthesizes a dummy - signature for unsigned tool calls, which must not suppress a real - text-part signature (e.g. replaying gemini-2.5 history to a newer model). + Only real signatures count here; a synthesized placeholder must not + suppress a genuine text-part signature (e.g. replaying gemini-2.5 history + to a newer model). """ signatures: tuple[str, ...] = () diff --git a/litellm/main.py b/litellm/main.py index f66767a8d42..7cfd322f3d0 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -5974,7 +5974,7 @@ def embedding( # Optional params dimensions: int | None = None, encoding_format: str | None = None, - timeout=600, # default to 10 minutes + timeout: float = 600, # default to 10 minutes # set api_base, api_version, api_key api_base: str | None = None, api_version: str | None = None, @@ -6000,7 +6000,7 @@ def embedding( # Optional params dimensions: int | None = None, encoding_format: str | None = None, - timeout=600, # default to 10 minutes + timeout: float = 600, # default to 10 minutes # set api_base, api_version, api_key api_base: str | None = None, api_version: str | None = None, @@ -6027,7 +6027,7 @@ def embedding( # Optional params dimensions: int | None = None, encoding_format: str | None = None, - timeout=600, # default to 10 minutes + timeout: float = 600, # default to 10 minutes # set api_base, api_version, api_key api_base: str | None = None, api_version: str | None = None, diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 5a1c988c21a..91c10d13e8e 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -759,7 +759,9 @@ "supports_vision": true, "supports_native_structured_output": true, "supports_parallel_tool_use_config": true, - "prompt_cache_min_tokens": 4096 + "prompt_cache_min_tokens": 4096, + "input_cost_per_token_batches": 5e-07, + "output_cost_per_token_batches": 2.5e-06 }, "anthropic.claude-haiku-4-5@20251001": { "cache_creation_input_token_cost": 1.25e-06, @@ -2487,7 +2489,9 @@ "supports_vision": true, "supports_native_structured_output": true, "supports_parallel_tool_use_config": true, - "prompt_cache_min_tokens": 1024 + "prompt_cache_min_tokens": 1024, + "input_cost_per_token_batches": 1.5e-06, + "output_cost_per_token_batches": 7.5e-06 }, "anthropic.claude-v1": { "input_cost_per_token": 8e-06, @@ -2743,7 +2747,9 @@ "supports_vision": true, "supports_native_structured_output": true, "supports_parallel_tool_use_config": true, - "prompt_cache_min_tokens": 4096 + "prompt_cache_min_tokens": 4096, + "input_cost_per_token_batches": 5.5e-07, + "output_cost_per_token_batches": 2.75e-06 }, "apac.anthropic.claude-3-sonnet-20240229-v1:0": { "deprecation_date": "2026-07-30", @@ -2839,7 +2845,9 @@ "supports_vision": true, "supports_native_structured_output": true, "supports_parallel_tool_use_config": true, - "prompt_cache_min_tokens": 1024 + "prompt_cache_min_tokens": 1024, + "input_cost_per_token_batches": 1.65e-06, + "output_cost_per_token_batches": 8.25e-06 }, "azure/ada": { "input_cost_per_token": 1e-07, @@ -12401,8 +12409,8 @@ "input_cost_per_token": 3e-06, "litellm_provider": "anthropic", "max_input_tokens": 1000000, - "max_output_tokens": 64000, - "max_tokens": 64000, + "max_output_tokens": 128000, + "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 1.5e-05, "search_context_cost_per_query": { @@ -12452,7 +12460,9 @@ "supports_tool_choice": true, "supports_vision": true, "supports_parallel_tool_use_config": true, - "prompt_cache_min_tokens": 1024 + "prompt_cache_min_tokens": 1024, + "input_cost_per_token_batches": 1.5e-06, + "output_cost_per_token_batches": 7.5e-06 }, "claude-opus-4-1": { "cache_creation_input_token_cost": 1.875e-05, @@ -12787,7 +12797,8 @@ "us": 1.1 }, "supports_output_config": true, - "prompt_cache_min_tokens": 512 + "prompt_cache_min_tokens": 512, + "supports_native_structured_output": true }, "claude-opus-5": { "deprecation_date": "2027-07-24", @@ -13350,7 +13361,8 @@ "max_output_tokens": 4096, "max_tokens": 4096, "mode": "completion", - "output_cost_per_token": 2e-06 + "output_cost_per_token": 2e-06, + "deprecation_date": "2025-09-15" }, "command-a-03-2025": { "input_cost_per_token": 2.5e-06, @@ -13371,7 +13383,8 @@ "max_tokens": 4096, "mode": "chat", "output_cost_per_token": 6e-07, - "supports_tool_choice": true + "supports_tool_choice": true, + "deprecation_date": "2025-09-15" }, "command-nightly": { "input_cost_per_token": 1e-06, @@ -13391,7 +13404,8 @@ "mode": "chat", "output_cost_per_token": 6e-07, "supports_function_calling": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "deprecation_date": "2025-09-15" }, "command-r-08-2024": { "input_cost_per_token": 1.5e-07, @@ -13413,7 +13427,8 @@ "mode": "chat", "output_cost_per_token": 1e-05, "supports_function_calling": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "deprecation_date": "2025-09-15" }, "command-r-plus-08-2024": { "input_cost_per_token": 2.5e-06, @@ -17027,7 +17042,9 @@ "supports_vision": true, "supports_native_structured_output": true, "supports_parallel_tool_use_config": true, - "prompt_cache_min_tokens": 4096 + "prompt_cache_min_tokens": 4096, + "input_cost_per_token_batches": 5.5e-07, + "output_cost_per_token_batches": 2.75e-06 }, "eu.anthropic.claude-3-5-sonnet-20240620-v1:0": { "input_cost_per_token": 3e-06, @@ -17250,7 +17267,9 @@ "supports_vision": true, "supports_native_structured_output": true, "supports_parallel_tool_use_config": true, - "prompt_cache_min_tokens": 1024 + "prompt_cache_min_tokens": 1024, + "input_cost_per_token_batches": 1.65e-06, + "output_cost_per_token_batches": 8.25e-06 }, "eu.meta.llama3-2-1b-instruct-v1:0": { "input_cost_per_token": 1.3e-07, @@ -17400,7 +17419,7 @@ "fal_ai/openai/gpt-image-2": { "litellm_provider": "fal_ai", "metadata": { - "notes": "OpenAI gpt-image-2 served through fal.ai. fal bills by token, so the flat output_cost_per_image here is fal's published per-image rate for a default request (quality=high, image_size=landscape_4_3 at 1024x768). Other canonical sizes at high quality: 1024x1024 $0.211, 1024x1536 $0.165, 1920x1080 $0.158, 2560x1440 $0.222, 3840x2160 $0.401" + "notes": "OpenAI gpt-image-2 served through fal.ai. fal bills by token but publishes deterministic per-image prices per size and quality, mirrored here as keyed entries fal_ai/{quality}/{width}-x-{height}/openai/gpt-image-2 that litellm's fal_ai cost calculator picks from the request params. This flat entry is the fallback when no keyed entry matches and carries the default request rate (quality=high, image_size=landscape_4_3 at 1024x768). quality=auto is priced as high" }, "mode": "image_generation", "output_cost_per_image": 0.145, @@ -17410,10 +17429,190 @@ ], "supports_vision": true }, + "fal_ai/low/1024-x-768/openai/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.005, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/low/1024-x-1024/openai/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.006, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/low/1024-x-1536/openai/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.005, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/low/1920-x-1080/openai/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.005, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/low/2560-x-1440/openai/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.007, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/low/3840-x-2160/openai/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.012, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/medium/1024-x-768/openai/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.037, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/medium/1024-x-1024/openai/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.053, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/medium/1024-x-1536/openai/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.042, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/medium/1920-x-1080/openai/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.04, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/medium/2560-x-1440/openai/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.056, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/medium/3840-x-2160/openai/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.101, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/high/1024-x-768/openai/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.145, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/high/1024-x-1024/openai/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.211, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/high/1024-x-1536/openai/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.165, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/high/1920-x-1080/openai/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.158, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/high/2560-x-1440/openai/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.222, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/high/3840-x-2160/openai/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.401, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, "fal_ai/gpt-image-2": { "litellm_provider": "fal_ai", "metadata": { - "notes": "Alias of fal_ai/openai/gpt-image-2, which litellm also accepts without the openai/ prefix. Same rate, see that entry for the size and quality caveat" + "notes": "Alias of fal_ai/openai/gpt-image-2, which litellm also accepts without the openai/ prefix. Same rates, including the keyed fal_ai/{quality}/{width}-x-{height}/gpt-image-2 entries; see that entry for details" }, "mode": "image_generation", "output_cost_per_image": 0.145, @@ -17423,13 +17622,373 @@ ], "supports_vision": true }, + "fal_ai/low/1024-x-768/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.005, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/low/1024-x-1024/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.006, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/low/1024-x-1536/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.005, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/low/1920-x-1080/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.005, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/low/2560-x-1440/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.007, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/low/3840-x-2160/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.012, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/medium/1024-x-768/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.037, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/medium/1024-x-1024/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.053, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/medium/1024-x-1536/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.042, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/medium/1920-x-1080/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.04, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/medium/2560-x-1440/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.056, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/medium/3840-x-2160/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.101, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/high/1024-x-768/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.145, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/high/1024-x-1024/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.211, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/high/1024-x-1536/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.165, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/high/1920-x-1080/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.158, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/high/2560-x-1440/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.222, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/high/3840-x-2160/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.401, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, "fal_ai/openai/gpt-image-2/edit": { "litellm_provider": "fal_ai", "metadata": { - "notes": "Editing endpoint of gpt-image-2 on fal.ai, reached through the image generation path with fal's image_urls param since /v1/images/edits is not wired for fal_ai. Same rate as fal_ai/openai/gpt-image-2, see that entry for the size and quality caveat" + "notes": "Editing endpoint of gpt-image-2 on fal.ai, reached through the image generation path with fal's image_urls param since /v1/images/edits is not wired for fal_ai. Prices include one input image and live in keyed entries fal_ai/{quality}/{width}-x-{height}/openai/gpt-image-2/edit. This flat entry is the fallback for the default edit request (quality=high, image_size=auto, inferred from the input image, priced as 1024x768 high)" }, "mode": "image_generation", - "output_cost_per_image": 0.145, + "output_cost_per_image": 0.151, + "source": "https://fal.ai/models/openai/gpt-image-2/edit", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/low/1024-x-768/openai/gpt-image-2/edit": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.011, + "source": "https://fal.ai/models/openai/gpt-image-2/edit", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/low/1024-x-1024/openai/gpt-image-2/edit": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.015, + "source": "https://fal.ai/models/openai/gpt-image-2/edit", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/low/1024-x-1536/openai/gpt-image-2/edit": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.018, + "source": "https://fal.ai/models/openai/gpt-image-2/edit", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/low/1920-x-1080/openai/gpt-image-2/edit": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.017, + "source": "https://fal.ai/models/openai/gpt-image-2/edit", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/low/2560-x-1440/openai/gpt-image-2/edit": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.019, + "source": "https://fal.ai/models/openai/gpt-image-2/edit", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/low/3840-x-2160/openai/gpt-image-2/edit": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.024, + "source": "https://fal.ai/models/openai/gpt-image-2/edit", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/medium/1024-x-768/openai/gpt-image-2/edit": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.043, + "source": "https://fal.ai/models/openai/gpt-image-2/edit", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/medium/1024-x-1024/openai/gpt-image-2/edit": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.061, + "source": "https://fal.ai/models/openai/gpt-image-2/edit", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/medium/1024-x-1536/openai/gpt-image-2/edit": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.054, + "source": "https://fal.ai/models/openai/gpt-image-2/edit", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/medium/1920-x-1080/openai/gpt-image-2/edit": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.053, + "source": "https://fal.ai/models/openai/gpt-image-2/edit", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/medium/2560-x-1440/openai/gpt-image-2/edit": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.068, + "source": "https://fal.ai/models/openai/gpt-image-2/edit", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/medium/3840-x-2160/openai/gpt-image-2/edit": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.113, + "source": "https://fal.ai/models/openai/gpt-image-2/edit", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/high/1024-x-768/openai/gpt-image-2/edit": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.151, + "source": "https://fal.ai/models/openai/gpt-image-2/edit", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/high/1024-x-1024/openai/gpt-image-2/edit": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.219, + "source": "https://fal.ai/models/openai/gpt-image-2/edit", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/high/1024-x-1536/openai/gpt-image-2/edit": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.178, + "source": "https://fal.ai/models/openai/gpt-image-2/edit", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/high/1920-x-1080/openai/gpt-image-2/edit": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.158, + "source": "https://fal.ai/models/openai/gpt-image-2/edit", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/high/2560-x-1440/openai/gpt-image-2/edit": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.234, + "source": "https://fal.ai/models/openai/gpt-image-2/edit", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/high/3840-x-2160/openai/gpt-image-2/edit": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.413, "source": "https://fal.ai/models/openai/gpt-image-2/edit", "supported_endpoints": [ "/v1/images/generations" @@ -18932,6 +19491,106 @@ }, "web_search_billing_unit": "per_query" }, + "gemini-3.1-flash-lite-image": { + "input_cost_per_image": 0.00028, + "input_cost_per_token": 2.5e-07, + "litellm_provider": "vertex_ai-language-models", + "max_input_tokens": 65536, + "max_output_tokens": 4096, + "max_tokens": 4096, + "mode": "image_generation", + "output_cost_per_image": 0.0336, + "output_cost_per_image_token": 3e-05, + "output_cost_per_token": 1.5e-06, + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#gemini-models", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text", + "image" + ], + "supports_function_calling": false, + "supports_prompt_caching": true, + "supports_response_schema": false, + "supports_reasoning": true, + "supports_system_messages": true, + "supports_vision": true + }, + "gemini/gemini-3.1-flash-lite-image": { + "rpm": 1000, + "tpm": 4000000, + "input_cost_per_image": 0.00028, + "input_cost_per_token": 2.5e-07, + "input_cost_per_token_batches": 1.25e-07, + "litellm_provider": "gemini", + "max_input_tokens": 65536, + "max_output_tokens": 4096, + "max_tokens": 4096, + "mode": "image_generation", + "output_cost_per_image": 0.0336, + "output_cost_per_image_token": 3e-05, + "output_cost_per_token": 1.5e-06, + "output_cost_per_token_batches": 7.5e-07, + "source": "https://ai.google.dev/gemini-api/docs/pricing#gemini-3.1-flash-lite-image", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text", + "image" + ], + "supports_function_calling": true, + "supports_prompt_caching": false, + "supports_response_schema": false, + "supports_reasoning": true, + "supports_system_messages": true, + "supports_vision": true + }, + "vertex_ai/gemini-3.1-flash-lite-image": { + "input_cost_per_image": 0.00028, + "input_cost_per_token": 2.5e-07, + "litellm_provider": "vertex_ai-language-models", + "max_input_tokens": 65536, + "max_output_tokens": 4096, + "max_tokens": 4096, + "mode": "image_generation", + "output_cost_per_image": 0.0336, + "output_cost_per_image_token": 3e-05, + "output_cost_per_token": 1.5e-06, + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#gemini-models", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text", + "image" + ], + "supports_function_calling": false, + "supports_prompt_caching": true, + "supports_response_schema": false, + "supports_reasoning": true, + "supports_system_messages": true, + "supports_vision": true + }, "gemini-3.1-flash-image": { "deprecation_date": "2027-05-28", "input_cost_per_image": 0.00056, @@ -19795,7 +20454,7 @@ "deprecation_date": "2027-05-19", "cache_read_input_token_cost": 1.5e-07, "input_cost_per_token": 1.5e-06, - "input_cost_per_audio_token": 1e-06, + "input_cost_per_audio_token": 1.5e-06, "litellm_provider": "vertex_ai", "max_input_tokens": 1048576, "max_output_tokens": 65535, @@ -19834,7 +20493,7 @@ "supports_web_search": true, "supports_native_streaming": true, "input_cost_per_token_priority": 2.7e-06, - "input_cost_per_audio_token_priority": 1.8e-06, + "input_cost_per_audio_token_priority": 2.7e-06, "output_cost_per_token_priority": 1.62e-05, "cache_read_input_token_cost_priority": 2.7e-07, "search_context_cost_per_query": { @@ -19842,7 +20501,12 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, - "web_search_billing_unit": "per_query" + "web_search_billing_unit": "per_query", + "input_cost_per_token_batches": 7.5e-07, + "output_cost_per_token_batches": 4.5e-06, + "input_cost_per_token_flex": 7.5e-07, + "output_cost_per_token_flex": 4.5e-06, + "cache_read_input_token_cost_flex": 7.5e-08 }, "vertex_ai/gemini-3.6-flash": { "prompt_cache_min_tokens": 4096, @@ -21527,7 +22191,7 @@ "gemini/gemini-3.5-flash": { "prompt_cache_min_tokens": 4096, "cache_read_input_token_cost": 1.5e-07, - "input_cost_per_audio_token": 1e-06, + "input_cost_per_audio_token": 1.5e-06, "input_cost_per_token": 1.5e-06, "litellm_provider": "gemini", "max_input_tokens": 1048576, @@ -21569,7 +22233,7 @@ "supports_native_streaming": true, "tpm": 800000, "input_cost_per_token_priority": 2.7e-06, - "input_cost_per_audio_token_priority": 1.8e-06, + "input_cost_per_audio_token_priority": 2.7e-06, "output_cost_per_token_priority": 1.62e-05, "cache_read_input_token_cost_priority": 2.7e-07, "search_context_cost_per_query": { @@ -21577,7 +22241,12 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, - "web_search_billing_unit": "per_query" + "web_search_billing_unit": "per_query", + "input_cost_per_token_batches": 7.5e-07, + "output_cost_per_token_batches": 4.5e-06, + "input_cost_per_token_flex": 7.5e-07, + "output_cost_per_token_flex": 4.5e-06, + "cache_read_input_token_cost_flex": 8e-08 }, "gemini/gemini-3.6-flash": { "prompt_cache_min_tokens": 4096, @@ -21929,7 +22598,7 @@ "prompt_cache_min_tokens": 4096, "deprecation_date": "2027-05-19", "cache_read_input_token_cost": 1.5e-07, - "input_cost_per_audio_token": 1e-06, + "input_cost_per_audio_token": 1.5e-06, "input_cost_per_token": 1.5e-06, "litellm_provider": "vertex_ai-language-models", "max_input_tokens": 1048576, @@ -21969,7 +22638,7 @@ "supports_web_search": true, "supports_native_streaming": true, "input_cost_per_token_priority": 2.7e-06, - "input_cost_per_audio_token_priority": 1.8e-06, + "input_cost_per_audio_token_priority": 2.7e-06, "output_cost_per_token_priority": 1.62e-05, "cache_read_input_token_cost_priority": 2.7e-07, "search_context_cost_per_query": { @@ -21977,7 +22646,12 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, - "web_search_billing_unit": "per_query" + "web_search_billing_unit": "per_query", + "input_cost_per_token_batches": 7.5e-07, + "output_cost_per_token_batches": 4.5e-06, + "input_cost_per_token_flex": 7.5e-07, + "output_cost_per_token_flex": 4.5e-06, + "cache_read_input_token_cost_flex": 7.5e-08 }, "gemini-3.6-flash": { "prompt_cache_min_tokens": 4096, @@ -23307,7 +23981,9 @@ "supports_vision": true, "supports_native_structured_output": true, "supports_parallel_tool_use_config": true, - "prompt_cache_min_tokens": 1024 + "prompt_cache_min_tokens": 1024, + "input_cost_per_token_batches": 1.5e-06, + "output_cost_per_token_batches": 7.5e-06 }, "global.anthropic.claude-sonnet-4-20250514-v1:0": { "cache_creation_input_token_cost": 3.75e-06, @@ -23365,7 +24041,9 @@ "supports_vision": true, "supports_native_structured_output": true, "supports_parallel_tool_use_config": true, - "prompt_cache_min_tokens": 4096 + "prompt_cache_min_tokens": 4096, + "input_cost_per_token_batches": 5e-07, + "output_cost_per_token_batches": 2.5e-06 }, "global.amazon.nova-2-lite-v1:0": { "cache_read_input_token_cost": 7.5e-08, @@ -24181,7 +24859,8 @@ "supports_response_schema": false, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": false + "supports_vision": false, + "deprecation_date": "2027-01-20" }, "gpt-4o-mini": { "cache_read_input_token_cost": 7.5e-08, @@ -25464,6 +26143,7 @@ "supported_output_modalities": [ "text" ], + "supports_computer_use": true, "supports_function_calling": true, "supports_minimal_reasoning_effort": false, "supports_native_streaming": true, @@ -25606,6 +26286,155 @@ "supports_web_search": true, "supports_xhigh_reasoning_effort": true }, + "gpt-5.6-cyber": { + "cache_creation_input_token_cost": 1.5625e-05, + "cache_creation_input_token_cost_above_272k_tokens": 3.125e-05, + "cache_read_input_token_cost": 1.25e-06, + "cache_read_input_token_cost_above_272k_tokens": 2.5e-06, + "input_cost_per_token": 1.25e-05, + "input_cost_per_token_above_272k_tokens": 2.5e-05, + "litellm_provider": "openai", + "max_input_tokens": 400000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 7.5e-05, + "output_cost_per_token_above_272k_tokens": 0.0001125, + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "source": "https://platform.openai.com/docs/models/gpt-5.6-cyber", + "supports_computer_use": true, + "supports_parallel_function_calling": true + }, + "daybreak-red-latest": { + "cache_creation_input_token_cost": 1.5625e-05, + "cache_creation_input_token_cost_above_272k_tokens": 3.125e-05, + "cache_read_input_token_cost": 1.25e-06, + "cache_read_input_token_cost_above_272k_tokens": 2.5e-06, + "input_cost_per_token": 1.25e-05, + "input_cost_per_token_above_272k_tokens": 2.5e-05, + "litellm_provider": "openai", + "max_input_tokens": 400000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 7.5e-05, + "output_cost_per_token_above_272k_tokens": 0.0001125, + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "source": "https://platform.openai.com/docs/models/daybreak-red-latest", + "supports_computer_use": true, + "supports_parallel_function_calling": true + }, + "daybreak-blue-latest": { + "cache_creation_input_token_cost": 6.25e-06, + "cache_creation_input_token_cost_above_272k_tokens": 1.25e-05, + "cache_read_input_token_cost": 5e-07, + "cache_read_input_token_cost_above_272k_tokens": 1e-06, + "input_cost_per_token": 5e-06, + "input_cost_per_token_above_272k_tokens": 1e-05, + "litellm_provider": "openai", + "max_input_tokens": 1050000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 3e-05, + "output_cost_per_token_above_272k_tokens": 4.5e-05, + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_computer_use": true, + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "source": "https://platform.openai.com/docs/models/daybreak-blue-latest", + "supports_parallel_function_calling": true + }, + "chat-latest": { + "cache_read_input_token_cost": 5e-07, + "input_cost_per_token": 5e-06, + "litellm_provider": "openai", + "max_input_tokens": 400000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 3e-05, + "source": "https://platform.openai.com/docs/models/chat-latest", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true + }, "gpt-5.5": { "cache_read_input_token_cost": 5e-07, "cache_read_input_token_cost_above_272k_tokens": 1e-06, @@ -28120,7 +28949,9 @@ "supports_vision": true, "supports_native_structured_output": true, "supports_parallel_tool_use_config": true, - "prompt_cache_min_tokens": 1024 + "prompt_cache_min_tokens": 1024, + "input_cost_per_token_batches": 1.65e-06, + "output_cost_per_token_batches": 8.25e-06 }, "jp.anthropic.claude-haiku-4-5-20251001-v1:0": { "cache_creation_input_token_cost": 1.375e-06, @@ -28146,7 +28977,9 @@ "supports_vision": true, "supports_native_structured_output": true, "supports_parallel_tool_use_config": true, - "prompt_cache_min_tokens": 4096 + "prompt_cache_min_tokens": 4096, + "input_cost_per_token_batches": 5.5e-07, + "output_cost_per_token_batches": 2.75e-06 }, "crusoe/deepseek-ai/DeepSeek-R1-0528": { "input_cost_per_token": 3e-06, @@ -29369,28 +30202,30 @@ "mistral/codestral-2508": { "input_cost_per_token": 3e-07, "litellm_provider": "mistral", - "max_input_tokens": 256000, - "max_output_tokens": 256000, - "max_tokens": 256000, + "max_input_tokens": 128000, + "max_output_tokens": 128000, + "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 9e-07, - "source": "https://mistral.ai/news/codestral-25-08", + "source": "https://docs.mistral.ai/models/model-cards/codestral-25-08", "supports_assistant_prefill": true, "supports_function_calling": true, "supports_response_schema": true, "supports_tool_choice": true }, "mistral/codestral-latest": { - "input_cost_per_token": 1e-06, + "input_cost_per_token": 3e-07, "litellm_provider": "mistral", - "max_input_tokens": 32000, - "max_output_tokens": 8191, - "max_tokens": 8191, + "max_input_tokens": 128000, + "max_output_tokens": 128000, + "max_tokens": 128000, "mode": "chat", - "output_cost_per_token": 3e-06, + "output_cost_per_token": 9e-07, "supports_assistant_prefill": true, "supports_response_schema": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "source": "https://docs.mistral.ai/models/model-cards/codestral-25-08", + "supports_function_calling": true }, "mistral/codestral-mamba-latest": { "input_cost_per_token": 2.5e-07, @@ -29623,6 +30458,16 @@ ], "source": "https://mistral.ai/pricing#api-pricing" }, + "mistral/mistral-ocr-4-1": { + "annotation_cost_per_page": 0.005, + "litellm_provider": "mistral", + "mode": "ocr", + "ocr_cost_per_page": 0.004, + "source": "https://docs.mistral.ai/models/model-cards/ocr-4-1", + "supported_endpoints": [ + "/v1/ocr" + ] + }, "mistral/mistral-ocr-2505-completion": { "deprecation_date": "2026-05-31", "litellm_provider": "mistral", @@ -29947,18 +30792,19 @@ "supports_tool_choice": true }, "mistral/mistral-small-latest": { - "input_cost_per_token": 6e-08, + "input_cost_per_token": 1.5e-07, "litellm_provider": "mistral", - "max_input_tokens": 131072, - "max_output_tokens": 131072, - "max_tokens": 131072, + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, "mode": "chat", - "output_cost_per_token": 1.8e-07, - "source": "https://mistral.ai/pricing", + "output_cost_per_token": 6e-07, + "source": "https://docs.mistral.ai/models/model-cards/mistral-small-4-0-26-03", "supports_assistant_prefill": true, "supports_function_calling": true, "supports_response_schema": true, "supports_tool_choice": true, + "supports_reasoning": true, "supports_vision": true }, "mistral/mistral-small-3-2-2506": { @@ -32781,6 +33627,31 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true }, + "openrouter/anthropic/claude-opus-5": { + "prompt_cache_min_tokens": 512, + "supports_adaptive_thinking": true, + "cache_creation_input_token_cost": 6.25e-06, + "cache_read_input_token_cost": 5e-07, + "input_cost_per_token": 5e-06, + "litellm_provider": "openrouter", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 2.5e-05, + "source": "https://openrouter.ai/anthropic/claude-opus-5", + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_max_reasoning_effort": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true + }, "openrouter/bytedance/ui-tars-1.5-7b": { "input_cost_per_token": 1e-07, "litellm_provider": "openrouter", @@ -32889,6 +33760,38 @@ "supports_reasoning": true, "supports_tool_choice": true }, + "openrouter/deepseek/deepseek-v4-pro": { + "input_cost_per_token": 1.32e-06, + "input_cost_per_token_cache_hit": 4.4e-08, + "litellm_provider": "openrouter", + "max_input_tokens": 1048576, + "max_output_tokens": 384000, + "max_tokens": 384000, + "mode": "chat", + "output_cost_per_token": 3.96e-06, + "source": "https://openrouter.ai/deepseek/deepseek-v4-pro", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "openrouter/deepseek/deepseek-v4-pro-0813": { + "input_cost_per_token": 1.32e-06, + "input_cost_per_token_cache_hit": 4.4e-08, + "litellm_provider": "openrouter", + "max_input_tokens": 1048576, + "max_output_tokens": 384000, + "max_tokens": 384000, + "mode": "chat", + "output_cost_per_token": 3.96e-06, + "source": "https://openrouter.ai/deepseek/deepseek-v4-pro-0813", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, "openrouter/google/gemini-2.0-flash-001": { "deprecation_date": "2026-06-01", "input_cost_per_audio_token": 7e-07, @@ -34923,7 +35826,9 @@ "supports_function_calling": true, "supports_reasoning": true, "supports_tool_choice": true, - "supports_native_structured_output": true + "supports_native_structured_output": true, + "input_cost_per_token_batches": 1.1e-07, + "output_cost_per_token_batches": 4.4e-07 }, "qwen.qwen3-coder-30b-a3b-v1:0": { "input_cost_per_token": 1.5e-07, @@ -35458,7 +36363,8 @@ "max_output_tokens": 4096, "max_tokens": 4096, "mode": "rerank", - "output_cost_per_token": 0.0 + "output_cost_per_token": 0.0, + "deprecation_date": "2025-04-30" }, "rerank-english-v3.0": { "input_cost_per_query": 0.002, @@ -35478,7 +36384,8 @@ "max_output_tokens": 4096, "max_tokens": 4096, "mode": "rerank", - "output_cost_per_token": 0.0 + "output_cost_per_token": 0.0, + "deprecation_date": "2025-04-30" }, "rerank-multilingual-v3.0": { "input_cost_per_query": 0.002, @@ -37118,7 +38025,9 @@ "supports_vision": true, "supports_native_structured_output": true, "supports_parallel_tool_use_config": true, - "prompt_cache_min_tokens": 4096 + "prompt_cache_min_tokens": 4096, + "input_cost_per_token_batches": 5.5e-07, + "output_cost_per_token_batches": 2.75e-06 }, "us.anthropic.claude-3-5-sonnet-20240620-v1:0": { "input_cost_per_token": 3e-06, @@ -37284,7 +38193,9 @@ "supports_vision": true, "supports_native_structured_output": true, "supports_parallel_tool_use_config": true, - "prompt_cache_min_tokens": 1024 + "prompt_cache_min_tokens": 1024, + "input_cost_per_token_batches": 1.65e-06, + "output_cost_per_token_batches": 8.25e-06 }, "us-gov.anthropic.claude-sonnet-4-5-20250929-v1:0": { "cache_creation_input_token_cost": 4.5e-06, @@ -37339,7 +38250,9 @@ "supports_vision": true, "supports_native_structured_output": true, "supports_parallel_tool_use_config": true, - "prompt_cache_min_tokens": 4096 + "prompt_cache_min_tokens": 4096, + "input_cost_per_token_batches": 5.5e-07, + "output_cost_per_token_batches": 2.75e-06 }, "us.anthropic.claude-opus-4-20250514-v1:0": { "cache_creation_input_token_cost": 1.875e-05, @@ -39930,13 +40843,13 @@ "supports_tool_choice": true }, "vertex_ai/deepseek-ai/deepseek-v3.1-maas": { - "input_cost_per_token": 1.35e-06, + "input_cost_per_token": 6e-07, "litellm_provider": "vertex_ai-deepseek_models", "max_input_tokens": 163840, "max_output_tokens": 32768, "max_tokens": 32768, "mode": "chat", - "output_cost_per_token": 5.4e-06, + "output_cost_per_token": 1.7e-06, "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#partner-models", "supported_regions": [ "us-central1" @@ -40770,13 +41683,13 @@ "supports_vision": true }, "vertex_ai/openai/gpt-oss-120b-maas": { - "input_cost_per_token": 1.5e-07, + "input_cost_per_token": 9e-08, "litellm_provider": "vertex_ai-openai_models", "max_input_tokens": 131072, "max_output_tokens": 32768, "max_tokens": 32768, "mode": "chat", - "output_cost_per_token": 6e-07, + "output_cost_per_token": 3.6e-07, "source": "https://console.cloud.google.com/vertex-ai/publishers/openai/model-garden/gpt-oss-120b-maas", "supports_reasoning": true }, @@ -40858,13 +41771,13 @@ "supports_web_search": true }, "vertex_ai/qwen/qwen3-235b-a22b-instruct-2507-maas": { - "input_cost_per_token": 2.5e-07, + "input_cost_per_token": 2.2e-07, "litellm_provider": "vertex_ai-qwen_models", "max_input_tokens": 262144, "max_output_tokens": 16384, "max_tokens": 16384, "mode": "chat", - "output_cost_per_token": 1e-06, + "output_cost_per_token": 8.8e-07, "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", "supported_regions": [ "global", @@ -40874,13 +41787,13 @@ "supports_tool_choice": true }, "vertex_ai/qwen/qwen3-coder-480b-a35b-instruct-maas": { - "input_cost_per_token": 1e-06, + "input_cost_per_token": 2.2e-07, "litellm_provider": "vertex_ai-qwen_models", "max_input_tokens": 262144, "max_output_tokens": 32768, "max_tokens": 32768, "mode": "chat", - "output_cost_per_token": 4e-06, + "output_cost_per_token": 1.8e-06, "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", "supported_regions": [ "global" @@ -41855,7 +42768,8 @@ "supports_prompt_caching": true, "supports_response_schema": false, "supports_tool_choice": true, - "supports_web_search": true + "supports_web_search": true, + "deprecation_date": "2026-05-15" }, "xai/grok-3-mini": { "cache_read_input_token_cost": 7.5e-08, @@ -41973,7 +42887,8 @@ "supports_function_calling": true, "supports_prompt_caching": true, "supports_tool_choice": true, - "supports_web_search": true + "supports_web_search": true, + "deprecation_date": "2026-05-15" }, "xai/grok-4-fast-reasoning": { "cache_read_input_token_cost": 5e-08, @@ -42042,7 +42957,8 @@ "supports_function_calling": true, "supports_prompt_caching": true, "supports_tool_choice": true, - "supports_web_search": true + "supports_web_search": true, + "deprecation_date": "2026-05-15" }, "xai/grok-4-1-fast": { "cache_read_input_token_cost": 5e-08, @@ -42370,7 +43286,8 @@ "output_cost_per_token_above_200k_tokens": 4e-06, "cache_read_input_token_cost_above_200k_tokens": 4e-07, "supports_response_schema": true, - "supports_vision": true + "supports_vision": true, + "deprecation_date": "2026-05-15" }, "xai/grok-code-fast-1": { "cache_read_input_token_cost": 2e-07, @@ -42390,7 +43307,8 @@ "output_cost_per_token_above_200k_tokens": 4e-06, "cache_read_input_token_cost_above_200k_tokens": 4e-07, "supports_response_schema": true, - "supports_vision": true + "supports_vision": true, + "deprecation_date": "2026-05-15" }, "xai/grok-code-fast-1-0825": { "cache_read_input_token_cost": 2e-07, @@ -42410,7 +43328,8 @@ "output_cost_per_token_above_200k_tokens": 4e-06, "cache_read_input_token_cost_above_200k_tokens": 4e-07, "supports_response_schema": true, - "supports_vision": true + "supports_vision": true, + "deprecation_date": "2026-05-15" }, "xai/grok-vision-beta": { "input_cost_per_image": 5e-06, @@ -46760,7 +47679,8 @@ "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_system_messages": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "deprecation_date": "2027-01-20" }, "gpt-realtime-whisper": { "input_cost_per_second": 0.0002833333333333333, @@ -48558,6 +49478,36 @@ "supports_reasoning": true, "supports_vision": false }, + "cognition/swe-1.6": { + "input_cost_per_token": 5e-07, + "output_cost_per_token": 2.5e-06, + "cache_read_input_token_cost": 2e-07, + "litellm_provider": "cognition", + "mode": "chat", + "supports_function_calling": true, + "supports_prompt_caching": true, + "source": "https://docs.devin.ai/windsurf/plugins/cascade/models" + }, + "cognition/swe-1.7": { + "input_cost_per_token": 5e-07, + "output_cost_per_token": 2.5e-06, + "cache_read_input_token_cost": 2e-07, + "litellm_provider": "cognition", + "mode": "chat", + "supports_function_calling": true, + "supports_prompt_caching": true, + "source": "https://docs.devin.ai/desktop/models" + }, + "cognition/swe-1.7-lightning": { + "input_cost_per_token": 2.5e-06, + "output_cost_per_token": 1.25e-05, + "cache_read_input_token_cost": 1e-06, + "litellm_provider": "cognition", + "mode": "chat", + "supports_function_calling": true, + "supports_prompt_caching": true, + "source": "https://docs.devin.ai/desktop/models" + }, "pinstripes/ps/glm-4.5-air": { "max_tokens": 128000, "max_input_tokens": 128000, @@ -48817,7 +49767,8 @@ "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, - "supports_xhigh_reasoning_effort": true + "supports_xhigh_reasoning_effort": true, + "supports_native_structured_output": true }, "claude-mythos-preview": { "cache_creation_input_token_cost": 1.25e-05, @@ -48850,7 +49801,8 @@ "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, - "supports_xhigh_reasoning_effort": true + "supports_xhigh_reasoning_effort": true, + "supports_native_structured_output": true }, "gemini/gemini-robotics-er-2-streaming-preview": { "input_cost_per_audio_token": 2e-06, @@ -48896,7 +49848,8 @@ "supports_function_calling": true, "supports_reasoning": true, "supports_response_schema": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_vision": true }, "mistral/labs-leanstral-1-5": { "input_cost_per_token": 0.0, @@ -49023,5 +49976,424 @@ } } ] + }, + "gemini/gemini-3.5-live-translate-preview": { + "input_cost_per_audio_token": 3.5e-06, + "input_cost_per_token": 3.5e-06, + "litellm_provider": "gemini", + "mode": "chat", + "output_cost_per_audio_token": 2.1e-05, + "output_cost_per_token": 2.1e-05, + "rpm": 10, + "source": "https://ai.google.dev/gemini-api/docs/pricing", + "supported_endpoints": [ + "/v1/realtime" + ], + "supported_modalities": [ + "audio" + ], + "supported_output_modalities": [ + "audio" + ], + "supports_audio_input": true, + "supports_audio_output": true, + "tpm": 250000 + }, + "perplexity/pplx-embed-context-v1-0.6b": { + "input_cost_per_token": 8e-09, + "litellm_provider": "perplexity", + "max_input_tokens": 32768, + "max_tokens": 32768, + "mode": "embedding", + "output_cost_per_token": 0.0, + "output_vector_size": 1024, + "source": "https://docs.perplexity.ai/getting-started/pricing" + }, + "perplexity/pplx-embed-context-v1-4b": { + "input_cost_per_token": 5e-08, + "litellm_provider": "perplexity", + "max_input_tokens": 32768, + "max_tokens": 32768, + "mode": "embedding", + "output_cost_per_token": 0.0, + "output_vector_size": 2560, + "source": "https://docs.perplexity.ai/getting-started/pricing" + }, + "voyage/voyage-4-large": { + "input_cost_per_token": 1.2e-07, + "litellm_provider": "voyage", + "max_input_tokens": 32000, + "max_tokens": 32000, + "mode": "embedding", + "output_cost_per_token": 0.0, + "output_vector_size": 1024, + "source": "https://docs.voyageai.com/docs/pricing" + }, + "voyage/voyage-4": { + "input_cost_per_token": 6e-08, + "litellm_provider": "voyage", + "max_input_tokens": 32000, + "max_tokens": 32000, + "mode": "embedding", + "output_cost_per_token": 0.0, + "output_vector_size": 1024, + "source": "https://docs.voyageai.com/docs/pricing" + }, + "voyage/voyage-4-lite": { + "input_cost_per_token": 2e-08, + "litellm_provider": "voyage", + "max_input_tokens": 32000, + "max_tokens": 32000, + "mode": "embedding", + "output_cost_per_token": 0.0, + "output_vector_size": 1024, + "source": "https://docs.voyageai.com/docs/pricing" + }, + "voyage/voyage-code-4": { + "input_cost_per_token": 1.2e-07, + "litellm_provider": "voyage", + "max_input_tokens": 32000, + "max_tokens": 32000, + "mode": "embedding", + "output_cost_per_token": 0.0, + "output_vector_size": 1024, + "source": "https://docs.voyageai.com/docs/pricing" + }, + "voyage/voyage-context-4": { + "input_cost_per_token": 1.2e-07, + "litellm_provider": "voyage", + "max_input_tokens": 120000, + "max_tokens": 120000, + "mode": "embedding", + "output_cost_per_token": 0.0, + "output_vector_size": 1024, + "source": "https://docs.voyageai.com/docs/pricing" + }, + "voyage/voyage-multimodal-3.5": { + "input_cost_per_token": 1.2e-07, + "litellm_provider": "voyage", + "max_input_tokens": 32000, + "max_tokens": 32000, + "mode": "embedding", + "output_cost_per_token": 0.0, + "output_vector_size": 1024, + "source": "https://docs.voyageai.com/docs/pricing", + "supports_embedding_image_input": true + }, + "fireworks_ai/accounts/fireworks/models/deepseek-v4-flash-0731": { + "cache_read_input_token_cost": 2.8e-08, + "input_cost_per_token": 1.4e-07, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 2.8e-07, + "source": "https://docs.fireworks.ai/serverless/pricing", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": false + }, + "fireworks_ai/accounts/fireworks/models/kimi-k3": { + "cache_read_input_token_cost": 3e-07, + "input_cost_per_token": 3e-06, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 1.5e-05, + "source": "https://docs.fireworks.ai/serverless/pricing", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "fireworks_ai/deepseek-v4-flash-0731": { + "cache_read_input_token_cost": 2.8e-08, + "input_cost_per_token": 1.4e-07, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 2.8e-07, + "source": "https://docs.fireworks.ai/serverless/pricing", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": false + }, + "fireworks_ai/glm-5p2-fast": { + "cache_read_input_token_cost": 2.1e-07, + "input_cost_per_token": 2.1e-06, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 6.6e-06, + "source": "https://docs.fireworks.ai/serverless/pricing", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": false + }, + "fireworks_ai/glm-5p2-fast-us": { + "cache_read_input_token_cost": 2.1e-07, + "input_cost_per_token": 2.1e-06, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 6.6e-06, + "source": "https://docs.fireworks.ai/serverless/pricing", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": false + }, + "fireworks_ai/kimi-k3": { + "cache_read_input_token_cost": 3e-07, + "input_cost_per_token": 3e-06, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 1.5e-05, + "source": "https://docs.fireworks.ai/serverless/pricing", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "fireworks_ai/kimi-k3-fast": { + "cache_read_input_token_cost": 4.5e-07, + "input_cost_per_token": 4.5e-06, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 2.25e-05, + "source": "https://docs.fireworks.ai/serverless/pricing", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "fireworks_ai/kimi-k3-us": { + "cache_read_input_token_cost": 3.3e-07, + "input_cost_per_token": 3.3e-06, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 1.65e-05, + "source": "https://docs.fireworks.ai/serverless/pricing", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "fireworks_ai/qwen3p8-max": { + "cache_read_input_token_cost": 2.5e-07, + "input_cost_per_token": 2e-06, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 6e-06, + "source": "https://docs.fireworks.ai/serverless/pricing", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "fireworks_ai/muse-glimmer-30b": { + "cache_read_input_token_cost": 4e-08, + "input_cost_per_token": 3.5e-07, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 131072, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "output_cost_per_token": 1.5e-06, + "source": "https://docs.fireworks.ai/serverless/pricing", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "fireworks_ai/nemotron-lightning-3p5-30b-a3b": { + "cache_read_input_token_cost": 1e-08, + "input_cost_per_token": 5e-08, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 262144, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 2e-07, + "source": "https://docs.fireworks.ai/serverless/pricing", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": false + }, + "fireworks_ai/nemotron-3-ultra-nvfp4": { + "cache_read_input_token_cost": 1.2e-07, + "input_cost_per_token": 6e-07, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 262144, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 2.4e-06, + "source": "https://docs.fireworks.ai/serverless/pricing", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": false + }, + "fireworks_ai/accounts/fireworks/models/muse-glimmer-30b": { + "cache_read_input_token_cost": 4e-08, + "input_cost_per_token": 3.5e-07, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 131072, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "output_cost_per_token": 1.5e-06, + "source": "https://docs.fireworks.ai/serverless/pricing", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "fireworks_ai/accounts/fireworks/models/nemotron-lightning-3p5-30b-a3b": { + "cache_read_input_token_cost": 1e-08, + "input_cost_per_token": 5e-08, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 262144, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 2e-07, + "source": "https://docs.fireworks.ai/serverless/pricing", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": false + }, + "fireworks_ai/accounts/fireworks/models/nemotron-3-ultra-nvfp4": { + "cache_read_input_token_cost": 1.2e-07, + "input_cost_per_token": 6e-07, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 262144, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 2.4e-06, + "source": "https://docs.fireworks.ai/serverless/pricing", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": false + }, + "fireworks_ai/accounts/fireworks/models/qwen3p8-max": { + "cache_read_input_token_cost": 2.5e-07, + "input_cost_per_token": 2e-06, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 6e-06, + "source": "https://docs.fireworks.ai/serverless/pricing", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "fireworks_ai/accounts/fireworks/routers/glm-5p2-fast": { + "cache_read_input_token_cost": 2.1e-07, + "input_cost_per_token": 2.1e-06, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 6.6e-06, + "source": "https://docs.fireworks.ai/serverless/pricing", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": false + }, + "fireworks_ai/accounts/fireworks/routers/glm-5p2-fast-us": { + "cache_read_input_token_cost": 2.1e-07, + "input_cost_per_token": 2.1e-06, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 6.6e-06, + "source": "https://docs.fireworks.ai/serverless/pricing", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": false + }, + "fireworks_ai/accounts/fireworks/routers/kimi-k3-fast": { + "cache_read_input_token_cost": 4.5e-07, + "input_cost_per_token": 4.5e-06, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 2.25e-05, + "source": "https://docs.fireworks.ai/serverless/pricing", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "fireworks_ai/accounts/fireworks/routers/kimi-k3-us": { + "cache_read_input_token_cost": 3.3e-07, + "input_cost_per_token": 3.3e-06, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 1.65e-05, + "source": "https://docs.fireworks.ai/serverless/pricing", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true } } diff --git a/litellm/provider_endpoints_support_backup.json b/litellm/provider_endpoints_support_backup.json index dd7712aabca..b4d635c0fba 100644 --- a/litellm/provider_endpoints_support_backup.json +++ b/litellm/provider_endpoints_support_backup.json @@ -528,6 +528,23 @@ "interactions": true } }, + "cognition": { + "display_name": "Cognition (`cognition`)", + "url": "https://docs.litellm.ai/docs/providers/cognition", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": false + } + }, "cohere": { "display_name": "Cohere (`cohere`)", "url": "https://docs.litellm.ai/docs/providers/cohere", diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 669924077e3..00cbd13cfdc 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -3746,6 +3746,8 @@ class ProxyErrorTypes(str, enum.Enum): Project does not have access to the model """ + model_cost_map_missing = "model_cost_map_missing" + expired_key = "expired_key" """ Key has expired diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index e1cc91df1c7..12d6b44a648 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -456,6 +456,103 @@ def _is_cost_explicitly_configured(model: str, llm_router: "Router") -> bool: return False +_EMPTY_COST_ENTRY: Final[Mapping[str, object]] = MappingProxyType({}) + + +def _is_positive_cost(value: object) -> bool: + return isinstance(value, (int, float)) and not isinstance(value, bool) and value > 0 + + +def _entry_has_priced_metric(entry: Mapping[str, object]) -> bool: + if entry.get("tiered_pricing") is not None: + return True + for key, value in entry.items(): + if "cost_per" not in key: + continue + if _is_positive_cost(value): + return True + if isinstance(value, dict) and any(_is_positive_cost(nested) for nested in value.values()): + return True + return False + + +def _entry_declares_price(entry: Mapping[str, object]) -> bool: + return any("cost_per" in key or key == "tiered_pricing" for key in entry) + + +def _model_group_has_pricing(model: str, llm_router: "Router") -> bool: + """ + A model group counts as priced when a deployment overrides any *cost_per* field or + tiered_pricing in its litellm_params, even at zero, or when its resolved model info carries + tiered_pricing or a positive price on any billed metric (tokens, characters, seconds, pages, + images, queries, ...), so models billed by a non-token metric are not treated as unpriced. + """ + for deployment in llm_router.get_model_list(model_name=model) or (): + litellm_params = deployment.get("litellm_params") or _EMPTY_COST_ENTRY + if _entry_declares_price(litellm_params): + return True + + model_id = (deployment.get("model_info") or _EMPTY_COST_ENTRY).get("id") + if model_id is None: + continue + + model_info = llm_router.get_deployment_model_info( + model_id=model_id, model_name=litellm_params.get("model") or "" + ) + if model_info is not None and _entry_has_priced_metric(model_info): + return True + + return False + + +def _group_declares_explicit_cost(model: str, llm_router: "Router") -> bool: + """ + Alias-aware counterpart to ``_is_cost_explicitly_configured``, which resolves the model group + the same way ``_model_group_has_pricing`` does. A deployment that prices itself through its + ``model_info`` block lands in the cost map under its deployment id rather than in its + litellm_params, and reaching that entry through the router's own resolution keeps an alias + pointing at such a group from being read as unpriced. + """ + for deployment in llm_router.get_model_list(model_name=model) or (): + model_id = (deployment.get("model_info") or _EMPTY_COST_ENTRY).get("id") + if model_id is None: + continue + raw_entry = litellm.model_cost.get(model_id, _EMPTY_COST_ENTRY) + if "input_cost_per_token" in raw_entry or "output_cost_per_token" in raw_entry: + return True + return False + + +def model_has_no_cost_mapping(model: str | None, llm_router: Router | None) -> bool: + if not model or llm_router is None: + return False + + if llm_router.get_model_group_info(model_group=model) is None: + return False + + if _model_group_has_pricing(model=model, llm_router=llm_router): + return False + + return not _group_declares_explicit_cost(model=model, llm_router=llm_router) + + +def _unpriced_models_in_request(model: str | list[str] | None, llm_router: Router | None) -> tuple[str, ...]: + candidates: Final = (model,) if isinstance(model, str) else tuple(model or ()) + return tuple( + candidate for candidate in candidates if model_has_no_cost_mapping(model=candidate, llm_router=llm_router) + ) + + +def _unpriced_models_block_message(models: tuple[str, ...]) -> str: + names: Final = ", ".join(f"'{model}'" for model in models) + subject: Final = f"Model {names} has" if len(models) == 1 else f"Models {names} have" + return ( + f"{subject} no pricing in the cost map, so litellm cannot price the request. " + "Requests for unpriced models are blocked because 'block_requests_for_models_without_pricing' " + "is enabled. Add pricing (input_cost_per_token/output_cost_per_token) to allow the request." + ) + + async def _run_project_checks( project_object: LiteLLM_ProjectTableCachedObj | None, _model: str | list[str] | None, @@ -726,6 +823,19 @@ async def common_checks( and (route in MODEL_DISCOVERY_ROUTES or not RouteChecks.is_llm_api_route(route=route)) ) + unpriced_models: Final = ( + _unpriced_models_in_request(model=_model, llm_router=llm_router) + if litellm.block_requests_for_models_without_pricing and RouteChecks.is_llm_api_route(route=route) + else () + ) + if unpriced_models: + raise ProxyException( + message=_unpriced_models_block_message(unpriced_models), + type=ProxyErrorTypes.model_cost_map_missing, + param="model", + code=status.HTTP_403_FORBIDDEN, + ) + # 1. If team is blocked if team_object is not None and team_object.blocked is True: raise Exception(f"Team={team_object.team_id} is blocked. Update via `/team/unblock` if you're an admin.") diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index a0b69ecb0bf..e5714ef66fb 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -43,6 +43,9 @@ from litellm.litellm_core_utils.llm_response_utils.get_headers import ( get_response_headers, ) from litellm.litellm_core_utils.safe_json_dumps import safe_dumps +from litellm.litellm_core_utils.streaming_handler import ( + backfill_missing_cache_usage_fields, +) from litellm.proxy._types import ProxyException, UserAPIKeyAuth from litellm.proxy.auth.auth_checks import can_key_call_resolved_model from litellm.proxy.auth.auth_utils import check_response_size_is_safe @@ -158,7 +161,7 @@ ProxyRouteType: TypeAlias = Literal[ "acancel_run", "adelete_run", ] -from litellm.types.utils import ServerToolUse +from litellm.llms.anthropic.chat.transformation import AnthropicConfig # Type alias for streaming chunk serializer (chunk after hooks + cost injection -> wire format) StreamChunkSerializer = Callable[[Any], str] @@ -274,6 +277,43 @@ def _deferred_stream_logging_is_armed(request_data: dict) -> bool: ) +def _assembled_model_came_from_a_later_chunk(chunks: list, assembled_model: object) -> bool: + """Report whether stream_chunk_builder picked a model the first chunk did not carry. + + Azure Model Router puts the routed model on the chunks after the first one, and the + proxy deliberately leaves those chunks unrestamped so the builder can recover it. + + A stored chunk that carries usage is a pre-restamp copy of the one the proxy saw, so + an alias-restamped stream reaches the builder with the same shape: a first chunk that + disagrees with the rest. Those two are only told apart by what the client asked for. + """ + first_chunk: Final = chunks[0] + first_chunk_model: Final = ( + first_chunk.get("model") if isinstance(first_chunk, dict) else getattr(first_chunk, "model", None) + ) + return ( + isinstance(first_chunk_model, str) + and isinstance(assembled_model, str) + and bool(assembled_model) + and assembled_model != first_chunk_model + ) + + +def _assembled_model_is_the_name_the_client_asked_for(request_data: dict, assembled_model: object) -> bool: + """Report whether the assembled model is the public name the proxy stamps onto chunks. + + That stamp is what leaves an unpriced alias on the partial response, so the deployment's + own model has to go back on before the row is costed. Pre-call processing rewrites + `request_data["model"]` for aliasing and routing, so the client's own name wins when it + is there, in the same order the proxy picks the name it stamps. + """ + client_requested_model: Final = request_data.get("_litellm_client_requested_model") + stamped_model: Final = ( + client_requested_model if isinstance(client_requested_model, str) else request_data.get("model") + ) + return isinstance(stamped_model, str) and assembled_model == stamped_model + + async def _bill_partial_streamed_spend_on_disconnect(request_data: dict, response: object) -> bool: """ A client disconnect throws GeneratorExit/CancelledError into the streaming @@ -324,6 +364,15 @@ async def _bill_partial_streamed_spend_on_disconnect(request_data: dict, respons return False if partial_response is None: return False + wrapper_model: Final = getattr(response, "model", None) + builder_recovered_the_routed_model: Final = _assembled_model_came_from_a_later_chunk( + chunks, partial_response.model + ) and not _assembled_model_is_the_name_the_client_asked_for(request_data, partial_response.model) + if isinstance(wrapper_model, str) and wrapper_model and not builder_recovered_the_routed_model: + partial_response.model = wrapper_model + partial_usage: Final = getattr(partial_response, "usage", None) + if isinstance(partial_usage, Usage): + backfill_missing_cache_usage_fields(partial_usage) try: await logging_obj.dispatch_success_handlers( partial_response, @@ -3321,7 +3370,9 @@ class ProxyBaseLLMRequestProcessing: str_so_far += str(chunk.get("content", "")) model_name = request_data.get("model", "") - chunk = ProxyBaseLLMRequestProcessing._process_chunk_with_cost_injection(chunk, model_name) + chunk = ProxyBaseLLMRequestProcessing._process_chunk_with_cost_injection( + chunk, model_name, request_data.get("litellm_logging_obj") + ) # Set before the yield: an async generator suspends at the yield, # so a GeneratorExit on client disconnect is raised there and any @@ -3418,20 +3469,27 @@ class ProxyBaseLLMRequestProcessing: @overload @staticmethod - def _process_chunk_with_cost_injection(chunk: bytes, model_name: str) -> bytes: ... + def _process_chunk_with_cost_injection( + chunk: bytes, model_name: str, litellm_logging_obj: LiteLLMLoggingObj | None = None + ) -> bytes: ... @overload @staticmethod - def _process_chunk_with_cost_injection(chunk: object, model_name: str) -> object: ... + def _process_chunk_with_cost_injection( + chunk: object, model_name: str, litellm_logging_obj: LiteLLMLoggingObj | None = None + ) -> object: ... @staticmethod - def _process_chunk_with_cost_injection(chunk: object, model_name: str) -> object: + def _process_chunk_with_cost_injection( + chunk: object, model_name: str, litellm_logging_obj: LiteLLMLoggingObj | None = None + ) -> object: """ Process a streaming chunk and inject cost information if enabled. Args: chunk: The streaming chunk (dict, str, bytes, or bytearray) model_name: Model name for cost calculation + litellm_logging_obj: The call's logging object, used for pricing Returns: The processed chunk with cost information injected if applicable @@ -3441,21 +3499,27 @@ class ProxyBaseLLMRequestProcessing: try: if isinstance(chunk, dict): - maybe_modified: Final = ProxyBaseLLMRequestProcessing._inject_cost_into_usage_dict(chunk, model_name) + maybe_modified: Final = ProxyBaseLLMRequestProcessing._inject_cost_into_usage_dict( + chunk, model_name, litellm_logging_obj + ) if maybe_modified is not None: return maybe_modified elif isinstance(chunk, (bytes, bytearray)): try: s: Final = chunk.decode("utf-8") if s.endswith(("\n\n", "\r\n\r\n")): - maybe_mod = ProxyBaseLLMRequestProcessing._inject_cost_into_sse_frame_str(s, model_name) + maybe_mod = ProxyBaseLLMRequestProcessing._inject_cost_into_sse_frame_str( + s, model_name, litellm_logging_obj + ) if maybe_mod is not None: return maybe_mod.encode("utf-8") except Exception: pass elif isinstance(chunk, str): # Try to parse SSE frame and inject cost into the data line - maybe_mod = ProxyBaseLLMRequestProcessing._inject_cost_into_sse_frame_str(chunk, model_name) + maybe_mod = ProxyBaseLLMRequestProcessing._inject_cost_into_sse_frame_str( + chunk, model_name, litellm_logging_obj + ) if maybe_mod is not None: # Ensure trailing frame separator return maybe_mod if maybe_mod.endswith("\n\n") else (maybe_mod + "\n\n") @@ -3466,13 +3530,16 @@ class ProxyBaseLLMRequestProcessing: return chunk @staticmethod - def _inject_cost_into_sse_frame_str(frame_str: str, model_name: str) -> str | None: + def _inject_cost_into_sse_frame_str( + frame_str: str, model_name: str, litellm_logging_obj: LiteLLMLoggingObj | None = None + ) -> str | None: """ Inject cost information into an SSE frame string by modifying the JSON in the 'data:' line. Args: frame_str: SSE frame string that may contain multiple lines model_name: Model name for cost calculation + litellm_logging_obj: The call's logging object, forwarded for pricing Returns: Modified SSE frame string with cost injected, or None if no modification needed @@ -3486,7 +3553,9 @@ class ProxyBaseLLMRequestProcessing: json_part = stripped_ln.split("data:", 1)[1].strip() if json_part and json_part != "[DONE]": obj = json.loads(json_part) - maybe_modified = ProxyBaseLLMRequestProcessing._inject_cost_into_usage_dict(obj, model_name) + maybe_modified = ProxyBaseLLMRequestProcessing._inject_cost_into_usage_dict( + obj, model_name, litellm_logging_obj + ) if maybe_modified is not None: lines[idx] = "data: " + safe_dumps(maybe_modified) + ("\r" if ln.endswith("\r") else "") return "\n".join(lines) @@ -3494,34 +3563,6 @@ class ProxyBaseLLMRequestProcessing: except Exception: return None - @staticmethod - def _anthropic_stream_usage_kwargs(usage: Mapping[str, Any]) -> Mapping[str, Any]: - prompt_tokens: Final = int(usage.get("input_tokens", 0) or 0) - completion_tokens: Final = int(usage.get("output_tokens", 0) or 0) - total_tokens: Final = int( - usage.get("total_tokens", prompt_tokens + completion_tokens) or (prompt_tokens + completion_tokens) - ) - web_search_requests: Final = usage.get("web_search_requests") - server_tool_use: Final = ( - ServerToolUse(web_search_requests=web_search_requests) if web_search_requests is not None else None - ) - return MappingProxyType( - { - key: value - for key, value in ( - ("prompt_tokens", prompt_tokens), - ("completion_tokens", completion_tokens), - ("total_tokens", total_tokens), - ("completion_tokens_details", usage.get("completion_tokens_details")), - ("prompt_tokens_details", usage.get("prompt_tokens_details")), - ("cache_creation_input_tokens", usage.get("cache_creation_input_tokens")), - ("cache_read_input_tokens", usage.get("cache_read_input_tokens")), - ("server_tool_use", server_tool_use), - ) - if value is not None - } - ) - @staticmethod def _openai_stream_usage_kwargs(usage: Mapping[str, Any]) -> Mapping[str, Any]: prompt_tokens: Final = int(usage.get("prompt_tokens", 0) or 0) @@ -3544,11 +3585,13 @@ class ProxyBaseLLMRequestProcessing: ) @staticmethod - def _stream_usage_kwargs_for_event(obj: Mapping[str, object], usage: Mapping[str, Any]) -> Mapping[str, Any] | None: + def _stream_usage_for_event(obj: Mapping[str, object], usage: Mapping[str, Any]) -> Usage | None: + # Anthropic reports input_tokens excluding cache tokens, so reuse the non-streaming + # transformation to total the prompt and keep the 5m/1h cache creation split if obj.get("type") == "message_delta": - return ProxyBaseLLMRequestProcessing._anthropic_stream_usage_kwargs(usage) + return AnthropicConfig().calculate_usage(usage_object=usage, reasoning_content=None) if obj.get("object") == "chat.completion.chunk": - return ProxyBaseLLMRequestProcessing._openai_stream_usage_kwargs(usage) + return Usage(**ProxyBaseLLMRequestProcessing._openai_stream_usage_kwargs(usage)) return None @staticmethod @@ -3563,7 +3606,54 @@ class ProxyBaseLLMRequestProcessing: return None @staticmethod - def _inject_cost_into_usage_dict(obj: dict, model_name: str) -> dict | None: + def _logging_obj_cost_or_none( + model_response: ModelResponse, litellm_logging_obj: LiteLLMLoggingObj + ) -> float | None: + # Pricing a frame stamps cost_breakdown and, on failure, the cost-failure debug key onto + # the live logging object. The pass-through handlers never recompute either one, so a + # frame-derived breakdown would outlive the stream and land in the spend log. Snapshot + # both and put them back, so pricing here stays a read as far as the request is concerned + breakdown_before: Final = getattr(litellm_logging_obj, "cost_breakdown", None) + call_details: Final = getattr(litellm_logging_obj, "model_call_details", None) + debug_key: Final = "response_cost_failure_debug_information" + debug_missing: Final = object() + debug_before: Final = call_details.get(debug_key, debug_missing) if isinstance(call_details, dict) else None + try: + cost: Final = litellm_logging_obj._response_cost_calculator(result=model_response) # pyright: ignore[reportPrivateUsage] # reuse the call's own cost calc for pricing parity with the logging callback + except Exception: # noqa: BLE001 # a pricing failure falls back to model-name pricing instead of breaking the stream + return None + finally: + if hasattr(litellm_logging_obj, "cost_breakdown"): + litellm_logging_obj.cost_breakdown = breakdown_before + if isinstance(call_details, dict): + if debug_before is debug_missing: + call_details.pop(debug_key, None) + else: + call_details[debug_key] = debug_before + return float(cost) if isinstance(cost, (int, float)) and not isinstance(cost, bool) else None + + @staticmethod + def _streamed_usage_cost( + model_response: ModelResponse, + model_name: str, + service_tier: str | None, + litellm_logging_obj: LiteLLMLoggingObj | None, + ) -> float | None: + # Pricing via the logging object inherits the deployment's custom pricing, so the + # streamed cost matches what the logging callback records instead of sticker price + cost_from_logging_obj: Final = ( + ProxyBaseLLMRequestProcessing._logging_obj_cost_or_none(model_response, litellm_logging_obj) + if litellm_logging_obj is not None + else None + ) + if cost_from_logging_obj is not None: + return cost_from_logging_obj + return ProxyBaseLLMRequestProcessing._completion_cost_or_none(model_response, model_name, service_tier) + + @staticmethod + def _inject_cost_into_usage_dict( + obj: dict, model_name: str, litellm_logging_obj: LiteLLMLoggingObj | None = None + ) -> dict | None: """ Inject cost information into the usage object of a streamed usage event (Anthropic ``message_delta`` or OpenAI ``chat.completion.chunk``). @@ -3571,6 +3661,7 @@ class ProxyBaseLLMRequestProcessing: Args: obj: Dictionary containing the SSE event data model_name: Model name for cost calculation + litellm_logging_obj: The call's logging object, used for pricing Returns: Modified dictionary with cost injected, or None if no modification needed @@ -3578,14 +3669,15 @@ class ProxyBaseLLMRequestProcessing: usage: Final = obj.get("usage") if not isinstance(usage, dict): return None - usage_kwargs: Final = ProxyBaseLLMRequestProcessing._stream_usage_kwargs_for_event(obj, usage) - if usage_kwargs is None: + stream_usage: Final = ProxyBaseLLMRequestProcessing._stream_usage_for_event(obj, usage) + if stream_usage is None: return None service_tier: Final = obj.get("service_tier") - cost_val: Final = ProxyBaseLLMRequestProcessing._completion_cost_or_none( - ModelResponse(usage=Usage(**usage_kwargs)), + cost_val: Final = ProxyBaseLLMRequestProcessing._streamed_usage_cost( + ModelResponse(usage=stream_usage), model_name, service_tier if isinstance(service_tier, str) else None, + litellm_logging_obj, ) if cost_val is None: return None diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py index b28b7291a4c..8fcb184b26a 100644 --- a/litellm/proxy/common_utils/reset_budget_job.py +++ b/litellm/proxy/common_utils/reset_budget_job.py @@ -4,6 +4,7 @@ import time from collections.abc import Awaitable, Callable, Iterable, Mapping, Sequence from dataclasses import dataclass from datetime import datetime, timezone +from enum import Enum from types import MappingProxyType from typing import Final, Literal, Protocol, TypeVar, assert_never @@ -14,7 +15,9 @@ from litellm.constants import ( GLOBAL_PROXY_SPEND_CACHE_KEY, LITELLM_PROXY_BUDGET_NAME, RESET_BUDGET_JOB_BATCH_SIZE, + RESET_BUDGET_JOB_LOCK_TTL_SECONDS, RESET_BUDGET_JOB_MAX_CHUNKS_PER_RUN, + RESET_BUDGET_JOB_NAME, ) from litellm.proxy._types import ( DB_RETRY_SAFE_ERROR_TYPES, @@ -30,6 +33,7 @@ from litellm.proxy.common_utils.timezone_utils import ( get_budget_reset_settings, ) from litellm.proxy.common_utils.user_api_key_cache import tag_cache_key +from litellm.proxy.db.db_transaction_queue.pod_lock_manager import PodLockManager from litellm.proxy.db.exception_handler import call_with_db_reconnect_retry from litellm.proxy.utils import PrismaClient, ProxyLogging from litellm.repositories.organization_repository import OrganizationRepository @@ -195,12 +199,94 @@ async def _run_phase_in_chunks(process_chunk: Callable[[], Awaitable[_ChunkOutco return +@dataclass(frozen=True, slots=True) +class _LazyJson: + """Serialize only if a log record is actually emitted. + + ``logger.debug("... %s", json.dumps(rows))`` evaluates the dump before the + logger decides to drop the record, so a chunk of rows is serialized on the + event loop on every tick at any log level. Passing this instead defers the + work to the formatter. + """ + + value: object + + def __str__(self) -> str: + return json.dumps(self.value, indent=4, default=str) + + +class _Lease(Enum): + """Whether this pod may sweep, and whether it owes a lock release.""" + + LEADER = "leader" + UNGUARDED = "unguarded" + FOLLOWER = "follower" + + +async def _write_key_windows(prisma_client: PrismaClient, row_id: str, payload: str) -> None: + await VerificationTokenRepository(prisma_client).table.update( + where={"token": row_id}, + data={"budget_limits": payload}, + ) + + +async def _write_team_windows(prisma_client: PrismaClient, row_id: str, payload: str) -> None: + await TeamRepository(prisma_client).table.update( + where={"team_id": row_id}, + data={"budget_limits": payload}, + ) + + +@dataclass(frozen=True, slots=True) +class _WindowSource: + """A table whose rows carry their own per-window budget limits.""" + + table: str + id_column: str + counter_prefix: str + log_subject: str + retry_subject: str + write: Callable[[PrismaClient, str, str], Awaitable[None]] + + def page_query(self) -> str: + """One keyset page, ordered by the primary key so the cursor never repeats a row. + + prisma-client-python cannot null-filter a ``Json?`` column (no DbNull / + JsonNull sentinel, RobertCraigie/prisma-client-py#714), so the read stays + raw SQL; the table and column names are module constants, never input. + Writes still go through the ORM. + """ + return ( + f'SELECT {self.id_column}, budget_limits FROM "{self.table}" ' + f"WHERE budget_limits IS NOT NULL AND {self.id_column} > $1 " + f"ORDER BY {self.id_column} LIMIT $2" + ) + + +_WINDOW_SOURCES: Final[tuple[_WindowSource, ...]] = ( + _WindowSource( + table="LiteLLM_VerificationToken", + id_column="token", + counter_prefix="spend:key", + log_subject="keys", + retry_subject="key", + write=_write_key_windows, + ), + _WindowSource( + table="LiteLLM_TeamTable", + id_column="team_id", + counter_prefix="spend:team", + log_subject="teams", + retry_subject="team", + write=_write_team_windows, + ), +) + + def _budget_cascade_event_metadata(cascade: _BudgetCascade) -> dict[str, object]: return { "num_budgets_found": len(cascade.budgets), - "budgets_found": json.dumps(cascade.budgets, indent=4, default=str), "num_endusers_found": len(cascade.endusers), - "endusers_found": json.dumps(cascade.endusers, indent=4, default=str), } @@ -214,10 +300,61 @@ class ResetBudgetJob: proxy_logging_obj: ProxyLogging, prisma_client: PrismaClient, reset_settings: BudgetResetSettings | None = None, + pod_lock_manager: PodLockManager | None = None, ): self.proxy_logging_obj: ProxyLogging = proxy_logging_obj self.prisma_client: PrismaClient = prisma_client self.reset_settings: BudgetResetSettings = reset_settings or get_budget_reset_settings() + self.pod_lock_manager: PodLockManager | None = pod_lock_manager + + async def _lease_is_held(self, lock_manager: PodLockManager) -> bool: + """True only when the lease is readable and someone holds it. + + An unreadable lock reports as unheld so the caller sweeps rather than + skipping; being wrong here costs a duplicate sweep, and the alternative + strands every expired budget at its cap. + """ + if lock_manager.redis_cache is None: + return False + try: + lock_key: Final = lock_manager.get_redis_lock_key(RESET_BUDGET_JOB_NAME) + return bool(await lock_manager.redis_cache.async_get_cache(lock_key)) + except Exception as exc: # noqa: BLE001 # an unreadable lease must not strand the sweep + verbose_proxy_logger.warning("Reset budget job: could not read the reset lease: %s", exc) + return False + + async def _acquire_lease(self) -> _Lease: + """Elect one sweeper per tick. + + Every pod schedules this job, and each one otherwise re-reads the whole + due population and writes it back at the same calendar boundary, so a + fleet multiplies one sweep's Postgres load by its replica count. A + deployment with no Redis-backed lock manager runs unguarded, as it + always has. + """ + lock_manager: Final = self.pod_lock_manager + if lock_manager is None or lock_manager.redis_cache is None: + return _Lease.UNGUARDED + + if await lock_manager.acquire_lock( + cronjob_id=RESET_BUDGET_JOB_NAME, + ttl=RESET_BUDGET_JOB_LOCK_TTL_SECONDS, + ): + return _Lease.LEADER + + if await self._lease_is_held(lock_manager): + verbose_proxy_logger.debug("Reset budget job: another pod holds the reset lease, skipping this tick") + return _Lease.FOLLOWER + + # acquire_lock reports contention and an unreachable Redis identically, so + # treating a failed acquire as contention would skip the sweep on every pod + # at once for as long as Redis is down. Sweeping unguarded costs duplicate + # work; not sweeping leaves every expired budget pinned at its cap. + verbose_proxy_logger.warning( + "Reset budget job: could not take the reset lease and no other pod holds it, " + "sweeping unguarded rather than skipping the tick" + ) + return _Lease.UNGUARDED async def reset_budget( self, @@ -228,15 +365,25 @@ class ResetBudgetJob: Resets their spend Updates db + + Runs on one pod per tick where a Redis lease is available. """ if self.prisma_client is None: return - await self.reset_budget_for_litellm_keys() - await self.reset_budget_for_litellm_users() - await self.reset_budget_for_litellm_teams() - await self.reset_budget_for_litellm_budget_table() - await self.reset_budget_windows() + lease: Final = await self._acquire_lease() + if lease is _Lease.FOLLOWER: + return + + try: + await self.reset_budget_for_litellm_keys() + await self.reset_budget_for_litellm_users() + await self.reset_budget_for_litellm_teams() + await self.reset_budget_for_litellm_budget_table() + await self.reset_budget_windows() + finally: + if lease is _Lease.LEADER and self.pod_lock_manager is not None: + await self.pod_lock_manager.release_lock(cronjob_id=RESET_BUDGET_JOB_NAME) async def _with_db_retry(self, operation: Callable[[], Awaitable[_RowT]], *, reason: str) -> _RowT: """Reconnect and retry once on a transport error, so a dropped connection @@ -647,7 +794,7 @@ class ResetBudgetJob: ), reason="reset_budget_read_keys_failure", ) - verbose_proxy_logger.debug("Keys to reset %s", json.dumps(keys_to_reset, indent=4, default=str)) + verbose_proxy_logger.debug("Keys to reset %s", _LazyJson(keys_to_reset)) updated_keys: Final[list[LiteLLM_VerificationToken]] = [] failed_keys: Final = [] if keys_to_reset is not None and len(keys_to_reset) > 0: @@ -666,7 +813,7 @@ class ResetBudgetJob: failed_keys.append({"key": key, "error": str(e)}) verbose_proxy_logger.exception("Failed to reset budget for key: %s", key) - verbose_proxy_logger.debug("Updated keys %s", json.dumps(updated_keys, indent=4, default=str)) + verbose_proxy_logger.debug("Updated keys %s", _LazyJson(updated_keys)) if updated_keys: await self._write_key_reset_updates(updated_keys=updated_keys) @@ -691,7 +838,6 @@ class ResetBudgetJob: end_time=end_time, event_metadata={ "num_keys_found": len(keys_to_reset) if keys_to_reset else 0, - "keys_found": json.dumps(keys_to_reset, indent=4, default=str), }, ) return outcome @@ -705,11 +851,8 @@ class ResetBudgetJob: end_time=end_time, event_metadata={ "num_keys_found": len(keys_to_reset) if keys_to_reset else 0, - "keys_found": json.dumps(keys_to_reset, indent=4, default=str), "num_keys_updated": len(updated_keys), - "keys_updated": json.dumps(updated_keys, indent=4, default=str), "num_keys_failed": len(failed_keys), - "keys_failed": json.dumps(failed_keys, indent=4, default=str), }, ) ) @@ -725,7 +868,6 @@ class ResetBudgetJob: end_time=end_time, event_metadata={ "num_keys_found": len(keys_to_reset) if keys_to_reset else 0, - "keys_found": json.dumps(keys_to_reset, indent=4, default=str), }, ) ) @@ -777,7 +919,7 @@ class ResetBudgetJob: failed_users.append({"user": user, "error": str(e)}) verbose_proxy_logger.exception("Failed to reset budget for user: %s", user) - verbose_proxy_logger.debug("Updated users %s", json.dumps(updated_users, indent=4, default=str)) + verbose_proxy_logger.debug("Updated users %s", _LazyJson(updated_users)) if updated_users: await self._write_user_reset_updates(updated_users=updated_users) for u in updated_users: @@ -805,7 +947,6 @@ class ResetBudgetJob: end_time=end_time, event_metadata={ "num_users_found": len(users_to_reset) if users_to_reset else 0, - "users_found": json.dumps(users_to_reset, indent=4, default=str), }, ) return outcome @@ -819,11 +960,8 @@ class ResetBudgetJob: end_time=end_time, event_metadata={ "num_users_found": len(users_to_reset) if users_to_reset else 0, - "users_found": json.dumps(users_to_reset, indent=4, default=str), "num_users_updated": len(updated_users), - "users_updated": json.dumps(updated_users, indent=4, default=str), "num_users_failed": len(failed_users), - "users_failed": json.dumps(failed_users, indent=4, default=str), }, ) ) @@ -839,7 +977,6 @@ class ResetBudgetJob: end_time=end_time, event_metadata={ "num_users_found": len(users_to_reset) if users_to_reset else 0, - "users_found": json.dumps(users_to_reset, indent=4, default=str), }, ) ) @@ -891,7 +1028,7 @@ class ResetBudgetJob: failed_teams.append({"team": team, "error": str(e)}) verbose_proxy_logger.exception("Failed to reset budget for team: %s", team) - verbose_proxy_logger.debug("Updated teams %s", json.dumps(updated_teams, indent=4, default=str)) + verbose_proxy_logger.debug("Updated teams %s", _LazyJson(updated_teams)) if updated_teams: await self._write_team_reset_updates(updated_teams=updated_teams) for t in updated_teams: @@ -917,7 +1054,6 @@ class ResetBudgetJob: end_time=end_time, event_metadata={ "num_teams_found": len(teams_to_reset) if teams_to_reset else 0, - "teams_found": json.dumps(teams_to_reset, indent=4, default=str), }, ) return outcome @@ -931,11 +1067,8 @@ class ResetBudgetJob: end_time=end_time, event_metadata={ "num_teams_found": len(teams_to_reset) if teams_to_reset else 0, - "teams_found": json.dumps(teams_to_reset, indent=4, default=str), "num_teams_updated": len(updated_teams), - "teams_updated": json.dumps(updated_teams, indent=4, default=str), "num_teams_failed": len(failed_teams), - "teams_failed": json.dumps(failed_teams, indent=4, default=str), }, ) ) @@ -951,7 +1084,6 @@ class ResetBudgetJob: end_time=end_time, event_metadata={ "num_teams_found": len(teams_to_reset) if teams_to_reset else 0, - "teams_found": json.dumps(teams_to_reset, indent=4, default=str), }, ) ) @@ -995,82 +1127,82 @@ class ResetBudgetJob: from litellm.proxy.proxy_server import spend_counter_cache now: Final = datetime.utcnow() + for source in _WINDOW_SOURCES: + try: + await self._reset_windows_for(source=source, now=now, spend_counter_cache=spend_counter_cache) + except Exception as e: + verbose_proxy_logger.exception("Failed to reset budget windows for %s: %s", source.log_subject, e) - # Note on raw SQL: prisma-client-python does not support null-filtering - # on `Json?` columns (no DbNull/JsonNull sentinel — see - # RobertCraigie/prisma-client-py#714). We use `query_raw` with - # `IS NOT NULL` so we don't materialize every key/team row on each - # tick of the reset job. Writes still go through the ORM. + async def _reset_windows_for( + self, + source: _WindowSource, + now: datetime, + spend_counter_cache: DualCache, + ) -> None: + """Walk one table's windowed rows a page at a time, to the end. - # --- Keys --- - try: - key_rows: Final = await self._with_db_retry( - lambda: self.prisma_client.db.query_raw( - 'SELECT token, budget_limits FROM "LiteLLM_VerificationToken" WHERE budget_limits IS NOT NULL' - ), - reason="reset_budget_read_key_windows_failure", + Paging is what bounds the memory: the previous form pulled every row + carrying budget_limits into one result set on every tick, which grows + with the deployment's key count and is paid on the event loop. + + The walk deliberately has no per-run page cap. A cap has to remember + where it stopped, and that position cannot live in the process: the + lease is released after each sweep, so the next tick can elect a + different pod whose own position is unset. It would restart at the first + row and never reach the tail, pinning those windows at their cap for + good. The cursor strictly advances, so the walk terminates on its own + without needing a bound. + """ + cursor = "" + while True: + next_cursor = await self._reset_window_page( + source=source, + cursor=cursor, + now=now, + spend_counter_cache=spend_counter_cache, ) - for row in key_rows: - raw = row["budget_limits"] - if not raw: - continue - windows: list = raw if isinstance(raw, list) else json.loads(raw) - changed = False - for window in windows: - counter_key = f"spend:key:{row['token']}:window:{window['budget_duration']}" - if await ResetBudgetJob._reset_expired_window( - window, - counter_key, - spend_counter_cache, - now, - self.reset_settings, - ): - changed = True - if changed: - await self._with_db_write_retry( - lambda: VerificationTokenRepository(self.prisma_client).table.update( - where={"token": row["token"]}, - data={"budget_limits": json.dumps(windows)}, - ), - reason="reset_budget_write_key_windows_failure", - ) - except Exception as e: - verbose_proxy_logger.exception("Failed to reset budget windows for keys: %s", e) + if next_cursor is None: + return + cursor = next_cursor - # --- Teams --- - try: - team_rows: Final = await self._with_db_retry( - lambda: self.prisma_client.db.query_raw( - 'SELECT team_id, budget_limits FROM "LiteLLM_TeamTable" WHERE budget_limits IS NOT NULL' - ), - reason="reset_budget_read_team_windows_failure", - ) - for row in team_rows: - raw = row["budget_limits"] - if not raw: - continue - windows = raw if isinstance(raw, list) else json.loads(raw) - changed = False - for window in windows: - counter_key = f"spend:team:{row['team_id']}:window:{window['budget_duration']}" - if await ResetBudgetJob._reset_expired_window( - window, - counter_key, - spend_counter_cache, - now, - self.reset_settings, - ): - changed = True - if changed: - await self._with_db_write_retry( - lambda: TeamRepository(self.prisma_client).table.update( - where={"team_id": row["team_id"]}, - data={"budget_limits": json.dumps(windows)}, - ), - reason="reset_budget_write_team_windows_failure", - ) - except Exception as e: - verbose_proxy_logger.exception("Failed to reset budget windows for teams: %s", e) + async def _reset_window_page( + self, + source: _WindowSource, + cursor: str, + now: datetime, + spend_counter_cache: DualCache, + ) -> str | None: + """Reset one page of windows; return the next cursor, or None when drained.""" + rows: Final = await self._with_db_retry( + lambda: self.prisma_client.db.query_raw(source.page_query(), cursor, RESET_BUDGET_JOB_BATCH_SIZE), + reason=f"reset_budget_read_{source.retry_subject}_windows_failure", + ) + for row in rows: + raw = row["budget_limits"] + if not raw: + continue + row_id: str = row[source.id_column] + windows: list = raw if isinstance(raw, list) else json.loads(raw) + changed = False + for window in windows: + counter_key = f"{source.counter_prefix}:{row_id}:window:{window['budget_duration']}" + if await ResetBudgetJob._reset_expired_window( + window, + counter_key, + spend_counter_cache, + now, + self.reset_settings, + ): + changed = True + if changed: + await self._with_db_write_retry( + lambda: source.write(self.prisma_client, row_id, json.dumps(windows)), + reason=f"reset_budget_write_{source.retry_subject}_windows_failure", + ) + + if len(rows) < RESET_BUDGET_JOB_BATCH_SIZE: + return None + return rows[-1][source.id_column] @staticmethod async def _reset_budget_common( diff --git a/litellm/proxy/management_endpoints/cost_tracking_settings.py b/litellm/proxy/management_endpoints/cost_tracking_settings.py index 56439172b63..204051c3715 100644 --- a/litellm/proxy/management_endpoints/cost_tracking_settings.py +++ b/litellm/proxy/management_endpoints/cost_tracking_settings.py @@ -15,6 +15,7 @@ from dataclasses import dataclass from typing import Final from fastapi import APIRouter, Depends, HTTPException +from pydantic import BaseModel import litellm from litellm._logging import verbose_proxy_logger @@ -439,6 +440,76 @@ async def update_cost_margin_config( ) +class BlockUnpricedModelsRequest(BaseModel): + enabled: bool + + +class BlockUnpricedModelsResponse(BaseModel): + enabled: bool + + +@router.get( + "/config/block_requests_for_models_without_pricing", + tags=("Cost Tracking",), + dependencies=(Depends(user_api_key_auth),), + response_model=BlockUnpricedModelsResponse, +) +async def get_block_requests_for_models_without_pricing() -> BlockUnpricedModelsResponse: + return BlockUnpricedModelsResponse(enabled=bool(litellm.block_requests_for_models_without_pricing)) + + +@router.patch( + "/config/block_requests_for_models_without_pricing", + tags=("Cost Tracking",), + dependencies=(Depends(user_api_key_auth),), + response_model=BlockUnpricedModelsResponse, +) +async def update_block_requests_for_models_without_pricing( + request: BlockUnpricedModelsRequest, +) -> BlockUnpricedModelsResponse: + from litellm.proxy.proxy_server import ( + prisma_client, + proxy_config, + store_model_in_db, + ) + + if prisma_client is None: + raise HTTPException( + status_code=500, + detail={ # mutable-ok: HTTPException detail must be a plain mapping + "error": CommonProxyErrors.db_not_connected_error.value + }, + ) + + if store_model_in_db is not True: + raise HTTPException( + status_code=500, + detail={ # mutable-ok: HTTPException detail must be a plain mapping + "error": "Set `'STORE_MODEL_IN_DB='True'` in your env to enable this feature." + }, + ) + + try: + config = await proxy_config.get_config() + if "litellm_settings" not in config: + config["litellm_settings"] = {} # mutable-ok: config is a plain-dict payload for save_config + config["litellm_settings"]["block_requests_for_models_without_pricing"] = request.enabled + await proxy_config.save_config(new_config=config) + + litellm.block_requests_for_models_without_pricing = request.enabled + verbose_proxy_logger.info("Updated block_requests_for_models_without_pricing: %s", request.enabled) + + return BlockUnpricedModelsResponse(enabled=request.enabled) + except Exception as e: # noqa: BLE001 # any config persistence failure must surface as a 500 response, not a crash + verbose_proxy_logger.error("Error updating block_requests_for_models_without_pricing: %s", e) + raise HTTPException( + status_code=500, + detail={ # mutable-ok: HTTPException detail must be a plain mapping + "error": f"Failed to update setting: {e!s}" + }, + ) + + @router.post( "/cost/estimate", tags=["Cost Tracking"], diff --git a/litellm/proxy/management_endpoints/scim/scim_v2.py b/litellm/proxy/management_endpoints/scim/scim_v2.py index 8d255859571..7183e6cb402 100644 --- a/litellm/proxy/management_endpoints/scim/scim_v2.py +++ b/litellm/proxy/management_endpoints/scim/scim_v2.py @@ -487,13 +487,18 @@ class _UnknownMember(NamedTuple): value: str -_ClassifiedGroupMember = Union[_ResolvedUserMember, _SkippedGroupMember, _UnknownMember] +class _AmbiguousMember(NamedTuple): + value: str + + +_ClassifiedGroupMember = Union[_ResolvedUserMember, _SkippedGroupMember, _UnknownMember, _AmbiguousMember] class _PartitionedMembers(NamedTuple): resolved_ids: tuple[str, ...] skipped: tuple[_SkippedGroupMember, ...] unknown_ids: tuple[str, ...] + ambiguous_values: tuple[str, ...] def _member_value(member: SCIMMember) -> str: @@ -536,6 +541,44 @@ def _team_metadata_has_scim_provenance(team_metadata: object) -> bool: return bool(fields.get(SCIM_MANAGED_TEAM_METADATA_KEY)) or fields.get(SCIM_TEAM_DATA_METADATA_KEY) is not None +class _CaseInsensitiveMatch(TypedDict): + equals: ReadOnly[str] + mode: ReadOnly[str] + + +async def _users_named_by_member_value( + value: str, prisma_client: PrismaClient, *, take: int | None = 2 +) -> tuple[str, ...]: + """Every user id this member value names, by SSO identity or by email. + + Both fields are searched in one pass, because searching either first would hide a + value that names one account by its SSO identity and another by its email, and + hand the group to whichever field was searched first. + + They are not compared alike. An email is matched the way ``new_user`` matches one + before it accepts a new account, case-insensitively: matching more strictly than + the layer that would reject the placeholder is what turned a member id whose + casing differed from the stored email into a 500 on the whole push. An SSO + identity is matched exactly, because OIDC defines ``sub`` as case-sensitive and + nothing folds its case on the way in, so treating two subjects that differ in case + as one would hand the group to an account the provider never named. + + ``take`` bounds the read for a caller that only needs to know whether the value + names one account or several; ``user_email`` carries no index, so letting the scan + stop early is worth the two rows. A caller that has to know *which* accounts, as a + removal does, passes None. That set is the accounts sharing one identity, which is + a handful at worst. + """ + subject: Final = value.strip() + email: Final[_CaseInsensitiveMatch] = {"equals": subject, "mode": "insensitive"} + rows: Final = await _table(UserRepository(prisma_client)).find_many( + # mutable-ok: the Prisma serializer requires concrete dicts and a concrete list + where={"OR": [{"sso_user_id": subject}, {"user_email": email}]}, + take=take, + ) + return tuple(dict.fromkeys(row.user_id for row in rows)) + + async def _classify_group_member(member: SCIMMember, prisma_client: PrismaClient) -> _ClassifiedGroupMember: """ Decide what a single SCIM group member refers to. @@ -557,6 +600,20 @@ async def _classify_group_member(member: SCIMMember, prisma_client: PrismaClient one the identity provider writes. An id the IdP called a User is a user even if some team happens to share the id, and a team created here rather than through SCIM is not evidence of anything about the member. + + When those checks miss on an otherwise user-shaped member, its value is looked + up as an SSO identity or an email, and a match resolves to that user's + ``user_id``. A value that names more than one account is ambiguous rather than + unknown: it names a real person we cannot identify, so it is neither guessed at + nor provisioned. + + An exact ``user_id`` hit is checked the same way rather than trusted outright. A + value can be one account's id and another's SSO identity or email, and taking the + id on sight would hand the group to whichever account happened to be keyed by it. + The placeholders this bug provisioned are that shape exactly, since they are keyed + by the very id the provider keeps pushing, so on a tenant that already has them + the membership is refused and named rather than silently landing on the + placeholder again. """ value: Final = _member_value(member) member_type: Final = _normalized_member_type(member) @@ -566,6 +623,18 @@ async def _classify_group_member(member: SCIMMember, prisma_client: PrismaClient user: Final = await _table(UserRepository(prisma_client)).find_unique(where={"user_id": value}) if user is not None: + shared_with: Final = tuple( + other for other in await _users_named_by_member_value(value, prisma_client) if other != value + ) + if shared_with: + verbose_proxy_logger.warning( + "SCIM: group member '%s' is one account's user id and is also account '%s' by SSO identity or email, " + "so the membership cannot be attributed. A placeholder an earlier release provisioned under this id " + "looks exactly like this and should be deleted so the real account can be matched", + value, + shared_with[0], + ) + return _AmbiguousMember(value=value) return _ResolvedUserMember(user_id=value) if member_type is not None and member_type != "user": @@ -576,6 +645,22 @@ async def _classify_group_member(member: SCIMMember, prisma_client: PrismaClient if team is not None and _team_metadata_has_scim_provenance(team.metadata): return _SkippedGroupMember(value=value, reason="existing_team") + named: Final = await _users_named_by_member_value(value, prisma_client) + if len(named) == 1: + verbose_proxy_logger.info( + "SCIM: group member '%s' matched user_id '%s' by SSO identity or email", + value, + named[0], + ) + return _ResolvedUserMember(user_id=named[0]) + if len(named) > 1: + verbose_proxy_logger.warning( + "SCIM: group member '%s' names more than one account by SSO identity or email and cannot be resolved " + "unambiguously", + value, + ) + return _AmbiguousMember(value=value) + return _UnknownMember(value=value) @@ -583,11 +668,13 @@ def _bucketed_member(entry: _ClassifiedGroupMember) -> _PartitionedMembers: """The single-member partition one classified entry contributes.""" match entry: case _ResolvedUserMember(user_id=user_id): - return _PartitionedMembers(resolved_ids=(user_id,), skipped=(), unknown_ids=()) + return _PartitionedMembers(resolved_ids=(user_id,), skipped=(), unknown_ids=(), ambiguous_values=()) case _SkippedGroupMember(): - return _PartitionedMembers(resolved_ids=(), skipped=(entry,), unknown_ids=()) + return _PartitionedMembers(resolved_ids=(), skipped=(entry,), unknown_ids=(), ambiguous_values=()) case _UnknownMember(value=value): - return _PartitionedMembers(resolved_ids=(), skipped=(), unknown_ids=(value,)) + return _PartitionedMembers(resolved_ids=(), skipped=(), unknown_ids=(value,), ambiguous_values=()) + case _AmbiguousMember(value=value): + return _PartitionedMembers(resolved_ids=(), skipped=(), unknown_ids=(), ambiguous_values=(value,)) case _: assert_never(entry) @@ -599,6 +686,7 @@ def _partition_classified_members(classified: Iterable[_ClassifiedGroupMember]) resolved_ids=tuple(chain.from_iterable(bucket.resolved_ids for bucket in bucketed)), skipped=tuple(chain.from_iterable(bucket.skipped for bucket in bucketed)), unknown_ids=tuple(chain.from_iterable(bucket.unknown_ids for bucket in bucketed)), + ambiguous_values=tuple(chain.from_iterable(bucket.ambiguous_values for bucket in bucketed)), ) @@ -608,7 +696,7 @@ def _admitted_member_id(entry: _ClassifiedGroupMember, created_ids: frozenset[st return user_id case _UnknownMember(value=value): return value if value in created_ids else None - case _SkippedGroupMember(): + case _SkippedGroupMember() | _AmbiguousMember(): return None case _: assert_never(entry) @@ -662,6 +750,70 @@ async def _ensure_group_member_user( raise HTTPException(status_code=500, detail=detail) +def _roster_entries_named_by(value: str, roster: frozenset[str], resolved: tuple[str, ...]) -> tuple[str, ...]: + """The members of this group a removal value names. + + Both ways of naming one count together. The id as written counts when the roster + holds it verbatim, which is how an earlier release recorded a member it could not + match, and the accounts it resolves to count when they are on the roster. Counting + only the resolved ones would let a value that is one member's canonical id and + another member's email revoke both, since each looks singular on its own. + """ + return tuple( + dict.fromkeys( + chain( + (value,) if value in roster else (), + (user_id for user_id in resolved if user_id in roster), + ) + ) + ) + + +async def _member_ids_to_drop( + members: Sequence[SCIMMember], roster: frozenset[str], prisma_client: PrismaClient +) -> frozenset[str]: + """The members a ``remove`` clears, one per id the request names. + + The roster holds canonical user ids, so a directory that added someone by their + email or SSO identity has to be able to remove them by that same value, and a + member an earlier release recorded under the raw id has to stay removable by it. + + Ambiguity is a property of the table as it stands, not of the value, so a value + that named one person when they were admitted can name two later. Resolving a + removal against the whole table would then drop nobody while answering 200, and + the person the directory just took out of the group would keep the team. So a + removal keeps only the accounts already on the roster: one is unambiguous however + many strangers share the address, none means there is nothing to revoke, and only + a value naming two of this group's own members is genuinely undecidable. That last + case fails rather than reporting a removal it did not perform, or revoking both. + + Raises: + HTTPException: 400 when a member id names more than one current member. + """ + written: Final = frozenset(_member_value(member) for member in members) + matched: Final = tuple( + [ + ( + value, + _roster_entries_named_by( + value, roster, await _users_named_by_member_value(value, prisma_client, take=None) + ), + ) + for value in sorted(written) + ] + ) + undecidable: Final = tuple(value for value, entries in matched if len(entries) > 1) + if undecidable: + raise HTTPException( + status_code=400, + detail={ + "error": f"Member ID '{undecidable[0]}' names more than one member of this group, so the removal " + "cannot be attributed. Send the LiteLLM user ID as the member value, or resolve the duplicate." + }, + ) + return frozenset(chain.from_iterable(entries for _, entries in matched)) + + async def _resolve_group_member_ids( members: Sequence[SCIMMember], created_via: str, @@ -670,17 +822,18 @@ async def _resolve_group_member_ids( """ Resolve SCIM group members to LiteLLM user ids, dropping members that are not users. - Only the operations that put ids onto a roster resolve their members: an id - that resolves to nothing is created when litellm_settings.scim_upsert_user is - True (default) and rejected per SCIM 2.0 otherwise. Removals do not come - through here; dropping an id is idempotent, so it needs neither a lookup nor a - user to drop. + Member ids are matched by ``user_id`` first, then by SSO identity or email. An + id that resolves to nothing is created when litellm_settings.scim_upsert_user is + True (default) and rejected per SCIM 2.0 otherwise. Removals do not come through + here: they resolve through ``_member_ids_to_drop`` instead, which neither creates + a user nor fails on an id it cannot place. Raises: - HTTPException: 400 when a member id is empty, or when scim_upsert_user is - False and a member id is neither an existing user, an existing team, nor a - member declared to be something other than a user. 500 when a member's - user row can neither be created nor found. + HTTPException: 400 when a member id is empty, when a member id names more + than one user, or when scim_upsert_user is False and a member id is neither + an existing user, an existing team, nor a member declared to be something + other than a user. 500 when a member's user row can neither be created nor + found. """ classified: Final = tuple([await _classify_group_member(member, prisma_client) for member in members]) partition: Final = _partition_classified_members(classified) @@ -692,6 +845,16 @@ async def _resolve_group_member_ids( skipped.reason, ) + if partition.ambiguous_values: + raise HTTPException( + status_code=400, + detail={ + "error": f"Member ID '{partition.ambiguous_values[0]}' names more than one LiteLLM user, so the " + "group membership cannot be attributed. Resolve the duplicate, which for an id that also matches a " + "SCIM-provisioned placeholder means deleting that placeholder." + }, + ) + if partition.unknown_ids and not await _get_scim_upsert_user_setting(): raise HTTPException( status_code=400, @@ -702,6 +865,13 @@ async def _resolve_group_member_ids( ) unique_unknown_ids: Final = tuple(dict.fromkeys(partition.unknown_ids)) + for user_id in unique_unknown_ids: + verbose_proxy_logger.warning( + "SCIM: creating placeholder user for group member '%s'; matched no user by user_id, sso_user_id or " + "user_email. An SSO-provisioned user's real account stays teamless if this is a mismatch", + user_id, + ) + creations: Final = tuple( [ ( @@ -2428,7 +2598,9 @@ async def _process_group_patch_operations( ) if op_type == "remove": - final_members = final_members - {_member_value(member) for member in patched_members} + final_members = final_members - await _member_ids_to_drop( + patched_members, frozenset(final_members), prisma_client + ) else: member_result = await _resolve_group_member_ids( members=patched_members, diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 82e22bb5bbf..a8e545a8551 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -2640,51 +2640,63 @@ async def _process_team_members( return updated_users, updated_team_memberships +def _resolve_member_identity(member: Member, updated_users: Sequence[LiteLLM_UserTable]) -> Member: + """Return ``member`` with whichever of ``user_id`` / ``user_email`` the caller left out filled in. + + The roster entry is a snapshot, so whatever is missing here is missing for good. + Resolution runs both ways off the user rows the add just touched: added by email + -> stamp the user_id, added by user_id -> stamp the email. A value the caller + supplied is never overwritten. + """ + resolved_user_id: Final = member.user_id or next( + ( + user.user_id + for user in updated_users + if member.user_email is not None and user.user_email == member.user_email + ), + None, + ) + resolved_user_email: Final = member.user_email or next( + ( + user.user_email + for user in updated_users + if resolved_user_id is not None and user.user_id == resolved_user_id and user.user_email is not None + ), + None, + ) + return member.model_copy( + update={ # mutable-ok: pydantic update payload + "user_id": resolved_user_id, + "user_email": resolved_user_email, + } + ) + + +def _member_already_in_team(member: Member, complete_team_data: LiteLLM_TeamTable) -> bool: + return any( + (member.user_id is not None and existing_member.user_id == member.user_id) + or (member.user_email is not None and existing_member.user_email == member.user_email) + for existing_member in complete_team_data.members_with_roles + ) + + async def _update_team_members_list( data: TeamMemberAddRequest, complete_team_data: LiteLLM_TeamTable, updated_users: list[LiteLLM_UserTable], ) -> None: """Update the team's members_with_roles list.""" - if isinstance(data.member, Member): - new_member: Final = data.member.model_copy() + requested_members: Final[Sequence[Member]] = ( + (data.member,) if isinstance(data.member, Member) else tuple(data.member) + ) + resolved_members: Final = tuple(_resolve_member_identity(m, updated_users) for m in requested_members) - # get user id - if new_member.user_id is None and new_member.user_email is not None: - for user in updated_users: - if user.user_email is not None and user.user_email == new_member.user_email: - new_member.user_id = user.user_id - - # Check if member already exists in team before adding - member_already_exists = False - for existing_member in complete_team_data.members_with_roles: - if (new_member.user_id is not None and existing_member.user_id == new_member.user_id) or ( - new_member.user_email is not None and existing_member.user_email == new_member.user_email - ): - member_already_exists = True - break - - if not member_already_exists: - complete_team_data.members_with_roles.append(new_member) - - elif isinstance(data.member, list): - for nm in data.member: - if nm.user_id is None and nm.user_email is not None: - for user in updated_users: - if user.user_email is not None and user.user_email == nm.user_email: - nm.user_id = user.user_id - - # Check if member already exists in team before adding - member_already_exists = False - for existing_member in complete_team_data.members_with_roles: - if (nm.user_id is not None and existing_member.user_id == nm.user_id) or ( - nm.user_email is not None and existing_member.user_email == nm.user_email - ): - member_already_exists = True - break - - if not member_already_exists: - complete_team_data.members_with_roles.append(nm) + # extend() consumes the generator as it appends, so a member already added by this + # same call is seen by the next _member_already_in_team check - the batch dedupes + # against itself exactly as the append-one-at-a-time loop this replaced did. + complete_team_data.members_with_roles.extend( # rebind-ok: this helper's contract is to grow the caller's roster in place + m for m in resolved_members if not _member_already_in_team(m, complete_team_data) + ) async def _add_team_members_to_team( @@ -4086,6 +4098,39 @@ async def _add_team_member_budget_table( return team_info_response_object +async def _hydrate_member_emails( + prisma_client: PrismaClient, + members: Sequence[Member], +) -> tuple[Member, ...]: + """Fill in ``user_email`` for roster entries that were stored without one. + + ``members_with_roles`` is a denormalized snapshot written at add-time, so an entry + stored with ``user_email=None`` keeps that null even once the user row has an email. + Look the missing ones up in ``LiteLLM_UserTable`` (one indexed query) and fill them + in. A stored email is never overwritten - the snapshot stays the source of truth + wherever it has a value. + """ + missing_user_ids: Final = frozenset(m.user_id for m in members if not m.user_email and m.user_id is not None) + if not missing_user_ids: + return tuple(members) + + user_rows: Final[Sequence[LiteLLM_UserTable]] = await _user_db(prisma_client).find_many( + where={ # mutable-ok: Prisma query filters are dict-shaped + "user_id": { # mutable-ok: Prisma query filters are dict-shaped + "in": sorted(missing_user_ids) + } + } + ) + email_by_user_id: Final = MappingProxyType({u.user_id: u.user_email for u in user_rows if u.user_email}) + + return tuple( + m.model_copy(update={"user_email": email_by_user_id[m.user_id]}) # mutable-ok: pydantic update payload + if not m.user_email and m.user_id in email_by_user_id + else m + for m in members + ) + + async def _resolve_team_access_group_resources( _team_info: TeamInfoResponseObjectTeamTable, ) -> TeamInfoResponseObjectTeamTable: @@ -4221,9 +4266,22 @@ async def team_info( # Resolve resources inherited from access groups resolved_team_info: Final = await _resolve_team_access_group_resources(_team_info) + # Fill in emails the add-time roster snapshot never captured + hydrated_members: Final = await _hydrate_member_emails( + prisma_client=prisma_client, + members=resolved_team_info.members_with_roles, + ) + hydrated_team_info: Final = resolved_team_info.model_copy( + update={ # mutable-ok: pydantic update payload + # list(), not the tuple: model_copy skips validation, so the field has + # to be handed the list[Member] the response model declares. + "members_with_roles": list(hydrated_members) # mutable-ok: declared list[Member] + } + ) + response_object: Final = TeamInfoResponseObject( team_id=team_id, - team_info=resolved_team_info, + team_info=hydrated_team_info, keys=keys, team_memberships=returned_tm, ) diff --git a/litellm/proxy/pass_through_endpoints/streaming_handler.py b/litellm/proxy/pass_through_endpoints/streaming_handler.py index 697eb7b96eb..b71622fc33d 100644 --- a/litellm/proxy/pass_through_endpoints/streaming_handler.py +++ b/litellm/proxy/pass_through_endpoints/streaming_handler.py @@ -106,7 +106,7 @@ class PassThroughStreamingHandler: ) # rebind-ok: SSE frame reassembly buffer across transport chunks if complete_frames: yield ProxyBaseLLMRequestProcessing._process_chunk_with_cost_injection( - complete_frames, resolved_model_name + complete_frames, resolved_model_name, litellm_logging_obj ) if pending: yield pending diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 4b97cade7f0..9ee62f94647 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -6840,6 +6840,20 @@ class ProxyConfig: if self._should_load_db_object(object_type="config_overrides"): await self._init_hashicorp_vault_config_override(prisma_client=prisma_client) + await self._apply_safe_litellm_settings_overrides_from_db(prisma_client=prisma_client) + + async def _apply_safe_litellm_settings_overrides_from_db(self, prisma_client: PrismaClient) -> None: + config_record: Final = await get_config_param(prisma_client, "litellm_settings") + if config_record is None or config_record.param_value is None: + return + raw_settings: Final = config_record.param_value + litellm_settings: Final = json.loads(raw_settings) if isinstance(raw_settings, str) else raw_settings + if not isinstance(litellm_settings, dict): + return + for key, value in litellm_settings.items(): + if key in LITELLM_SETTINGS_SAFE_DB_OVERRIDES: + setattr(litellm, key, value) + async def _init_semantic_filter_settings_in_db(self, prisma_client: PrismaClient): """ Initialize MCP semantic filter settings from database. @@ -8875,6 +8889,7 @@ class ProxyStartupEvent: proxy_logging_obj=proxy_logging_obj, prisma_client=prisma_client, reset_settings=get_budget_reset_settings(), + pod_lock_manager=proxy_logging_obj.db_spend_update_writer.pod_lock_manager, ) scheduler.add_job( diff --git a/litellm/proxy/public_endpoints/provider_create_fields.json b/litellm/proxy/public_endpoints/provider_create_fields.json index ab13773614a..9c2e94dd861 100644 --- a/litellm/proxy/public_endpoints/provider_create_fields.json +++ b/litellm/proxy/public_endpoints/provider_create_fields.json @@ -772,6 +772,34 @@ ], "default_model_placeholder": "gpt-3.5-turbo" }, + { + "provider": "Cognition", + "provider_display_name": "Cognition", + "litellm_provider": "cognition", + "credential_fields": [ + { + "key": "api_base", + "label": "API Base", + "placeholder": "https://api.cognition.ai/v1", + "tooltip": null, + "required": false, + "field_type": "text", + "options": null, + "default_value": null + }, + { + "key": "api_key", + "label": "API Key", + "placeholder": null, + "tooltip": null, + "required": true, + "field_type": "password", + "options": null, + "default_value": null + } + ], + "default_model_placeholder": "cognition/swe-1.7" + }, { "provider": "Cohere", "provider_display_name": "Cohere", diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index 7fd72d80e60..140d92cde6b 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -10,6 +10,7 @@ from pydantic import BaseModel import litellm from litellm._logging import verbose_proxy_logger from litellm.constants import ( + LITELLM_PROXY_MASTER_KEY_ALIAS, LITELLM_TRUNCATED_PAYLOAD_FIELD, LITELLM_TRUNCATION_DB_SAFEGUARD_NOTE, REDACTED_BY_LITELM_STRING, @@ -22,6 +23,7 @@ from litellm.litellm_core_utils.core_helpers import ( reconstruct_model_name, ) from litellm.litellm_core_utils.internal_call_metadata import is_unbilled_non_inference_call +from litellm.litellm_core_utils.litellm_logging import is_valid_sha256_hash from litellm.litellm_core_utils.safe_json_dumps import safe_dumps, strip_null_bytes from litellm.proxy._types import SpendLogsMetadata, SpendLogsPayload from litellm.proxy.spend_tracking.spend_log_error_logger import spend_log_error @@ -54,13 +56,6 @@ def _get_max_string_length_prompt_in_db() -> int: return DEFAULT_MAX_STRING_LENGTH_PROMPT_IN_DB -def _hash_api_key_for_spend_log(api_key: str) -> str: - stripped: Final = api_key[7:] if api_key[:7].lower() == "bearer " else api_key - if stripped.startswith("sk-"): - return hash_token(stripped) - return stripped - - def _is_master_key(api_key: str | None, _master_key: str | None) -> bool: """ Raw-only constant-time master-key comparison. The hashed form is never @@ -71,6 +66,28 @@ def _is_master_key(api_key: str | None, _master_key: str | None) -> bool: return secrets.compare_digest(api_key, _master_key) +_HASHED_JWT_RE = re.compile(r"hashed-jwt-[a-fA-F0-9]{64}") + + +def _is_non_secret_key_value(value: str) -> bool: + return ( + value == LITELLM_PROXY_MASTER_KEY_ALIAS + or is_valid_sha256_hash(value) + or _HASHED_JWT_RE.fullmatch(value) is not None + ) + + +def _redact_logged_api_key(value: str | None, *, already_redacted: bool = False) -> str | None: + if not isinstance(value, str) or not value: + return None + stripped: Final = re.sub(r"(?i)^bearer ", "", value) + if not stripped: + return None + if already_redacted and _is_non_secret_key_value(stripped): + return stripped + return hash_token(stripped) + + def _get_spend_logs_metadata( metadata: dict | None, applied_guardrails: list[str] | None = None, @@ -124,9 +141,12 @@ def _get_spend_logs_metadata( # Filter the metadata dictionary to include only the specified keys clean_metadata: Final = SpendLogsMetadata(**{key: metadata.get(key) for key in SpendLogsMetadata.__annotations__}) - raw_user_api_key: Final = clean_metadata.get("user_api_key") - if raw_user_api_key is not None and isinstance(raw_user_api_key, str): - clean_metadata["user_api_key"] = _hash_api_key_for_spend_log(raw_user_api_key) + _raw_key: Final = clean_metadata.get("user_api_key") + _trusted_hash: Final = metadata.get("user_api_key_hash") + _already_redacted: Final = ( + isinstance(_trusted_hash, str) and _is_non_secret_key_value(_trusted_hash) and _trusted_hash == _raw_key + ) + clean_metadata["user_api_key"] = _redact_logged_api_key(_raw_key, already_redacted=_already_redacted) clean_metadata["applied_guardrails"] = applied_guardrails clean_metadata["batch_models"] = batch_models clean_metadata["mcp_tool_call_metadata"] = mcp_tool_call_metadata @@ -282,16 +302,23 @@ def get_logging_payload(kwargs, response_obj, start_time, end_time) -> SpendLogs standard_logging_prompt_tokens = standard_logging_payload.get("prompt_tokens", 0) standard_logging_completion_tokens = standard_logging_payload.get("completion_tokens", 0) standard_logging_total_tokens = standard_logging_payload.get("total_tokens", 0) - if api_key is not None and isinstance(api_key, str): - api_key = _hash_api_key_for_spend_log(api_key) + _trusted_hash = metadata.get("user_api_key_hash") + _key_already_redacted = ( + isinstance(_trusted_hash, str) and _is_non_secret_key_value(_trusted_hash) and _trusted_hash == api_key + ) + api_key = _redact_logged_api_key(api_key, already_redacted=_key_already_redacted) or "" if ( standard_logging_payload is not None ): # [TODO] migrate completely to sl payload. currently missing pass-through endpoint data - api_key = api_key or standard_logging_payload["metadata"].get("user_api_key_hash") or "" + api_key = ( + api_key + or _redact_logged_api_key( + standard_logging_payload["metadata"].get("user_api_key_hash"), already_redacted=True + ) + or "" + ) end_user_id = end_user_id or standard_logging_payload["metadata"].get("user_api_key_end_user_id") - # BUG FIX: Don't overwrite api_key when standard_logging_payload is None - # The api_key was already extracted from metadata (line 243) and hashed (lines 256-259) request_tags = safe_dumps(metadata.get("tags", [])) if isinstance(metadata.get("tags", []), list) else "[]" if ( standard_logging_payload is not None and standard_logging_payload.get("request_tags") is not None diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 77a17f2143e..b077b33e217 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -11,6 +11,7 @@ from typing import ( get_args, ) +import httpx from openai._models import BaseModel as OpenAIObject from openai.types.audio.transcription_create_params import ( FileTypes as FileTypes, @@ -49,7 +50,7 @@ from litellm.types.llms.base import ( ) from litellm.types.mcp import MCPServerCostInfo -from ..litellm_core_utils.core_helpers import map_finish_reason +from ..litellm_core_utils.core_helpers import map_finish_reason, process_response_headers from .agents import LiteLLMSendMessageResponse from .guardrails import GuardrailEventHooks from .llms.anthropic_messages.anthropic_response import AnthropicMessagesResponse @@ -1916,6 +1917,10 @@ class ModelResponseBase(OpenAIObject): _response_headers: dict | None = None + def set_provider_response_headers(self, headers: httpx.Headers) -> None: + """Surface a provider's raw response headers to the caller as `llm_provider-*` headers.""" + self._hidden_params["additional_headers"] = process_response_headers(headers) + def model_dump(self, **kwargs): """Default to exclude_unset to avoid Pydantic serializer warnings for OpenAIObject-derived types.""" if "exclude_unset" not in kwargs and "exclude_none" not in kwargs: @@ -3788,6 +3793,7 @@ class LlmProviders(str, Enum): TENSORMESH = "tensormesh" LIBERTAI = "libertai" PINSTRIPES = "pinstripes" + COGNITION = "cognition" DARKBLOOM = "darkbloom" META = "meta" LITELLM_AGENT = "litellm_agent" diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 5a1c988c21a..91c10d13e8e 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -759,7 +759,9 @@ "supports_vision": true, "supports_native_structured_output": true, "supports_parallel_tool_use_config": true, - "prompt_cache_min_tokens": 4096 + "prompt_cache_min_tokens": 4096, + "input_cost_per_token_batches": 5e-07, + "output_cost_per_token_batches": 2.5e-06 }, "anthropic.claude-haiku-4-5@20251001": { "cache_creation_input_token_cost": 1.25e-06, @@ -2487,7 +2489,9 @@ "supports_vision": true, "supports_native_structured_output": true, "supports_parallel_tool_use_config": true, - "prompt_cache_min_tokens": 1024 + "prompt_cache_min_tokens": 1024, + "input_cost_per_token_batches": 1.5e-06, + "output_cost_per_token_batches": 7.5e-06 }, "anthropic.claude-v1": { "input_cost_per_token": 8e-06, @@ -2743,7 +2747,9 @@ "supports_vision": true, "supports_native_structured_output": true, "supports_parallel_tool_use_config": true, - "prompt_cache_min_tokens": 4096 + "prompt_cache_min_tokens": 4096, + "input_cost_per_token_batches": 5.5e-07, + "output_cost_per_token_batches": 2.75e-06 }, "apac.anthropic.claude-3-sonnet-20240229-v1:0": { "deprecation_date": "2026-07-30", @@ -2839,7 +2845,9 @@ "supports_vision": true, "supports_native_structured_output": true, "supports_parallel_tool_use_config": true, - "prompt_cache_min_tokens": 1024 + "prompt_cache_min_tokens": 1024, + "input_cost_per_token_batches": 1.65e-06, + "output_cost_per_token_batches": 8.25e-06 }, "azure/ada": { "input_cost_per_token": 1e-07, @@ -12401,8 +12409,8 @@ "input_cost_per_token": 3e-06, "litellm_provider": "anthropic", "max_input_tokens": 1000000, - "max_output_tokens": 64000, - "max_tokens": 64000, + "max_output_tokens": 128000, + "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 1.5e-05, "search_context_cost_per_query": { @@ -12452,7 +12460,9 @@ "supports_tool_choice": true, "supports_vision": true, "supports_parallel_tool_use_config": true, - "prompt_cache_min_tokens": 1024 + "prompt_cache_min_tokens": 1024, + "input_cost_per_token_batches": 1.5e-06, + "output_cost_per_token_batches": 7.5e-06 }, "claude-opus-4-1": { "cache_creation_input_token_cost": 1.875e-05, @@ -12787,7 +12797,8 @@ "us": 1.1 }, "supports_output_config": true, - "prompt_cache_min_tokens": 512 + "prompt_cache_min_tokens": 512, + "supports_native_structured_output": true }, "claude-opus-5": { "deprecation_date": "2027-07-24", @@ -13350,7 +13361,8 @@ "max_output_tokens": 4096, "max_tokens": 4096, "mode": "completion", - "output_cost_per_token": 2e-06 + "output_cost_per_token": 2e-06, + "deprecation_date": "2025-09-15" }, "command-a-03-2025": { "input_cost_per_token": 2.5e-06, @@ -13371,7 +13383,8 @@ "max_tokens": 4096, "mode": "chat", "output_cost_per_token": 6e-07, - "supports_tool_choice": true + "supports_tool_choice": true, + "deprecation_date": "2025-09-15" }, "command-nightly": { "input_cost_per_token": 1e-06, @@ -13391,7 +13404,8 @@ "mode": "chat", "output_cost_per_token": 6e-07, "supports_function_calling": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "deprecation_date": "2025-09-15" }, "command-r-08-2024": { "input_cost_per_token": 1.5e-07, @@ -13413,7 +13427,8 @@ "mode": "chat", "output_cost_per_token": 1e-05, "supports_function_calling": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "deprecation_date": "2025-09-15" }, "command-r-plus-08-2024": { "input_cost_per_token": 2.5e-06, @@ -17027,7 +17042,9 @@ "supports_vision": true, "supports_native_structured_output": true, "supports_parallel_tool_use_config": true, - "prompt_cache_min_tokens": 4096 + "prompt_cache_min_tokens": 4096, + "input_cost_per_token_batches": 5.5e-07, + "output_cost_per_token_batches": 2.75e-06 }, "eu.anthropic.claude-3-5-sonnet-20240620-v1:0": { "input_cost_per_token": 3e-06, @@ -17250,7 +17267,9 @@ "supports_vision": true, "supports_native_structured_output": true, "supports_parallel_tool_use_config": true, - "prompt_cache_min_tokens": 1024 + "prompt_cache_min_tokens": 1024, + "input_cost_per_token_batches": 1.65e-06, + "output_cost_per_token_batches": 8.25e-06 }, "eu.meta.llama3-2-1b-instruct-v1:0": { "input_cost_per_token": 1.3e-07, @@ -17400,7 +17419,7 @@ "fal_ai/openai/gpt-image-2": { "litellm_provider": "fal_ai", "metadata": { - "notes": "OpenAI gpt-image-2 served through fal.ai. fal bills by token, so the flat output_cost_per_image here is fal's published per-image rate for a default request (quality=high, image_size=landscape_4_3 at 1024x768). Other canonical sizes at high quality: 1024x1024 $0.211, 1024x1536 $0.165, 1920x1080 $0.158, 2560x1440 $0.222, 3840x2160 $0.401" + "notes": "OpenAI gpt-image-2 served through fal.ai. fal bills by token but publishes deterministic per-image prices per size and quality, mirrored here as keyed entries fal_ai/{quality}/{width}-x-{height}/openai/gpt-image-2 that litellm's fal_ai cost calculator picks from the request params. This flat entry is the fallback when no keyed entry matches and carries the default request rate (quality=high, image_size=landscape_4_3 at 1024x768). quality=auto is priced as high" }, "mode": "image_generation", "output_cost_per_image": 0.145, @@ -17410,10 +17429,190 @@ ], "supports_vision": true }, + "fal_ai/low/1024-x-768/openai/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.005, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/low/1024-x-1024/openai/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.006, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/low/1024-x-1536/openai/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.005, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/low/1920-x-1080/openai/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.005, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/low/2560-x-1440/openai/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.007, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/low/3840-x-2160/openai/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.012, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/medium/1024-x-768/openai/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.037, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/medium/1024-x-1024/openai/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.053, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/medium/1024-x-1536/openai/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.042, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/medium/1920-x-1080/openai/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.04, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/medium/2560-x-1440/openai/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.056, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/medium/3840-x-2160/openai/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.101, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/high/1024-x-768/openai/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.145, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/high/1024-x-1024/openai/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.211, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/high/1024-x-1536/openai/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.165, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/high/1920-x-1080/openai/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.158, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/high/2560-x-1440/openai/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.222, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/high/3840-x-2160/openai/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.401, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, "fal_ai/gpt-image-2": { "litellm_provider": "fal_ai", "metadata": { - "notes": "Alias of fal_ai/openai/gpt-image-2, which litellm also accepts without the openai/ prefix. Same rate, see that entry for the size and quality caveat" + "notes": "Alias of fal_ai/openai/gpt-image-2, which litellm also accepts without the openai/ prefix. Same rates, including the keyed fal_ai/{quality}/{width}-x-{height}/gpt-image-2 entries; see that entry for details" }, "mode": "image_generation", "output_cost_per_image": 0.145, @@ -17423,13 +17622,373 @@ ], "supports_vision": true }, + "fal_ai/low/1024-x-768/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.005, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/low/1024-x-1024/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.006, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/low/1024-x-1536/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.005, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/low/1920-x-1080/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.005, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/low/2560-x-1440/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.007, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/low/3840-x-2160/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.012, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/medium/1024-x-768/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.037, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/medium/1024-x-1024/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.053, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/medium/1024-x-1536/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.042, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/medium/1920-x-1080/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.04, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/medium/2560-x-1440/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.056, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/medium/3840-x-2160/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.101, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/high/1024-x-768/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.145, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/high/1024-x-1024/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.211, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/high/1024-x-1536/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.165, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/high/1920-x-1080/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.158, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/high/2560-x-1440/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.222, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/high/3840-x-2160/gpt-image-2": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.401, + "source": "https://fal.ai/models/openai/gpt-image-2", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, "fal_ai/openai/gpt-image-2/edit": { "litellm_provider": "fal_ai", "metadata": { - "notes": "Editing endpoint of gpt-image-2 on fal.ai, reached through the image generation path with fal's image_urls param since /v1/images/edits is not wired for fal_ai. Same rate as fal_ai/openai/gpt-image-2, see that entry for the size and quality caveat" + "notes": "Editing endpoint of gpt-image-2 on fal.ai, reached through the image generation path with fal's image_urls param since /v1/images/edits is not wired for fal_ai. Prices include one input image and live in keyed entries fal_ai/{quality}/{width}-x-{height}/openai/gpt-image-2/edit. This flat entry is the fallback for the default edit request (quality=high, image_size=auto, inferred from the input image, priced as 1024x768 high)" }, "mode": "image_generation", - "output_cost_per_image": 0.145, + "output_cost_per_image": 0.151, + "source": "https://fal.ai/models/openai/gpt-image-2/edit", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/low/1024-x-768/openai/gpt-image-2/edit": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.011, + "source": "https://fal.ai/models/openai/gpt-image-2/edit", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/low/1024-x-1024/openai/gpt-image-2/edit": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.015, + "source": "https://fal.ai/models/openai/gpt-image-2/edit", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/low/1024-x-1536/openai/gpt-image-2/edit": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.018, + "source": "https://fal.ai/models/openai/gpt-image-2/edit", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/low/1920-x-1080/openai/gpt-image-2/edit": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.017, + "source": "https://fal.ai/models/openai/gpt-image-2/edit", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/low/2560-x-1440/openai/gpt-image-2/edit": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.019, + "source": "https://fal.ai/models/openai/gpt-image-2/edit", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/low/3840-x-2160/openai/gpt-image-2/edit": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.024, + "source": "https://fal.ai/models/openai/gpt-image-2/edit", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/medium/1024-x-768/openai/gpt-image-2/edit": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.043, + "source": "https://fal.ai/models/openai/gpt-image-2/edit", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/medium/1024-x-1024/openai/gpt-image-2/edit": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.061, + "source": "https://fal.ai/models/openai/gpt-image-2/edit", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/medium/1024-x-1536/openai/gpt-image-2/edit": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.054, + "source": "https://fal.ai/models/openai/gpt-image-2/edit", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/medium/1920-x-1080/openai/gpt-image-2/edit": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.053, + "source": "https://fal.ai/models/openai/gpt-image-2/edit", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/medium/2560-x-1440/openai/gpt-image-2/edit": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.068, + "source": "https://fal.ai/models/openai/gpt-image-2/edit", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/medium/3840-x-2160/openai/gpt-image-2/edit": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.113, + "source": "https://fal.ai/models/openai/gpt-image-2/edit", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/high/1024-x-768/openai/gpt-image-2/edit": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.151, + "source": "https://fal.ai/models/openai/gpt-image-2/edit", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/high/1024-x-1024/openai/gpt-image-2/edit": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.219, + "source": "https://fal.ai/models/openai/gpt-image-2/edit", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/high/1024-x-1536/openai/gpt-image-2/edit": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.178, + "source": "https://fal.ai/models/openai/gpt-image-2/edit", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/high/1920-x-1080/openai/gpt-image-2/edit": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.158, + "source": "https://fal.ai/models/openai/gpt-image-2/edit", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/high/2560-x-1440/openai/gpt-image-2/edit": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.234, + "source": "https://fal.ai/models/openai/gpt-image-2/edit", + "supported_endpoints": [ + "/v1/images/generations" + ], + "supports_vision": true + }, + "fal_ai/high/3840-x-2160/openai/gpt-image-2/edit": { + "litellm_provider": "fal_ai", + "mode": "image_generation", + "output_cost_per_image": 0.413, "source": "https://fal.ai/models/openai/gpt-image-2/edit", "supported_endpoints": [ "/v1/images/generations" @@ -18932,6 +19491,106 @@ }, "web_search_billing_unit": "per_query" }, + "gemini-3.1-flash-lite-image": { + "input_cost_per_image": 0.00028, + "input_cost_per_token": 2.5e-07, + "litellm_provider": "vertex_ai-language-models", + "max_input_tokens": 65536, + "max_output_tokens": 4096, + "max_tokens": 4096, + "mode": "image_generation", + "output_cost_per_image": 0.0336, + "output_cost_per_image_token": 3e-05, + "output_cost_per_token": 1.5e-06, + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#gemini-models", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text", + "image" + ], + "supports_function_calling": false, + "supports_prompt_caching": true, + "supports_response_schema": false, + "supports_reasoning": true, + "supports_system_messages": true, + "supports_vision": true + }, + "gemini/gemini-3.1-flash-lite-image": { + "rpm": 1000, + "tpm": 4000000, + "input_cost_per_image": 0.00028, + "input_cost_per_token": 2.5e-07, + "input_cost_per_token_batches": 1.25e-07, + "litellm_provider": "gemini", + "max_input_tokens": 65536, + "max_output_tokens": 4096, + "max_tokens": 4096, + "mode": "image_generation", + "output_cost_per_image": 0.0336, + "output_cost_per_image_token": 3e-05, + "output_cost_per_token": 1.5e-06, + "output_cost_per_token_batches": 7.5e-07, + "source": "https://ai.google.dev/gemini-api/docs/pricing#gemini-3.1-flash-lite-image", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text", + "image" + ], + "supports_function_calling": true, + "supports_prompt_caching": false, + "supports_response_schema": false, + "supports_reasoning": true, + "supports_system_messages": true, + "supports_vision": true + }, + "vertex_ai/gemini-3.1-flash-lite-image": { + "input_cost_per_image": 0.00028, + "input_cost_per_token": 2.5e-07, + "litellm_provider": "vertex_ai-language-models", + "max_input_tokens": 65536, + "max_output_tokens": 4096, + "max_tokens": 4096, + "mode": "image_generation", + "output_cost_per_image": 0.0336, + "output_cost_per_image_token": 3e-05, + "output_cost_per_token": 1.5e-06, + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#gemini-models", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text", + "image" + ], + "supports_function_calling": false, + "supports_prompt_caching": true, + "supports_response_schema": false, + "supports_reasoning": true, + "supports_system_messages": true, + "supports_vision": true + }, "gemini-3.1-flash-image": { "deprecation_date": "2027-05-28", "input_cost_per_image": 0.00056, @@ -19795,7 +20454,7 @@ "deprecation_date": "2027-05-19", "cache_read_input_token_cost": 1.5e-07, "input_cost_per_token": 1.5e-06, - "input_cost_per_audio_token": 1e-06, + "input_cost_per_audio_token": 1.5e-06, "litellm_provider": "vertex_ai", "max_input_tokens": 1048576, "max_output_tokens": 65535, @@ -19834,7 +20493,7 @@ "supports_web_search": true, "supports_native_streaming": true, "input_cost_per_token_priority": 2.7e-06, - "input_cost_per_audio_token_priority": 1.8e-06, + "input_cost_per_audio_token_priority": 2.7e-06, "output_cost_per_token_priority": 1.62e-05, "cache_read_input_token_cost_priority": 2.7e-07, "search_context_cost_per_query": { @@ -19842,7 +20501,12 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, - "web_search_billing_unit": "per_query" + "web_search_billing_unit": "per_query", + "input_cost_per_token_batches": 7.5e-07, + "output_cost_per_token_batches": 4.5e-06, + "input_cost_per_token_flex": 7.5e-07, + "output_cost_per_token_flex": 4.5e-06, + "cache_read_input_token_cost_flex": 7.5e-08 }, "vertex_ai/gemini-3.6-flash": { "prompt_cache_min_tokens": 4096, @@ -21527,7 +22191,7 @@ "gemini/gemini-3.5-flash": { "prompt_cache_min_tokens": 4096, "cache_read_input_token_cost": 1.5e-07, - "input_cost_per_audio_token": 1e-06, + "input_cost_per_audio_token": 1.5e-06, "input_cost_per_token": 1.5e-06, "litellm_provider": "gemini", "max_input_tokens": 1048576, @@ -21569,7 +22233,7 @@ "supports_native_streaming": true, "tpm": 800000, "input_cost_per_token_priority": 2.7e-06, - "input_cost_per_audio_token_priority": 1.8e-06, + "input_cost_per_audio_token_priority": 2.7e-06, "output_cost_per_token_priority": 1.62e-05, "cache_read_input_token_cost_priority": 2.7e-07, "search_context_cost_per_query": { @@ -21577,7 +22241,12 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, - "web_search_billing_unit": "per_query" + "web_search_billing_unit": "per_query", + "input_cost_per_token_batches": 7.5e-07, + "output_cost_per_token_batches": 4.5e-06, + "input_cost_per_token_flex": 7.5e-07, + "output_cost_per_token_flex": 4.5e-06, + "cache_read_input_token_cost_flex": 8e-08 }, "gemini/gemini-3.6-flash": { "prompt_cache_min_tokens": 4096, @@ -21929,7 +22598,7 @@ "prompt_cache_min_tokens": 4096, "deprecation_date": "2027-05-19", "cache_read_input_token_cost": 1.5e-07, - "input_cost_per_audio_token": 1e-06, + "input_cost_per_audio_token": 1.5e-06, "input_cost_per_token": 1.5e-06, "litellm_provider": "vertex_ai-language-models", "max_input_tokens": 1048576, @@ -21969,7 +22638,7 @@ "supports_web_search": true, "supports_native_streaming": true, "input_cost_per_token_priority": 2.7e-06, - "input_cost_per_audio_token_priority": 1.8e-06, + "input_cost_per_audio_token_priority": 2.7e-06, "output_cost_per_token_priority": 1.62e-05, "cache_read_input_token_cost_priority": 2.7e-07, "search_context_cost_per_query": { @@ -21977,7 +22646,12 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, - "web_search_billing_unit": "per_query" + "web_search_billing_unit": "per_query", + "input_cost_per_token_batches": 7.5e-07, + "output_cost_per_token_batches": 4.5e-06, + "input_cost_per_token_flex": 7.5e-07, + "output_cost_per_token_flex": 4.5e-06, + "cache_read_input_token_cost_flex": 7.5e-08 }, "gemini-3.6-flash": { "prompt_cache_min_tokens": 4096, @@ -23307,7 +23981,9 @@ "supports_vision": true, "supports_native_structured_output": true, "supports_parallel_tool_use_config": true, - "prompt_cache_min_tokens": 1024 + "prompt_cache_min_tokens": 1024, + "input_cost_per_token_batches": 1.5e-06, + "output_cost_per_token_batches": 7.5e-06 }, "global.anthropic.claude-sonnet-4-20250514-v1:0": { "cache_creation_input_token_cost": 3.75e-06, @@ -23365,7 +24041,9 @@ "supports_vision": true, "supports_native_structured_output": true, "supports_parallel_tool_use_config": true, - "prompt_cache_min_tokens": 4096 + "prompt_cache_min_tokens": 4096, + "input_cost_per_token_batches": 5e-07, + "output_cost_per_token_batches": 2.5e-06 }, "global.amazon.nova-2-lite-v1:0": { "cache_read_input_token_cost": 7.5e-08, @@ -24181,7 +24859,8 @@ "supports_response_schema": false, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": false + "supports_vision": false, + "deprecation_date": "2027-01-20" }, "gpt-4o-mini": { "cache_read_input_token_cost": 7.5e-08, @@ -25464,6 +26143,7 @@ "supported_output_modalities": [ "text" ], + "supports_computer_use": true, "supports_function_calling": true, "supports_minimal_reasoning_effort": false, "supports_native_streaming": true, @@ -25606,6 +26286,155 @@ "supports_web_search": true, "supports_xhigh_reasoning_effort": true }, + "gpt-5.6-cyber": { + "cache_creation_input_token_cost": 1.5625e-05, + "cache_creation_input_token_cost_above_272k_tokens": 3.125e-05, + "cache_read_input_token_cost": 1.25e-06, + "cache_read_input_token_cost_above_272k_tokens": 2.5e-06, + "input_cost_per_token": 1.25e-05, + "input_cost_per_token_above_272k_tokens": 2.5e-05, + "litellm_provider": "openai", + "max_input_tokens": 400000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 7.5e-05, + "output_cost_per_token_above_272k_tokens": 0.0001125, + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "source": "https://platform.openai.com/docs/models/gpt-5.6-cyber", + "supports_computer_use": true, + "supports_parallel_function_calling": true + }, + "daybreak-red-latest": { + "cache_creation_input_token_cost": 1.5625e-05, + "cache_creation_input_token_cost_above_272k_tokens": 3.125e-05, + "cache_read_input_token_cost": 1.25e-06, + "cache_read_input_token_cost_above_272k_tokens": 2.5e-06, + "input_cost_per_token": 1.25e-05, + "input_cost_per_token_above_272k_tokens": 2.5e-05, + "litellm_provider": "openai", + "max_input_tokens": 400000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 7.5e-05, + "output_cost_per_token_above_272k_tokens": 0.0001125, + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "source": "https://platform.openai.com/docs/models/daybreak-red-latest", + "supports_computer_use": true, + "supports_parallel_function_calling": true + }, + "daybreak-blue-latest": { + "cache_creation_input_token_cost": 6.25e-06, + "cache_creation_input_token_cost_above_272k_tokens": 1.25e-05, + "cache_read_input_token_cost": 5e-07, + "cache_read_input_token_cost_above_272k_tokens": 1e-06, + "input_cost_per_token": 5e-06, + "input_cost_per_token_above_272k_tokens": 1e-05, + "litellm_provider": "openai", + "max_input_tokens": 1050000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 3e-05, + "output_cost_per_token_above_272k_tokens": 4.5e-05, + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_computer_use": true, + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "source": "https://platform.openai.com/docs/models/daybreak-blue-latest", + "supports_parallel_function_calling": true + }, + "chat-latest": { + "cache_read_input_token_cost": 5e-07, + "input_cost_per_token": 5e-06, + "litellm_provider": "openai", + "max_input_tokens": 400000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 3e-05, + "source": "https://platform.openai.com/docs/models/chat-latest", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true + }, "gpt-5.5": { "cache_read_input_token_cost": 5e-07, "cache_read_input_token_cost_above_272k_tokens": 1e-06, @@ -28120,7 +28949,9 @@ "supports_vision": true, "supports_native_structured_output": true, "supports_parallel_tool_use_config": true, - "prompt_cache_min_tokens": 1024 + "prompt_cache_min_tokens": 1024, + "input_cost_per_token_batches": 1.65e-06, + "output_cost_per_token_batches": 8.25e-06 }, "jp.anthropic.claude-haiku-4-5-20251001-v1:0": { "cache_creation_input_token_cost": 1.375e-06, @@ -28146,7 +28977,9 @@ "supports_vision": true, "supports_native_structured_output": true, "supports_parallel_tool_use_config": true, - "prompt_cache_min_tokens": 4096 + "prompt_cache_min_tokens": 4096, + "input_cost_per_token_batches": 5.5e-07, + "output_cost_per_token_batches": 2.75e-06 }, "crusoe/deepseek-ai/DeepSeek-R1-0528": { "input_cost_per_token": 3e-06, @@ -29369,28 +30202,30 @@ "mistral/codestral-2508": { "input_cost_per_token": 3e-07, "litellm_provider": "mistral", - "max_input_tokens": 256000, - "max_output_tokens": 256000, - "max_tokens": 256000, + "max_input_tokens": 128000, + "max_output_tokens": 128000, + "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 9e-07, - "source": "https://mistral.ai/news/codestral-25-08", + "source": "https://docs.mistral.ai/models/model-cards/codestral-25-08", "supports_assistant_prefill": true, "supports_function_calling": true, "supports_response_schema": true, "supports_tool_choice": true }, "mistral/codestral-latest": { - "input_cost_per_token": 1e-06, + "input_cost_per_token": 3e-07, "litellm_provider": "mistral", - "max_input_tokens": 32000, - "max_output_tokens": 8191, - "max_tokens": 8191, + "max_input_tokens": 128000, + "max_output_tokens": 128000, + "max_tokens": 128000, "mode": "chat", - "output_cost_per_token": 3e-06, + "output_cost_per_token": 9e-07, "supports_assistant_prefill": true, "supports_response_schema": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "source": "https://docs.mistral.ai/models/model-cards/codestral-25-08", + "supports_function_calling": true }, "mistral/codestral-mamba-latest": { "input_cost_per_token": 2.5e-07, @@ -29623,6 +30458,16 @@ ], "source": "https://mistral.ai/pricing#api-pricing" }, + "mistral/mistral-ocr-4-1": { + "annotation_cost_per_page": 0.005, + "litellm_provider": "mistral", + "mode": "ocr", + "ocr_cost_per_page": 0.004, + "source": "https://docs.mistral.ai/models/model-cards/ocr-4-1", + "supported_endpoints": [ + "/v1/ocr" + ] + }, "mistral/mistral-ocr-2505-completion": { "deprecation_date": "2026-05-31", "litellm_provider": "mistral", @@ -29947,18 +30792,19 @@ "supports_tool_choice": true }, "mistral/mistral-small-latest": { - "input_cost_per_token": 6e-08, + "input_cost_per_token": 1.5e-07, "litellm_provider": "mistral", - "max_input_tokens": 131072, - "max_output_tokens": 131072, - "max_tokens": 131072, + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, "mode": "chat", - "output_cost_per_token": 1.8e-07, - "source": "https://mistral.ai/pricing", + "output_cost_per_token": 6e-07, + "source": "https://docs.mistral.ai/models/model-cards/mistral-small-4-0-26-03", "supports_assistant_prefill": true, "supports_function_calling": true, "supports_response_schema": true, "supports_tool_choice": true, + "supports_reasoning": true, "supports_vision": true }, "mistral/mistral-small-3-2-2506": { @@ -32781,6 +33627,31 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true }, + "openrouter/anthropic/claude-opus-5": { + "prompt_cache_min_tokens": 512, + "supports_adaptive_thinking": true, + "cache_creation_input_token_cost": 6.25e-06, + "cache_read_input_token_cost": 5e-07, + "input_cost_per_token": 5e-06, + "litellm_provider": "openrouter", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 2.5e-05, + "source": "https://openrouter.ai/anthropic/claude-opus-5", + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_max_reasoning_effort": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true + }, "openrouter/bytedance/ui-tars-1.5-7b": { "input_cost_per_token": 1e-07, "litellm_provider": "openrouter", @@ -32889,6 +33760,38 @@ "supports_reasoning": true, "supports_tool_choice": true }, + "openrouter/deepseek/deepseek-v4-pro": { + "input_cost_per_token": 1.32e-06, + "input_cost_per_token_cache_hit": 4.4e-08, + "litellm_provider": "openrouter", + "max_input_tokens": 1048576, + "max_output_tokens": 384000, + "max_tokens": 384000, + "mode": "chat", + "output_cost_per_token": 3.96e-06, + "source": "https://openrouter.ai/deepseek/deepseek-v4-pro", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "openrouter/deepseek/deepseek-v4-pro-0813": { + "input_cost_per_token": 1.32e-06, + "input_cost_per_token_cache_hit": 4.4e-08, + "litellm_provider": "openrouter", + "max_input_tokens": 1048576, + "max_output_tokens": 384000, + "max_tokens": 384000, + "mode": "chat", + "output_cost_per_token": 3.96e-06, + "source": "https://openrouter.ai/deepseek/deepseek-v4-pro-0813", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, "openrouter/google/gemini-2.0-flash-001": { "deprecation_date": "2026-06-01", "input_cost_per_audio_token": 7e-07, @@ -34923,7 +35826,9 @@ "supports_function_calling": true, "supports_reasoning": true, "supports_tool_choice": true, - "supports_native_structured_output": true + "supports_native_structured_output": true, + "input_cost_per_token_batches": 1.1e-07, + "output_cost_per_token_batches": 4.4e-07 }, "qwen.qwen3-coder-30b-a3b-v1:0": { "input_cost_per_token": 1.5e-07, @@ -35458,7 +36363,8 @@ "max_output_tokens": 4096, "max_tokens": 4096, "mode": "rerank", - "output_cost_per_token": 0.0 + "output_cost_per_token": 0.0, + "deprecation_date": "2025-04-30" }, "rerank-english-v3.0": { "input_cost_per_query": 0.002, @@ -35478,7 +36384,8 @@ "max_output_tokens": 4096, "max_tokens": 4096, "mode": "rerank", - "output_cost_per_token": 0.0 + "output_cost_per_token": 0.0, + "deprecation_date": "2025-04-30" }, "rerank-multilingual-v3.0": { "input_cost_per_query": 0.002, @@ -37118,7 +38025,9 @@ "supports_vision": true, "supports_native_structured_output": true, "supports_parallel_tool_use_config": true, - "prompt_cache_min_tokens": 4096 + "prompt_cache_min_tokens": 4096, + "input_cost_per_token_batches": 5.5e-07, + "output_cost_per_token_batches": 2.75e-06 }, "us.anthropic.claude-3-5-sonnet-20240620-v1:0": { "input_cost_per_token": 3e-06, @@ -37284,7 +38193,9 @@ "supports_vision": true, "supports_native_structured_output": true, "supports_parallel_tool_use_config": true, - "prompt_cache_min_tokens": 1024 + "prompt_cache_min_tokens": 1024, + "input_cost_per_token_batches": 1.65e-06, + "output_cost_per_token_batches": 8.25e-06 }, "us-gov.anthropic.claude-sonnet-4-5-20250929-v1:0": { "cache_creation_input_token_cost": 4.5e-06, @@ -37339,7 +38250,9 @@ "supports_vision": true, "supports_native_structured_output": true, "supports_parallel_tool_use_config": true, - "prompt_cache_min_tokens": 4096 + "prompt_cache_min_tokens": 4096, + "input_cost_per_token_batches": 5.5e-07, + "output_cost_per_token_batches": 2.75e-06 }, "us.anthropic.claude-opus-4-20250514-v1:0": { "cache_creation_input_token_cost": 1.875e-05, @@ -39930,13 +40843,13 @@ "supports_tool_choice": true }, "vertex_ai/deepseek-ai/deepseek-v3.1-maas": { - "input_cost_per_token": 1.35e-06, + "input_cost_per_token": 6e-07, "litellm_provider": "vertex_ai-deepseek_models", "max_input_tokens": 163840, "max_output_tokens": 32768, "max_tokens": 32768, "mode": "chat", - "output_cost_per_token": 5.4e-06, + "output_cost_per_token": 1.7e-06, "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#partner-models", "supported_regions": [ "us-central1" @@ -40770,13 +41683,13 @@ "supports_vision": true }, "vertex_ai/openai/gpt-oss-120b-maas": { - "input_cost_per_token": 1.5e-07, + "input_cost_per_token": 9e-08, "litellm_provider": "vertex_ai-openai_models", "max_input_tokens": 131072, "max_output_tokens": 32768, "max_tokens": 32768, "mode": "chat", - "output_cost_per_token": 6e-07, + "output_cost_per_token": 3.6e-07, "source": "https://console.cloud.google.com/vertex-ai/publishers/openai/model-garden/gpt-oss-120b-maas", "supports_reasoning": true }, @@ -40858,13 +41771,13 @@ "supports_web_search": true }, "vertex_ai/qwen/qwen3-235b-a22b-instruct-2507-maas": { - "input_cost_per_token": 2.5e-07, + "input_cost_per_token": 2.2e-07, "litellm_provider": "vertex_ai-qwen_models", "max_input_tokens": 262144, "max_output_tokens": 16384, "max_tokens": 16384, "mode": "chat", - "output_cost_per_token": 1e-06, + "output_cost_per_token": 8.8e-07, "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", "supported_regions": [ "global", @@ -40874,13 +41787,13 @@ "supports_tool_choice": true }, "vertex_ai/qwen/qwen3-coder-480b-a35b-instruct-maas": { - "input_cost_per_token": 1e-06, + "input_cost_per_token": 2.2e-07, "litellm_provider": "vertex_ai-qwen_models", "max_input_tokens": 262144, "max_output_tokens": 32768, "max_tokens": 32768, "mode": "chat", - "output_cost_per_token": 4e-06, + "output_cost_per_token": 1.8e-06, "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", "supported_regions": [ "global" @@ -41855,7 +42768,8 @@ "supports_prompt_caching": true, "supports_response_schema": false, "supports_tool_choice": true, - "supports_web_search": true + "supports_web_search": true, + "deprecation_date": "2026-05-15" }, "xai/grok-3-mini": { "cache_read_input_token_cost": 7.5e-08, @@ -41973,7 +42887,8 @@ "supports_function_calling": true, "supports_prompt_caching": true, "supports_tool_choice": true, - "supports_web_search": true + "supports_web_search": true, + "deprecation_date": "2026-05-15" }, "xai/grok-4-fast-reasoning": { "cache_read_input_token_cost": 5e-08, @@ -42042,7 +42957,8 @@ "supports_function_calling": true, "supports_prompt_caching": true, "supports_tool_choice": true, - "supports_web_search": true + "supports_web_search": true, + "deprecation_date": "2026-05-15" }, "xai/grok-4-1-fast": { "cache_read_input_token_cost": 5e-08, @@ -42370,7 +43286,8 @@ "output_cost_per_token_above_200k_tokens": 4e-06, "cache_read_input_token_cost_above_200k_tokens": 4e-07, "supports_response_schema": true, - "supports_vision": true + "supports_vision": true, + "deprecation_date": "2026-05-15" }, "xai/grok-code-fast-1": { "cache_read_input_token_cost": 2e-07, @@ -42390,7 +43307,8 @@ "output_cost_per_token_above_200k_tokens": 4e-06, "cache_read_input_token_cost_above_200k_tokens": 4e-07, "supports_response_schema": true, - "supports_vision": true + "supports_vision": true, + "deprecation_date": "2026-05-15" }, "xai/grok-code-fast-1-0825": { "cache_read_input_token_cost": 2e-07, @@ -42410,7 +43328,8 @@ "output_cost_per_token_above_200k_tokens": 4e-06, "cache_read_input_token_cost_above_200k_tokens": 4e-07, "supports_response_schema": true, - "supports_vision": true + "supports_vision": true, + "deprecation_date": "2026-05-15" }, "xai/grok-vision-beta": { "input_cost_per_image": 5e-06, @@ -46760,7 +47679,8 @@ "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_system_messages": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "deprecation_date": "2027-01-20" }, "gpt-realtime-whisper": { "input_cost_per_second": 0.0002833333333333333, @@ -48558,6 +49478,36 @@ "supports_reasoning": true, "supports_vision": false }, + "cognition/swe-1.6": { + "input_cost_per_token": 5e-07, + "output_cost_per_token": 2.5e-06, + "cache_read_input_token_cost": 2e-07, + "litellm_provider": "cognition", + "mode": "chat", + "supports_function_calling": true, + "supports_prompt_caching": true, + "source": "https://docs.devin.ai/windsurf/plugins/cascade/models" + }, + "cognition/swe-1.7": { + "input_cost_per_token": 5e-07, + "output_cost_per_token": 2.5e-06, + "cache_read_input_token_cost": 2e-07, + "litellm_provider": "cognition", + "mode": "chat", + "supports_function_calling": true, + "supports_prompt_caching": true, + "source": "https://docs.devin.ai/desktop/models" + }, + "cognition/swe-1.7-lightning": { + "input_cost_per_token": 2.5e-06, + "output_cost_per_token": 1.25e-05, + "cache_read_input_token_cost": 1e-06, + "litellm_provider": "cognition", + "mode": "chat", + "supports_function_calling": true, + "supports_prompt_caching": true, + "source": "https://docs.devin.ai/desktop/models" + }, "pinstripes/ps/glm-4.5-air": { "max_tokens": 128000, "max_input_tokens": 128000, @@ -48817,7 +49767,8 @@ "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, - "supports_xhigh_reasoning_effort": true + "supports_xhigh_reasoning_effort": true, + "supports_native_structured_output": true }, "claude-mythos-preview": { "cache_creation_input_token_cost": 1.25e-05, @@ -48850,7 +49801,8 @@ "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, - "supports_xhigh_reasoning_effort": true + "supports_xhigh_reasoning_effort": true, + "supports_native_structured_output": true }, "gemini/gemini-robotics-er-2-streaming-preview": { "input_cost_per_audio_token": 2e-06, @@ -48896,7 +49848,8 @@ "supports_function_calling": true, "supports_reasoning": true, "supports_response_schema": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_vision": true }, "mistral/labs-leanstral-1-5": { "input_cost_per_token": 0.0, @@ -49023,5 +49976,424 @@ } } ] + }, + "gemini/gemini-3.5-live-translate-preview": { + "input_cost_per_audio_token": 3.5e-06, + "input_cost_per_token": 3.5e-06, + "litellm_provider": "gemini", + "mode": "chat", + "output_cost_per_audio_token": 2.1e-05, + "output_cost_per_token": 2.1e-05, + "rpm": 10, + "source": "https://ai.google.dev/gemini-api/docs/pricing", + "supported_endpoints": [ + "/v1/realtime" + ], + "supported_modalities": [ + "audio" + ], + "supported_output_modalities": [ + "audio" + ], + "supports_audio_input": true, + "supports_audio_output": true, + "tpm": 250000 + }, + "perplexity/pplx-embed-context-v1-0.6b": { + "input_cost_per_token": 8e-09, + "litellm_provider": "perplexity", + "max_input_tokens": 32768, + "max_tokens": 32768, + "mode": "embedding", + "output_cost_per_token": 0.0, + "output_vector_size": 1024, + "source": "https://docs.perplexity.ai/getting-started/pricing" + }, + "perplexity/pplx-embed-context-v1-4b": { + "input_cost_per_token": 5e-08, + "litellm_provider": "perplexity", + "max_input_tokens": 32768, + "max_tokens": 32768, + "mode": "embedding", + "output_cost_per_token": 0.0, + "output_vector_size": 2560, + "source": "https://docs.perplexity.ai/getting-started/pricing" + }, + "voyage/voyage-4-large": { + "input_cost_per_token": 1.2e-07, + "litellm_provider": "voyage", + "max_input_tokens": 32000, + "max_tokens": 32000, + "mode": "embedding", + "output_cost_per_token": 0.0, + "output_vector_size": 1024, + "source": "https://docs.voyageai.com/docs/pricing" + }, + "voyage/voyage-4": { + "input_cost_per_token": 6e-08, + "litellm_provider": "voyage", + "max_input_tokens": 32000, + "max_tokens": 32000, + "mode": "embedding", + "output_cost_per_token": 0.0, + "output_vector_size": 1024, + "source": "https://docs.voyageai.com/docs/pricing" + }, + "voyage/voyage-4-lite": { + "input_cost_per_token": 2e-08, + "litellm_provider": "voyage", + "max_input_tokens": 32000, + "max_tokens": 32000, + "mode": "embedding", + "output_cost_per_token": 0.0, + "output_vector_size": 1024, + "source": "https://docs.voyageai.com/docs/pricing" + }, + "voyage/voyage-code-4": { + "input_cost_per_token": 1.2e-07, + "litellm_provider": "voyage", + "max_input_tokens": 32000, + "max_tokens": 32000, + "mode": "embedding", + "output_cost_per_token": 0.0, + "output_vector_size": 1024, + "source": "https://docs.voyageai.com/docs/pricing" + }, + "voyage/voyage-context-4": { + "input_cost_per_token": 1.2e-07, + "litellm_provider": "voyage", + "max_input_tokens": 120000, + "max_tokens": 120000, + "mode": "embedding", + "output_cost_per_token": 0.0, + "output_vector_size": 1024, + "source": "https://docs.voyageai.com/docs/pricing" + }, + "voyage/voyage-multimodal-3.5": { + "input_cost_per_token": 1.2e-07, + "litellm_provider": "voyage", + "max_input_tokens": 32000, + "max_tokens": 32000, + "mode": "embedding", + "output_cost_per_token": 0.0, + "output_vector_size": 1024, + "source": "https://docs.voyageai.com/docs/pricing", + "supports_embedding_image_input": true + }, + "fireworks_ai/accounts/fireworks/models/deepseek-v4-flash-0731": { + "cache_read_input_token_cost": 2.8e-08, + "input_cost_per_token": 1.4e-07, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 2.8e-07, + "source": "https://docs.fireworks.ai/serverless/pricing", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": false + }, + "fireworks_ai/accounts/fireworks/models/kimi-k3": { + "cache_read_input_token_cost": 3e-07, + "input_cost_per_token": 3e-06, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 1.5e-05, + "source": "https://docs.fireworks.ai/serverless/pricing", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "fireworks_ai/deepseek-v4-flash-0731": { + "cache_read_input_token_cost": 2.8e-08, + "input_cost_per_token": 1.4e-07, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 2.8e-07, + "source": "https://docs.fireworks.ai/serverless/pricing", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": false + }, + "fireworks_ai/glm-5p2-fast": { + "cache_read_input_token_cost": 2.1e-07, + "input_cost_per_token": 2.1e-06, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 6.6e-06, + "source": "https://docs.fireworks.ai/serverless/pricing", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": false + }, + "fireworks_ai/glm-5p2-fast-us": { + "cache_read_input_token_cost": 2.1e-07, + "input_cost_per_token": 2.1e-06, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 6.6e-06, + "source": "https://docs.fireworks.ai/serverless/pricing", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": false + }, + "fireworks_ai/kimi-k3": { + "cache_read_input_token_cost": 3e-07, + "input_cost_per_token": 3e-06, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 1.5e-05, + "source": "https://docs.fireworks.ai/serverless/pricing", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "fireworks_ai/kimi-k3-fast": { + "cache_read_input_token_cost": 4.5e-07, + "input_cost_per_token": 4.5e-06, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 2.25e-05, + "source": "https://docs.fireworks.ai/serverless/pricing", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "fireworks_ai/kimi-k3-us": { + "cache_read_input_token_cost": 3.3e-07, + "input_cost_per_token": 3.3e-06, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 1.65e-05, + "source": "https://docs.fireworks.ai/serverless/pricing", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "fireworks_ai/qwen3p8-max": { + "cache_read_input_token_cost": 2.5e-07, + "input_cost_per_token": 2e-06, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 6e-06, + "source": "https://docs.fireworks.ai/serverless/pricing", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "fireworks_ai/muse-glimmer-30b": { + "cache_read_input_token_cost": 4e-08, + "input_cost_per_token": 3.5e-07, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 131072, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "output_cost_per_token": 1.5e-06, + "source": "https://docs.fireworks.ai/serverless/pricing", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "fireworks_ai/nemotron-lightning-3p5-30b-a3b": { + "cache_read_input_token_cost": 1e-08, + "input_cost_per_token": 5e-08, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 262144, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 2e-07, + "source": "https://docs.fireworks.ai/serverless/pricing", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": false + }, + "fireworks_ai/nemotron-3-ultra-nvfp4": { + "cache_read_input_token_cost": 1.2e-07, + "input_cost_per_token": 6e-07, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 262144, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 2.4e-06, + "source": "https://docs.fireworks.ai/serverless/pricing", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": false + }, + "fireworks_ai/accounts/fireworks/models/muse-glimmer-30b": { + "cache_read_input_token_cost": 4e-08, + "input_cost_per_token": 3.5e-07, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 131072, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "output_cost_per_token": 1.5e-06, + "source": "https://docs.fireworks.ai/serverless/pricing", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "fireworks_ai/accounts/fireworks/models/nemotron-lightning-3p5-30b-a3b": { + "cache_read_input_token_cost": 1e-08, + "input_cost_per_token": 5e-08, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 262144, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 2e-07, + "source": "https://docs.fireworks.ai/serverless/pricing", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": false + }, + "fireworks_ai/accounts/fireworks/models/nemotron-3-ultra-nvfp4": { + "cache_read_input_token_cost": 1.2e-07, + "input_cost_per_token": 6e-07, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 262144, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 2.4e-06, + "source": "https://docs.fireworks.ai/serverless/pricing", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": false + }, + "fireworks_ai/accounts/fireworks/models/qwen3p8-max": { + "cache_read_input_token_cost": 2.5e-07, + "input_cost_per_token": 2e-06, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 6e-06, + "source": "https://docs.fireworks.ai/serverless/pricing", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "fireworks_ai/accounts/fireworks/routers/glm-5p2-fast": { + "cache_read_input_token_cost": 2.1e-07, + "input_cost_per_token": 2.1e-06, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 6.6e-06, + "source": "https://docs.fireworks.ai/serverless/pricing", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": false + }, + "fireworks_ai/accounts/fireworks/routers/glm-5p2-fast-us": { + "cache_read_input_token_cost": 2.1e-07, + "input_cost_per_token": 2.1e-06, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 6.6e-06, + "source": "https://docs.fireworks.ai/serverless/pricing", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": false + }, + "fireworks_ai/accounts/fireworks/routers/kimi-k3-fast": { + "cache_read_input_token_cost": 4.5e-07, + "input_cost_per_token": 4.5e-06, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 2.25e-05, + "source": "https://docs.fireworks.ai/serverless/pricing", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "fireworks_ai/accounts/fireworks/routers/kimi-k3-us": { + "cache_read_input_token_cost": 3.3e-07, + "input_cost_per_token": 3.3e-06, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 1.65e-05, + "source": "https://docs.fireworks.ai/serverless/pricing", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true } } diff --git a/provider_endpoints_support.json b/provider_endpoints_support.json index ec0b1c27344..7c1ca34c23c 100644 --- a/provider_endpoints_support.json +++ b/provider_endpoints_support.json @@ -563,6 +563,23 @@ "interactions": true } }, + "cognition": { + "display_name": "Cognition (`cognition`)", + "url": "https://docs.litellm.ai/docs/providers/cognition", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": false + } + }, "cohere": { "display_name": "Cohere (`cohere`)", "url": "https://docs.litellm.ai/docs/providers/cohere", diff --git a/ruff-tests.toml b/ruff-tests.toml index c760a53ed91..60438d355f0 100644 --- a/ruff-tests.toml +++ b/ruff-tests.toml @@ -14,6 +14,18 @@ # B018 a bare attribute access or literal, usually a call missing its parens # PLW0127 `x = x` self-assignment, dead code that reads like a narrowing or a fixup # PLR0133 comparison of two constants, e.g. `assert True == True` +# B017 `pytest.raises(Exception)` accepts the TypeError a refactor introduced just as +# readily as the rejection under test, so a crash reads as a pass. Narrow to the +# real type, or add `match=` where the code genuinely raises a bare Exception +# PT012 a `pytest.raises` block that runs on past the raising call. Everything after +# that call is dead, so an `assert` sitting there is never checked. Keep the +# block to the call itself and put the assertions below it +# PT011 `pytest.raises(Exception)` / `(ValueError)` / `(OSError)` with no `match=`. The +# block passes on any error that broad, so the TypeError a refactor introduced +# reads as the rejection under test. Pin the message the code actually raises +# PT014 the same `parametrize` case listed twice. The copy re-runs an assertion that +# already passed and adds no coverage, and it usually marks a case someone meant +# to vary and forgot to edit # # No target-version here on purpose: it resolves from requires-python (>=3.10), so # 3.11-only builtins like BaseExceptionGroup are correctly flagged in a tree that @@ -21,4 +33,4 @@ line-length = 120 -lint.select = ["F821", "B011", "B015", "B018", "PT015", "PLR0133", "PLW0127"] +lint.select = ["F821", "B011", "B015", "B017", "B018", "PT011", "PT012", "PT014", "PT015", "PLR0133", "PLW0127"] diff --git a/test-quality-budget.json b/test-quality-budget.json index fcb29c3191d..1613c8c75cb 100644 --- a/test-quality-budget.json +++ b/test-quality-budget.json @@ -9,7 +9,7 @@ "limit": 1078 }, "TQ004": { - "limit": 770 + "limit": 768 }, "TQ005": { "limit": 2832 diff --git a/tests/e2e/batches/batch_client.py b/tests/e2e/batches/batch_client.py index 5cc5d1dae3b..968a357e8af 100644 --- a/tests/e2e/batches/batch_client.py +++ b/tests/e2e/batches/batch_client.py @@ -40,8 +40,15 @@ class FileObject(BaseModel): class FileList(BaseModel): + """GET /v1/files page. The cursors are modelled because they are part of the + page's isolation contract: they must address rows in `data`, never rows the + caller was not allowed to see.""" + object: str | None = None data: list[FileObject] = [] + first_id: str | None = None + last_id: str | None = None + has_more: bool | None = None class BatchObject(BaseModel): diff --git a/tests/e2e/batches/test_batches_e2e.py b/tests/e2e/batches/test_batches_e2e.py index 53bf9739983..12b848dd063 100644 --- a/tests/e2e/batches/test_batches_e2e.py +++ b/tests/e2e/batches/test_batches_e2e.py @@ -572,6 +572,41 @@ class TestOpenAIFiles: f"listed file must round-trip the upload purpose, got {match.purpose!r}" ) + @pytest.mark.covers( + "llm.files.openai.list_isolation.nonstream.works", + exercised_on=["files"], + ) + def test_list_page_cursors_address_only_the_callers_own_files( + self, client: BatchClient, resources: ResourceManager + ) -> None: + """Pins GitHub issue #36087: a list page's pagination cursors must address + rows in that page. + + The proxy fronts one shared provider account, so the upstream page is the + whole organization's. The gateway narrows `data` to the files the caller + owns, and `first_id` / `last_id` have to be narrowed with it: left as the + upstream org's, they hand any caller raw provider file ids belonging to + other tenants, which is the handle the file routes accept. + """ + key = resources.key(user_id=f"e2e-file-list-{unique_marker()}") + + listed = unwrap(client.list_files(key=key)) + + expected_first = listed.data[0].id if listed.data else None + expected_last = listed.data[-1].id if listed.data else None + assert listed.first_id == expected_first, ( + f"first_id {listed.first_id!r} is not the first row this caller can see " + f"({expected_first!r}); the page leaked another caller's file id" + ) + assert listed.last_id == expected_last, ( + f"last_id {listed.last_id!r} is not the last row this caller can see " + f"({expected_last!r}); the page leaked another caller's file id" + ) + assert listed.has_more is not True, ( + "the page advertises another page, but the proxy never forwards a cursor " + "upstream, so following it re-serves this same page forever" + ) + @pytest.mark.covers( "llm.files.openai.retrieve.nonstream.works", exercised_on=["files"], diff --git a/tests/e2e/coverage_registry/llm_conversational.yaml b/tests/e2e/coverage_registry/llm_conversational.yaml index 61ce30e3b81..44ed5765e38 100644 --- a/tests/e2e/coverage_registry/llm_conversational.yaml +++ b/tests/e2e/coverage_registry/llm_conversational.yaml @@ -63,6 +63,7 @@ - {id: llm.responses.openai.input_validation.nonstream.works, module: llm, tier: P1, subject_endpoint: responses, route: openai, capability: input_validation, streaming: nonstream, assertions: [works], source: "vendor strategy §9.9 / LIT-4778", rationale: "Responses missing/empty input and missing model are rejected"} - {id: llm.responses.openai.basic.stream.works, module: llm, tier: P0, subject_endpoint: responses, route: openai, capability: basic, streaming: stream, assertions: [works], source: "response_api_endpoints/endpoints.py:26", rationale: "Streaming via /v1/responses"} - {id: llm.responses.openai.basic.nonstream.cost_logged, module: llm, tier: P0, subject_endpoint: responses, route: openai, capability: basic, streaming: nonstream, assertions: [works, cost_logged], source: "response_api_endpoints/endpoints.py:26", rationale: "Cost logged on responses"} +- {id: llm.responses.openai.passthrough.stream.cost_logged, module: llm, tier: P0, subject_endpoint: responses, route: openai, capability: basic, streaming: stream, assertions: [cost_logged], source: "test_passthrough_e2e.py", rationale: "A streamed POST /openai_passthrough/v1/responses is costed and keyed by the provider response id; it used to log a zero-cost row under a random id (GitHub issue #36523)"} - {id: llm.responses.openai.tool_use.nonstream.works, module: llm, tier: P0, subject_endpoint: responses, route: openai, capability: tool_use, streaming: nonstream, assertions: [works], source: "model_prices json", rationale: "Tool calls via Responses API"} - {id: llm.responses.openai.vision.nonstream.works, module: llm, tier: P0, subject_endpoint: responses, route: openai, capability: vision, streaming: nonstream, assertions: [works], source: "model_prices json", rationale: "Vision via Responses API"} - {id: llm.responses.anthropic.basic.nonstream.works, module: llm, tier: P1, subject_endpoint: responses, route: anthropic, capability: basic, streaming: nonstream, assertions: [works], source: "response_api_endpoints/endpoints.py:26", rationale: "Responses w/ Anthropic translation (smoke)"} diff --git a/tests/e2e/coverage_registry/llm_nonconversational.yaml b/tests/e2e/coverage_registry/llm_nonconversational.yaml index bb7169509eb..8ae7dd01b5a 100644 --- a/tests/e2e/coverage_registry/llm_nonconversational.yaml +++ b/tests/e2e/coverage_registry/llm_nonconversational.yaml @@ -3,6 +3,7 @@ - {id: llm.embeddings.openai.basic.nonstream.works, module: llm, tier: P0, subject_endpoint: embeddings, route: openai, capability: basic, streaming: nonstream, assertions: [works], source: "test_embeddings_endpoint_e2e.py:23", rationale: "Core endpoint, live vector response"} - {id: llm.embeddings.openai.input_validation.nonstream.works, module: llm, tier: P1, subject_endpoint: embeddings, route: openai, capability: input_validation, streaming: nonstream, assertions: [works], source: "vendor strategy §9.3 / LIT-4778", rationale: "Missing model/input on /embeddings return client errors"} - {id: llm.embeddings.openai.basic.nonstream.cost_logged, module: llm, tier: P0, subject_endpoint: embeddings, route: openai, capability: basic, streaming: nonstream, assertions: [cost_logged], source: "SPEND_TRACKING_COVERAGE_MATRIX.md:34", rationale: "Cost tracking on embeddings"} +- {id: llm.embeddings.openai.passthrough.nonstream.cost_logged, module: llm, tier: P0, subject_endpoint: embeddings, route: openai, capability: basic, streaming: nonstream, assertions: [cost_logged], source: "test_passthrough_e2e.py", rationale: "POST /openai_passthrough/v1/embeddings is costed; the route wrote no spend row at all, so budgets never saw traffic OpenAI was billing for (GitHub issue #36646)"} - {id: llm.embeddings.azure_openai.basic.nonstream.works, module: llm, tier: P0, subject_endpoint: embeddings, route: azure_openai, capability: basic, streaming: nonstream, assertions: [works], source: "llms/azure/azure.py", rationale: "Azure embeddings via translation"} - {id: llm.embeddings.bedrock.basic.nonstream.works, module: llm, tier: P1, subject_endpoint: embeddings, route: bedrock_converse, capability: basic, streaming: nonstream, assertions: [works], source: "llms/bedrock/embed/embedding.py", rationale: "Bedrock Titan embeddings"} - {id: llm.embeddings.vertex.basic.nonstream.works, module: llm, tier: P1, subject_endpoint: embeddings, route: vertex, capability: basic, streaming: nonstream, assertions: [works], source: "vertex_embeddings/embedding_handler.py", rationale: "Vertex embeddings"} @@ -13,6 +14,7 @@ - {id: llm.batches.openai.cancel.nonstream.works, module: llm, tier: P0, subject_endpoint: batches, route: openai, capability: basic, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py", rationale: "Batch cancel"} - {id: llm.batches.openai.list.nonstream.works, module: llm, tier: P0, subject_endpoint: batches, route: openai, capability: basic, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py", rationale: "Batch list envelope"} - {id: llm.batches.openai.file_lifecycle.nonstream.works, module: llm, tier: P0, subject_endpoint: batches, route: openai, capability: basic, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py", rationale: "File upload/retrieve/delete for batch flow"} +- {id: llm.batches.openai.passthrough.nonstream.works, module: llm, tier: P1, subject_endpoint: batches, route: openai, capability: basic, streaming: nonstream, assertions: [works], source: "test_passthrough_e2e.py", rationale: "GET /openai_passthrough/v1/batches relays OpenAI's own batch page; the dedicated prefix must not bind as a provider name on the /{provider}/v1/batches route (GitHub issue #36086)"} - {id: llm.batches.openai_encoded.basic.nonstream.works, module: llm, tier: P0, subject_endpoint: batches, route: openai, capability: basic, streaming: nonstream, assertions: [works], source: "batches/capabilities.py", rationale: "Encoded scenario lifecycle"} - {id: llm.batches.openai_unified.basic.nonstream.works, module: llm, tier: P0, subject_endpoint: batches, route: openai, capability: basic, streaming: nonstream, assertions: [works], source: "batches/capabilities.py", rationale: "Unified/managed-id scenario"} - {id: llm.batches.openai_model_param.basic.nonstream.works, module: llm, tier: P0, subject_endpoint: batches, route: openai, capability: basic, streaming: nonstream, assertions: [works], source: "batches/capabilities.py", rationale: "Model-param scenario"} @@ -29,6 +31,8 @@ - {id: llm.files.openai.retrieve.nonstream.works, module: llm, tier: P0, subject_endpoint: files, route: openai, capability: basic, streaming: nonstream, assertions: [works], source: "files_endpoints.py", rationale: "File retrieve by id"} - {id: llm.files.openai.delete.nonstream.works, module: llm, tier: P0, subject_endpoint: files, route: openai, capability: basic, streaming: nonstream, assertions: [works], source: "files_endpoints.py", rationale: "File delete returns deleted=true"} - {id: llm.files.openai.list.nonstream.works, module: llm, tier: P0, subject_endpoint: files, route: openai, capability: basic, streaming: nonstream, assertions: [works], source: "files_endpoints.py", rationale: "File list paginated"} +- {id: llm.files.openai.list_isolation.nonstream.works, module: llm, tier: P0, subject_endpoint: files, route: openai, capability: basic, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py", rationale: "GET /v1/files pagination cursors address only rows the caller owns; on a shared provider account the upstream cursors otherwise hand out other tenants' raw provider file ids (GitHub issue #36087)"} +- {id: llm.files.openai.passthrough.nonstream.works, module: llm, tier: P1, subject_endpoint: files, route: openai, capability: basic, streaming: nonstream, assertions: [works], source: "test_passthrough_e2e.py", rationale: "POST/DELETE /openai_passthrough/v1/files relay OpenAI's own file object; the dedicated prefix must not bind as a provider name on the /{provider}/v1/files route (GitHub issue #36086)"} - {id: llm.files.azure_openai.upload.nonstream.works, module: llm, tier: P0, subject_endpoint: files, route: azure_openai, capability: basic, streaming: nonstream, assertions: [works], source: "batches/capabilities.py:45", rationale: "Azure file upload managed backend"} - {id: llm.files.vertex.upload.nonstream.works, module: llm, tier: P0, subject_endpoint: files, route: vertex, capability: basic, streaming: nonstream, assertions: [works], source: "batches/capabilities.py:52", rationale: "Vertex file upload to GCS"} - {id: llm.files.bedrock.upload.nonstream.works, module: llm, tier: P0, subject_endpoint: files, route: bedrock_converse, capability: basic, streaming: nonstream, assertions: [works], source: "batches/capabilities.py:59", rationale: "Bedrock file upload to S3"} diff --git a/tests/e2e/lifecycle.py b/tests/e2e/lifecycle.py index 4ef25509905..c9a67ebdb8c 100644 --- a/tests/e2e/lifecycle.py +++ b/tests/e2e/lifecycle.py @@ -52,7 +52,7 @@ class ResourceManager: """ client: ResourceClient - _cleanups: List[Callable[[], None]] = field( + _cleanups: List[Callable[[], object]] = field( default_factory=list ) # mutable-ok: append-only teardown registry @@ -60,8 +60,11 @@ class ResourceManager: """No global setup needed today; present for lifecycle symmetry.""" return None - def defer(self, cleanup: Callable[[], None]) -> None: - """Register a teardown action for any resource the test just created.""" + def defer(self, cleanup: Callable[[], object]) -> None: + """Register a teardown action for any resource the test just created. + + Whatever the action returns is discarded, so a delete that answers with a + response model can be deferred directly.""" self._cleanups.append(cleanup) def key(self, models: list[str] | None = None, user_id: str | None = "e2e-test-user") -> str: diff --git a/tests/e2e/llm_translation/passthrough_client.py b/tests/e2e/llm_translation/passthrough_client.py index 439594f3624..e0dfae679a9 100644 --- a/tests/e2e/llm_translation/passthrough_client.py +++ b/tests/e2e/llm_translation/passthrough_client.py @@ -15,7 +15,7 @@ from dataclasses import dataclass from pydantic import BaseModel, Field from proxy_client import ProxyClient -from e2e_http import Headers, StreamingResponse +from e2e_http import FileUploadForm, Headers, NoBody, Result, StreamingResponse from models import ChatMessage @@ -113,6 +113,76 @@ class OpenAIChatBody(BaseModel): max_completion_tokens: int = 64 +class PassthroughFileObject(BaseModel): + id: str + object: str | None = None + purpose: str | None = None + filename: str | None = None + bytes: int | None = None + + +class PassthroughFileDeleted(BaseModel): + id: str + deleted: bool + + +class PassthroughListEntry(BaseModel): + id: str + + +class ResponsesUsage(BaseModel): + input_tokens: int + output_tokens: int + + +class ResponsesObject(BaseModel): + id: str + usage: ResponsesUsage | None = None + + +class ResponsesStreamEvent(BaseModel): + """One SSE frame of a native Responses stream. Only the terminal frames carry a + `response`, so it stays optional and the deltas validate as themselves.""" + + type: str + response: ResponsesObject | None = None + + +def completed_responses_object(result: StreamingResponse) -> ResponsesObject | None: + """The `response.completed` frame's response object, or None if the stream never + completed. Its `id` is what the spend row is keyed by on this route, and its + usage is what the row is priced from.""" + events = ( + ResponsesStreamEvent.model_validate_json(payload) + for payload in result.stream_events + ) + completed = tuple( + event.response + for event in events + if event.type == "response.completed" and event.response is not None + ) + return completed[-1] if completed else None + + +class OpenAIResponsesBody(BaseModel): + model: str + input: str + stream: bool = False + + +class OpenAIEmbeddingBody(BaseModel): + model: str + input: str + + +class PassthroughBatchList(BaseModel): + """OpenAI's own batch page, relayed verbatim. `object` is required so a body + that is not an OpenAI list fails validation instead of passing vacuously.""" + + object: str + data: list[PassthroughListEntry] + + def _tags_header(tags: list[str] | None) -> str | None: return ",".join(tags) if tags else None @@ -196,6 +266,66 @@ class PassthroughClient: stream=stream, ) + # ---- OpenAI file/batch routes under /openai_passthrough ------------- + # + # Relayed to OpenAI untouched, which is the whole point of the prefix: the + # customer opts out of the gateway's managed-file handling here. + + def openai_passthrough_upload_file( + self, key: str, *, content: bytes, filename: str + ) -> Result[PassthroughFileObject]: + return self.proxy.transport.upload( + "/openai_passthrough/v1/files", + headers=self.proxy.transport.bearer(key), + form=FileUploadForm(purpose="batch"), + filename=filename, + content=content, + response_type=PassthroughFileObject, + ) + + def openai_passthrough_delete_file( + self, key: str, file_id: str + ) -> Result[PassthroughFileDeleted]: + return self.proxy.transport.delete( + f"/openai_passthrough/v1/files/{file_id}", + headers=self.proxy.transport.bearer(key), + json=NoBody(), + response_type=PassthroughFileDeleted, + ) + + def openai_passthrough_list_batches(self, key: str) -> Result[PassthroughBatchList]: + return self.proxy.transport.get( + "/openai_passthrough/v1/batches", + headers=self.proxy.transport.bearer(key), + params=NoBody(), + response_type=PassthroughBatchList, + ) + + # ---- OpenAI inference routes under /openai_passthrough ------------- + # + # Relayed to OpenAI verbatim, but still costed by the gateway: the customer + # budgets against this traffic, so a 200 that logs no spend is money the + # gateway never sees. + + def openai_passthrough_responses( + self, key: str, model: str, text: str, *, stream: bool = False + ) -> StreamingResponse: + return self.proxy.transport.send( + "/openai_passthrough/v1/responses", + headers=self.proxy.transport.bearer(key), + json=OpenAIResponsesBody(model=model, input=text, stream=stream), + stream=stream, + ) + + def openai_passthrough_embed( + self, key: str, model: str, text: str + ) -> StreamingResponse: + return self.proxy.transport.send( + "/openai_passthrough/v1/embeddings", + headers=self.proxy.transport.bearer(key), + json=OpenAIEmbeddingBody(model=model, input=text), + ) + def openai_chat( self, key: str, model: str, text: str, *, max_completion_tokens: int = 64 ) -> StreamingResponse: diff --git a/tests/e2e/llm_translation/test_passthrough_e2e.py b/tests/e2e/llm_translation/test_passthrough_e2e.py index b57164df9bb..17b0dbe1ae5 100644 --- a/tests/e2e/llm_translation/test_passthrough_e2e.py +++ b/tests/e2e/llm_translation/test_passthrough_e2e.py @@ -13,8 +13,8 @@ A passthrough call returning non-2xx fails hard (never a skip); once it returns import pytest -from e2e_config import unique_marker -from e2e_http import StreamingResponse, require_successful_call +from e2e_config import CHEAP_OPENAI_MODEL, unique_marker +from e2e_http import StreamingResponse, require_successful_call, unwrap from lifecycle import ResourceManager from models import KeyGenerateBody, SpendLogRow from passthrough_client import ( @@ -24,8 +24,11 @@ from passthrough_client import ( JsonSchema, JsonSchemaProperty, PassthroughClient, + completed_responses_object, ) +EMBEDDING_MODEL = "text-embedding-3-small" + pytestmark = pytest.mark.e2e @@ -210,3 +213,129 @@ class TestPassthroughModelAllowlist: "a key restricted to gemini-2.5-flash must be denied a claude passthrough call, " f"got {result.status_code}: {result.body[:300]}" ) + + +class TestOpenAIPassthroughPrefix: + """The dedicated `/openai_passthrough` prefix must reach OpenAI, not be + swallowed by the provider-scoped `/{provider}/v1/...` routes. + + The customer fronts OpenAI's own file and batch APIs through this prefix + precisely to opt out of the gateway's managed-file handling. `/v1/files` and + `/v1/batches` also answer `/{provider}/v1/files` and `/{provider}/v1/batches`, + so `openai_passthrough` used to bind as a provider name and the request died + inside the gateway with a provider-lookup error, never reaching OpenAI. + """ + + @pytest.mark.covers("llm.files.openai.passthrough.nonstream.works") + def test_passthrough_prefix_uploads_a_file_to_openai( + self, client: PassthroughClient, resources: ResourceManager, scoped_key: str + ) -> None: + """Pins GitHub issue #36086: a file upload through the dedicated prefix + reaches OpenAI's file API instead of 500ing on a provider-name lookup.""" + content = f'{{"marker":"{unique_marker()}"}}\n'.encode() + uploaded = unwrap( + client.openai_passthrough_upload_file( + scoped_key, content=content, filename="e2e-passthrough-batch.jsonl" + ) + ) + resources.defer( + lambda: client.openai_passthrough_delete_file(scoped_key, uploaded.id) + ) + + assert uploaded.object == "file", ( + f"/openai_passthrough/v1/files did not relay OpenAI's file object: {uploaded}" + ) + assert uploaded.purpose == "batch" + assert uploaded.bytes == len(content) + + @pytest.mark.covers("llm.batches.openai.passthrough.nonstream.works") + def test_passthrough_prefix_lists_batches_from_openai( + self, client: PassthroughClient, scoped_key: str + ) -> None: + """Pins GitHub issue #36086 on the batches route: the dedicated prefix + relays OpenAI's own batch page instead of dying on the provider lookup.""" + listed = unwrap(client.openai_passthrough_list_batches(scoped_key)) + + assert listed.object == "list", ( + f"/openai_passthrough/v1/batches did not relay OpenAI's batch page: {listed}" + ) + + +class TestOpenAIPassthroughSpend: + """A call relayed to OpenAI's own endpoints must still be costed. + + The customer routes native OpenAI traffic through `/openai_passthrough` and + budgets against it, so a call that returns 200 while logging no spend is money + the gateway never sees and a budget that never trips. Streamed Responses calls + and embeddings each used to land exactly that way, on separate code paths. + """ + + @pytest.mark.covers("llm.responses.openai.passthrough.stream.cost_logged") + def test_streamed_responses_call_logs_its_cost( + self, client: PassthroughClient, scoped_key: str + ) -> None: + """Pins GitHub issue #36523: a streamed passthrough Responses call is billed + under the provider id the caller was served, never a $0 row under a random + id.""" + result = client.openai_passthrough_responses( + scoped_key, + CHEAP_OPENAI_MODEL, + f"Say hi in one word. {unique_marker()}", + stream=True, + ) + require_successful_call(result) + assert result.chunks > 0, "streamed responses passthrough produced no events" + + completed = completed_responses_object(result) + assert completed is not None, ( + f"the stream never delivered a response.completed frame, so there is no " + f"provider id to reconcile against: last events {result.stream_events[-3:]}" + ) + assert completed.usage is not None, ( + f"the completed response carried no usage to price from: {completed}" + ) + + rows = client.proxy.poll_logs_for_request_id( + completed.id, predicate=lambda rows: (rows[0].spend or 0) > 0 + ) + assert rows, ( + f"no spend row for the response the customer was served ({completed.id}); " + "a streamed passthrough call OpenAI bills them for is invisible to the " + "gateway's own spend and budgets" + ) + row = rows[0] + assert (row.spend or 0) > 0, f"streamed responses passthrough was not costed: {row}" + assert row.prompt_tokens == completed.usage.input_tokens, ( + f"logged {row.prompt_tokens} prompt tokens, the response the customer read " + f"reported {completed.usage.input_tokens}" + ) + assert row.completion_tokens == completed.usage.output_tokens, ( + f"logged {row.completion_tokens} completion tokens, the response the customer " + f"read reported {completed.usage.output_tokens}" + ) + + @pytest.mark.covers("llm.embeddings.openai.passthrough.nonstream.cost_logged") + def test_embeddings_call_logs_its_cost( + self, client: PassthroughClient, scoped_key: str + ) -> None: + """Pins GitHub issue #36646: a passthrough embeddings call writes a priced + spend row instead of no row at all.""" + result = client.openai_passthrough_embed( + scoped_key, EMBEDDING_MODEL, f"cost this sentence {unique_marker()}" + ) + require_successful_call(result) + assert result.call_id, "embeddings passthrough returned no x-litellm-call-id" + + rows = client.proxy.poll_logs_for_request_id( + result.call_id, predicate=lambda rows: (rows[0].spend or 0) > 0 + ) + assert rows, ( + f"no spend row for embeddings call {result.call_id}; the customer is billed " + "by OpenAI for tokens the gateway never counted against their budget" + ) + row = rows[0] + assert (row.spend or 0) > 0, f"embeddings passthrough was not costed: {row}" + assert (row.prompt_tokens or 0) > 0, ( + f"the embeddings row logged no prompt tokens, so whatever cost it carries " + f"was not computed from the real usage: {row}" + ) diff --git a/tests/enterprise/litellm_enterprise/integrations/test_prometheus.py b/tests/enterprise/litellm_enterprise/integrations/test_prometheus.py index 0cd6055e09d..bdf73b6ab03 100644 --- a/tests/enterprise/litellm_enterprise/integrations/test_prometheus.py +++ b/tests/enterprise/litellm_enterprise/integrations/test_prometheus.py @@ -400,7 +400,7 @@ def test_invalid_metric_name_validation(): litellm.prometheus_metrics_config = test_config # Creating PrometheusLogger should raise ValueError due to invalid metric - with pytest.raises(ValueError) as exc_info: + with pytest.raises(ValueError, match='Configuration validation failed') as exc_info: PrometheusLogger() # Verify error message contains information about invalid metric @@ -429,7 +429,7 @@ def test_invalid_labels_validation(): litellm.prometheus_metrics_config = test_config # Creating PrometheusLogger should raise ValueError due to invalid labels - with pytest.raises(ValueError) as exc_info: + with pytest.raises(ValueError, match='Configuration validation failed') as exc_info: PrometheusLogger() # Verify error message contains information about invalid labels @@ -598,7 +598,7 @@ def test_invalid_exclude_metric_name_raises(reset_prometheus_exclude_settings): litellm.prometheus_exclude_labels = None litellm.prometheus_exclude_metrics = ["not_a_real_metric"] - with pytest.raises(ValueError) as exc_info: + with pytest.raises(ValueError, match='Prometheus exclude configuration validation failed') as exc_info: PrometheusLogger() assert "not_a_real_metric" in str(exc_info.value) @@ -612,7 +612,7 @@ def test_invalid_exclude_label_name_raises(reset_prometheus_exclude_settings): litellm.prometheus_exclude_metrics = None litellm.prometheus_exclude_labels = ["not_a_real_label"] - with pytest.raises(ValueError) as exc_info: + with pytest.raises(ValueError, match='Prometheus exclude configuration validation failed') as exc_info: PrometheusLogger() assert "not_a_real_label" in str(exc_info.value) diff --git a/tests/enterprise/litellm_enterprise/proxy/guardrails/test_bedrock_apply_guardrail.py b/tests/enterprise/litellm_enterprise/proxy/guardrails/test_bedrock_apply_guardrail.py index f257b47404e..6b6b5d768dd 100644 --- a/tests/enterprise/litellm_enterprise/proxy/guardrails/test_bedrock_apply_guardrail.py +++ b/tests/enterprise/litellm_enterprise/proxy/guardrails/test_bedrock_apply_guardrail.py @@ -141,7 +141,7 @@ async def test_bedrock_apply_guardrail_api_failure(): mock_api_request.side_effect = Exception("API connection failed") # Test the apply_guardrail method should raise an exception - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Bedrock guardrail failed: API connection failed') as exc_info: await guardrail.apply_guardrail( inputs={"texts": ["This is a test message"]}, request_data={}, diff --git a/tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py b/tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py index 714f3be6df9..2d845a445b5 100644 --- a/tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py +++ b/tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py @@ -1653,7 +1653,7 @@ async def test_afile_retrieve_raises_error_when_no_router_and_file_object_none() unified_file_id = "test-unified-file-id" - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='LiteLLM Managed File object with id=test-unified-file-id') as exc_info: await proxy_managed_files.afile_retrieve( file_id=unified_file_id, litellm_parent_otel_span=None, @@ -1719,7 +1719,7 @@ async def test_afile_retrieve_raises_error_for_non_managed_file(): # Mock get_unified_file_id to return None (file not found) proxy_managed_files.get_unified_file_id = AsyncMock(return_value=None) - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='LiteLLM Managed File object with id=non-existent-file-id') as exc_info: await proxy_managed_files.afile_retrieve( file_id="non-existent-file-id", litellm_parent_otel_span=None, @@ -2027,7 +2027,7 @@ async def test_list_batches_from_managed_objects_table_provider_filter_raises_ex ) # Filtering by provider should raise Exception - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match="Filtering by 'provider' is not supported when using managed") as exc_info: await proxy_managed_files.list_user_batches( user_api_key_dict=UserAPIKeyAuth(user_id="test-user"), limit=10, @@ -2053,7 +2053,7 @@ async def test_list_batches_from_managed_objects_table_target_model_name_filter_ ) # Filtering by provider should raise Exception - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match="Filtering by 'target_model_names' is not supported when") as exc_info: await proxy_managed_files.list_user_batches( user_api_key_dict=UserAPIKeyAuth(user_id="test-user"), limit=10, diff --git a/tests/enterprise/litellm_enterprise/proxy/management_endpoints/test_project_endpoints_prisma.py b/tests/enterprise/litellm_enterprise/proxy/management_endpoints/test_project_endpoints_prisma.py index c29b4c68bb0..ed6735a7126 100644 --- a/tests/enterprise/litellm_enterprise/proxy/management_endpoints/test_project_endpoints_prisma.py +++ b/tests/enterprise/litellm_enterprise/proxy/management_endpoints/test_project_endpoints_prisma.py @@ -29,6 +29,7 @@ from litellm_enterprise.proxy.management_endpoints.project_endpoints import ( from litellm.proxy.proxy_server import ( LitellmUserRoles, ) +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.utils import PrismaClient, ProxyLogging verbose_proxy_logger.setLevel(level=logging.DEBUG) @@ -447,7 +448,7 @@ def test_check_team_project_limits_models_not_in_team(): models=["gpt-5.5", "claude-3"], # claude-3 not in team ) - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match="not in team's allowed models\\. Team allowed models") as exc_info: _check_team_project_limits(team_object=team, data=data) assert "claude-3" in str(exc_info.value.detail) @@ -475,7 +476,7 @@ def test_check_team_project_limits_budget_exceeds_team(): max_budget=150.0, # exceeds team's 100.0 ) - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Project max_budget') as exc_info: _check_team_project_limits(team_object=team, data=data) assert "exceeds team's max_budget" in str(exc_info.value.detail) @@ -550,7 +551,7 @@ def test_check_team_project_limits_tpm_exceeds_team(): tpm_limit=20000, # exceeds team's 10000 ) - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Project tpm_limit') as exc_info: _check_team_project_limits(team_object=team, data=data) assert "exceeds team's tpm_limit" in str(exc_info.value.detail) @@ -576,7 +577,7 @@ def test_check_team_project_limits_negative_budget(): max_budget=-10.0, ) - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='max_budget cannot be negative\\. Received') as exc_info: _check_team_project_limits(team_object=team, data=data) assert "cannot be negative" in str(exc_info.value.detail) @@ -603,7 +604,7 @@ def test_check_team_project_limits_soft_budget_gte_max(): soft_budget=100.0, # equal to max, should fail ) - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='must be strictly lower than max_budget') as exc_info: _check_team_project_limits(team_object=team, data=data) assert "must be strictly lower" in str(exc_info.value.detail) @@ -1039,3 +1040,70 @@ async def test_project_eviction_publishes_cross_worker_invalidation(monkeypatch) ) mock_publish.assert_awaited_once_with(cache_key=f"project_id:{project_id}") + + +def _project_update_mocks(monkeypatch, stored_metadata: dict) -> mock.MagicMock: + existing_row = mock.MagicMock( + team_id=None, budget_id=None, object_permission_id=None, metadata=stored_metadata + ) + mock_prisma = mock.MagicMock() + mock_prisma.jsonify_object = lambda data: data + mock_prisma.db.litellm_projecttable.find_unique = mock.AsyncMock(return_value=existing_row) + mock_prisma.db.litellm_projecttable.update = mock.AsyncMock(return_value=mock.MagicMock()) + + monkeypatch.setattr(litellm.proxy.proxy_server, "premium_user", True) + monkeypatch.setattr(litellm.proxy.proxy_server, "prisma_client", mock_prisma) + monkeypatch.setattr(litellm.proxy.proxy_server, "user_api_key_cache", UserApiKeyCache()) + return mock_prisma + + +async def _run_project_update(project_id: str, **fields) -> None: + await update_project( + data=UpdateProjectRequest(project_id=project_id, **fields), + http_request=Request(scope={"type": "http"}), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + api_key="sk-1234", + user_id="1234", + ), + ) + + +def _written_project_data(mock_prisma: mock.MagicMock) -> dict: + return mock_prisma.db.litellm_projecttable.update.await_args.kwargs["data"] + + +@pytest.mark.asyncio +async def test_update_project_clears_model_itpm_limit_sent_as_an_empty_map(monkeypatch): + """ + LIT-4693 regression: an omitted key means "leave this alone", so the only way to drop a + per-model input/output TPM quota is to send it as an explicitly empty map. The written + metadata must stop carrying the quota, otherwise the proxy keeps enforcing a limit the + operator has already removed in the UI. + """ + project_id = f"project-{uuid.uuid4()}" + mock_prisma = _project_update_mocks( + monkeypatch, + {"owner": "platform", "model_itpm_limit": {"gpt-4": 60}, "model_otpm_limit": {"gpt-4": 40}}, + ) + + await _run_project_update(project_id, model_itpm_limit={}, model_otpm_limit={}) + + written_metadata = _written_project_data(mock_prisma)["metadata"] + assert written_metadata["model_itpm_limit"] == {} + assert written_metadata["model_otpm_limit"] == {} + + +@pytest.mark.asyncio +async def test_update_project_leaves_metadata_untouched_when_no_limit_is_sent(monkeypatch): + """ + The other half of the same contract: an update that says nothing about the limits must not + write metadata at all. That is what makes a dropped key silently preserve the old quota, so + the UI has to send the empty map instead of omitting it. + """ + project_id = f"project-{uuid.uuid4()}" + mock_prisma = _project_update_mocks(monkeypatch, {"model_itpm_limit": {"gpt-4": 60}}) + + await _run_project_update(project_id, description="renamed only") + + assert "metadata" not in _written_project_data(mock_prisma) diff --git a/tests/guardrails_tests/test_bedrock_guardrails.py b/tests/guardrails_tests/test_bedrock_guardrails.py index 823ee05839f..8b22cc0eb73 100644 --- a/tests/guardrails_tests/test_bedrock_guardrails.py +++ b/tests/guardrails_tests/test_bedrock_guardrails.py @@ -197,6 +197,7 @@ async def test_bedrock_guardrails_block_responses_api(): @pytest.mark.asyncio async def test_bedrock_guardrails_with_streaming(): + from fastapi import HTTPException from litellm.proxy.utils import ProxyLogging from litellm.types.guardrails import GuardrailEventHooks @@ -204,7 +205,7 @@ async def test_bedrock_guardrails_with_streaming(): mock_user_api_key_cache = MagicMock(spec=DualCache) mock_user_api_key_dict = UserAPIKeyAuth() - with pytest.raises(Exception): # Assert that this raises an exception + async def _stream_through_guardrail(): proxy_logging_obj = ProxyLogging( user_api_key_cache=mock_user_api_key_cache, premium_user=True, @@ -239,6 +240,9 @@ async def test_bedrock_guardrails_with_streaming(): async for chunk in response: print(chunk) + with pytest.raises(HTTPException): + await _stream_through_guardrail() + @pytest.mark.asyncio async def test_bedrock_guardrails_with_streaming_no_violation(): @@ -1501,7 +1505,7 @@ async def test_bedrock_guardrail_disable_exception_on_block_streaming(): mock_post.return_value = mock_bedrock_response # Should raise exception during streaming processing - with pytest.raises(HTTPException): + async def _drain(): result_generator = ( guardrail_default.async_post_call_streaming_iterator_hook( user_api_key_dict=mock_user_api_key_dict, @@ -1510,10 +1514,12 @@ async def test_bedrock_guardrail_disable_exception_on_block_streaming(): ) ) - # Try to consume the generator - should raise exception async for chunk in result_generator: pass + with pytest.raises(HTTPException): + await _drain() + # Test 2: disable_exception_on_block=True. Streaming can't raise up to the # endpoint handler (SSE headers already flushed), so the block is delivered # as a synthetic stream with finish_reason=content_filter and the block diff --git a/tests/guardrails_tests/test_dynamoai_guardrails.py b/tests/guardrails_tests/test_dynamoai_guardrails.py index 98f676a71d5..6f0ea00165b 100644 --- a/tests/guardrails_tests/test_dynamoai_guardrails.py +++ b/tests/guardrails_tests/test_dynamoai_guardrails.py @@ -61,7 +61,7 @@ async def test_dynamoai_blocks_content_with_block_action(): guardrail.should_run_guardrail = MagicMock(return_value=True) # Test that the guardrail raises ValueError for blocked content - with pytest.raises(ValueError) as exc_info: + with pytest.raises(ValueError, match='violation\\(s\\) detected') as exc_info: await guardrail.async_pre_call_hook( data=request_data, user_api_key_dict=UserAPIKeyAuth(), diff --git a/tests/guardrails_tests/test_eu_ai_act_article5.py b/tests/guardrails_tests/test_eu_ai_act_article5.py index bda7bf6f517..f7384667481 100644 --- a/tests/guardrails_tests/test_eu_ai_act_article5.py +++ b/tests/guardrails_tests/test_eu_ai_act_article5.py @@ -20,6 +20,7 @@ from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.content_fil from litellm.types.proxy.guardrails.guardrail_hooks.litellm_content_filter import ( ContentFilterCategoryConfig, ) +from fastapi import HTTPException # Test cases: (sentence, expected_result, reason) @@ -210,7 +211,7 @@ class TestEUAIActArticle5ConditionalMatching: # Apply guardrail if expected == "BLOCK": # Should raise an exception or return modified response indicating block - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Content blocked: eu_ai_act_article') as exc_info: await content_filter_guardrail.apply_guardrail( inputs={"texts": [sentence]}, request_data=request_data, @@ -275,7 +276,7 @@ class TestEUAIActEdgeCases: for sentence in sentences: request_data = {"messages": [{"role": "user", "content": sentence}]} - with pytest.raises(Exception): + with pytest.raises(HTTPException): await content_filter_guardrail.apply_guardrail( inputs={"texts": [sentence]}, request_data=request_data, @@ -289,7 +290,7 @@ class TestEUAIActEdgeCases: request_data = {"messages": [{"role": "user", "content": sentence}]} # Should block (contains multiple violations) - with pytest.raises(Exception): + with pytest.raises(HTTPException): await content_filter_guardrail.apply_guardrail( inputs={"texts": [sentence]}, request_data=request_data, diff --git a/tests/guardrails_tests/test_eu_ai_act_french_3_scenarios.py b/tests/guardrails_tests/test_eu_ai_act_french_3_scenarios.py index bc121330a45..221ca5aa6e6 100644 --- a/tests/guardrails_tests/test_eu_ai_act_french_3_scenarios.py +++ b/tests/guardrails_tests/test_eu_ai_act_french_3_scenarios.py @@ -19,6 +19,7 @@ from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.content_fil from litellm.types.proxy.guardrails.guardrail_hooks.litellm_content_filter import ( ContentFilterCategoryConfig, ) +from fastapi import HTTPException @pytest.fixture @@ -82,7 +83,7 @@ class TestEUAIActFrench3Scenarios: print(f"{'='*70}\n") # Should raise an exception (blocked) - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match="prohibited_practices_fr conditional match 'concevoir \\+") as exc_info: await content_filter_guardrail.apply_guardrail( inputs={"texts": [sentence]}, request_data=request_data, @@ -122,7 +123,7 @@ class TestEUAIActFrench3Scenarios: print(f"{'='*70}\n") # Should raise an exception (blocked) - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match="prohibited_practices_fr conditional match 'créer \\+") as exc_info: await content_filter_guardrail.apply_guardrail( inputs={"texts": [sentence]}, request_data=request_data, @@ -193,7 +194,7 @@ class TestEUAIActFrench3Scenarios: print(f"{'='*70}\n") # Should raise an exception (blocked by conditional matching) - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match="prohibited_practices_fr conditional match 'développer \\+") as exc_info: await content_filter_guardrail.apply_guardrail( inputs={"texts": [sentence]}, request_data=request_data, @@ -228,7 +229,7 @@ class TestFrenchEdgeCases: request_data = {"messages": [{"role": "user", "content": sentence}]} # Should block (contains "build" and "système de crédit social") - with pytest.raises(Exception): + with pytest.raises(HTTPException): await content_filter_guardrail.apply_guardrail( inputs={"texts": [sentence]}, request_data=request_data, @@ -257,7 +258,7 @@ class TestFrenchEdgeCases: request_data = {"messages": [{"role": "user", "content": sentence}]} # Should block (case-insensitive) - with pytest.raises(Exception): + with pytest.raises(HTTPException): await content_filter_guardrail.apply_guardrail( inputs={"texts": [sentence]}, request_data=request_data, @@ -277,7 +278,7 @@ class TestFrenchEdgeCases: request_data = {"messages": [{"role": "user", "content": sentence}]} # Should still block (no exception bypass) - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match="prohibited_practices_fr conditional match 'créer \\+ crédit") as exc_info: await content_filter_guardrail.apply_guardrail( inputs={"texts": [sentence]}, request_data=request_data, diff --git a/tests/guardrails_tests/test_semantic_guard.py b/tests/guardrails_tests/test_semantic_guard.py index a7e6230d029..c9f4a902895 100644 --- a/tests/guardrails_tests/test_semantic_guard.py +++ b/tests/guardrails_tests/test_semantic_guard.py @@ -10,6 +10,7 @@ sys.path.insert(0, os.path.abspath("../..")) from unittest.mock import MagicMock import pytest +from fastapi import HTTPException class TestRouteLoader: @@ -307,7 +308,7 @@ class TestContentFilterSqlInjectionTemplate: @pytest.mark.asyncio async def test_sql_always_block(self, sql_injection_guardrail, sentence, reason): request_data = {"messages": [{"role": "user", "content": sentence}]} - with pytest.raises(Exception): + with pytest.raises(HTTPException): await sql_injection_guardrail.apply_guardrail( inputs={"texts": [sentence]}, request_data=request_data, @@ -343,7 +344,7 @@ class TestContentFilterSqlInjectionTemplate: self, sql_injection_guardrail, sentence, reason ): request_data = {"messages": [{"role": "user", "content": sentence}]} - with pytest.raises(Exception): + with pytest.raises(HTTPException): await sql_injection_guardrail.apply_guardrail( inputs={"texts": [sentence]}, request_data=request_data, @@ -552,7 +553,7 @@ class TestContentFilterPromptInjectionTemplate: @pytest.mark.asyncio async def test_always_block(self, content_filter_guardrail, sentence, reason): request_data = {"messages": [{"role": "user", "content": sentence}]} - with pytest.raises(Exception): + with pytest.raises(HTTPException): await content_filter_guardrail.apply_guardrail( inputs={"texts": [sentence]}, request_data=request_data, diff --git a/tests/guardrails_tests/test_sg_mas_ai_guardrails.py b/tests/guardrails_tests/test_sg_mas_ai_guardrails.py index 668ee704692..e587d666a79 100644 --- a/tests/guardrails_tests/test_sg_mas_ai_guardrails.py +++ b/tests/guardrails_tests/test_sg_mas_ai_guardrails.py @@ -55,7 +55,7 @@ def _make_guardrail(yaml_filename: str, category_name: str) -> ContentFilterGuar async def _expect_block(guardrail: ContentFilterGuardrail, sentence: str, reason: str): request_data = {"messages": [{"role": "user", "content": sentence}]} - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Content blocked: sg_mas_') as exc_info: await guardrail.apply_guardrail( inputs={"texts": [sentence]}, request_data=request_data, diff --git a/tests/guardrails_tests/test_sg_pdpa_guardrails.py b/tests/guardrails_tests/test_sg_pdpa_guardrails.py index fd7133bc745..42c3a15f9f6 100644 --- a/tests/guardrails_tests/test_sg_pdpa_guardrails.py +++ b/tests/guardrails_tests/test_sg_pdpa_guardrails.py @@ -62,7 +62,7 @@ def _make_guardrail(yaml_filename: str, category_name: str) -> ContentFilterGuar async def _expect_block(guardrail: ContentFilterGuardrail, sentence: str, reason: str): """Assert that the guardrail BLOCKS the sentence.""" request_data = {"messages": [{"role": "user", "content": sentence}]} - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Content blocked: sg_pdpa_') as exc_info: await guardrail.apply_guardrail( inputs={"texts": [sentence]}, request_data=request_data, diff --git a/tests/image_gen_tests/test_image_generation.py b/tests/image_gen_tests/test_image_generation.py index ad141a651e8..9047557c493 100644 --- a/tests/image_gen_tests/test_image_generation.py +++ b/tests/image_gen_tests/test_image_generation.py @@ -444,7 +444,7 @@ async def test_azure_image_generation_request_body(): ) as mock_post: mock_post.side_effect = Exception("test") - with pytest.raises(Exception): + with pytest.raises(litellm.APIConnectionError): await aimage_generation( model="azure/gpt-image-1", prompt="test prompt", diff --git a/tests/litellm_utils_tests/test_hashicorp.py b/tests/litellm_utils_tests/test_hashicorp.py index 9aff7ddc10e..fa39a045227 100644 --- a/tests/litellm_utils_tests/test_hashicorp.py +++ b/tests/litellm_utils_tests/test_hashicorp.py @@ -432,7 +432,7 @@ def test_hashicorp_get_url_rejects_path_traversal(monkeypatch, malicious_secret_ monkeypatch.setenv("HCP_VAULT_TOKEN", "test-token-for-get-url-only") manager = HashicorpSecretManager() - with pytest.raises(ValueError): + with pytest.raises(ValueError, match='Invalid secret_name'): manager.get_url(malicious_secret_name) diff --git a/tests/litellm_utils_tests/test_proxy_budget_reset.py b/tests/litellm_utils_tests/test_proxy_budget_reset.py index b13b7342c25..a188fcf9d72 100644 --- a/tests/litellm_utils_tests/test_proxy_budget_reset.py +++ b/tests/litellm_utils_tests/test_proxy_budget_reset.py @@ -748,8 +748,9 @@ async def test_service_logger_keys_failure(): ) = proxy_logging_obj.service_logging_obj.async_service_failure_hook.call_args event_metadata = kwargs.get("event_metadata", {}) assert event_metadata.get("num_keys_found") == len(keys) - keys_found_str = event_metadata.get("keys_found", "") - assert "key1" in keys_found_str + # the row payload is deliberately absent: serializing every found row on the + # event loop is what blocked auth on the sweeping pod + assert "keys_found" not in event_metadata # Success hook should not be called. proxy_logging_obj.service_logging_obj.async_service_success_hook.assert_not_called() @@ -866,8 +867,7 @@ async def test_service_logger_users_failure(): ) = proxy_logging_obj.service_logging_obj.async_service_failure_hook.call_args event_metadata = kwargs.get("event_metadata", {}) assert event_metadata.get("num_users_found") == len(users) - users_found_str = event_metadata.get("users_found", "") - assert "user1" in users_found_str + assert "users_found" not in event_metadata proxy_logging_obj.service_logging_obj.async_service_success_hook.assert_not_called() @@ -983,8 +983,7 @@ async def test_service_logger_teams_failure(): ) = proxy_logging_obj.service_logging_obj.async_service_failure_hook.call_args event_metadata = kwargs.get("event_metadata", {}) assert event_metadata.get("num_teams_found") == len(teams) - teams_found_str = event_metadata.get("teams_found", "") - assert "team1" in teams_found_str + assert "teams_found" not in event_metadata proxy_logging_obj.service_logging_obj.async_service_success_hook.assert_not_called() @@ -1113,8 +1112,8 @@ async def test_service_logger_endusers_failure(): event_metadata = kwargs.get("event_metadata", {}) assert event_metadata.get("num_budgets_found") == len(budgets) assert event_metadata.get("num_endusers_found") == len(endusers) - endusers_found_str = event_metadata.get("endusers_found", "") - assert "user1" in endusers_found_str + assert "endusers_found" not in event_metadata + assert "budgets_found" not in event_metadata proxy_logging_obj.service_logging_obj.async_service_success_hook.assert_not_called() diff --git a/tests/litellm_utils_tests/test_utils.py b/tests/litellm_utils_tests/test_utils.py index 697c3837602..3c73224d7a1 100644 --- a/tests/litellm_utils_tests/test_utils.py +++ b/tests/litellm_utils_tests/test_utils.py @@ -1334,7 +1334,7 @@ def test_validate_chat_completion_user_messages(messages, expected_bool): validate_chat_completion_user_messages(messages=messages) else: ## Invalid message - with pytest.raises(Exception): + with pytest.raises(Exception, match="Invalid user message at index 0"): validate_chat_completion_user_messages(messages=messages) @@ -1354,7 +1354,7 @@ def test_validate_chat_completion_tool_choice(tool_choice, expected_bool): if expected_bool: validate_chat_completion_tool_choice(tool_choice=tool_choice) else: - with pytest.raises(Exception): + with pytest.raises(Exception, match="Invalid tool choice"): validate_chat_completion_tool_choice(tool_choice=tool_choice) @@ -2147,7 +2147,7 @@ def test_validate_user_messages_invalid_content_type(): messages = [{"content": [{"type": "invalid_type", "text": "Hello"}]}] - with pytest.raises(Exception) as e: + with pytest.raises(Exception, match='Please ensure all messages are valid OpenAI chat completion') as e: validate_chat_completion_user_messages(messages) assert "Invalid message" in str(e) diff --git a/tests/litellm_utils_tests/test_validate_tool_choice.py b/tests/litellm_utils_tests/test_validate_tool_choice.py index 0e6294a7cd4..07f8c9ed8f4 100644 --- a/tests/litellm_utils_tests/test_validate_tool_choice.py +++ b/tests/litellm_utils_tests/test_validate_tool_choice.py @@ -37,27 +37,27 @@ def test_validate_tool_choice_cursor_format(): def test_validate_tool_choice_invalid_dict(): """Test that invalid dict formats raise exceptions.""" # Missing both type and function - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Invalid tool choice, tool_choice=\\{\\}\\. Please ensure') as exc_info: validate_chat_completion_tool_choice({}) assert "Invalid tool choice" in str(exc_info.value) # Invalid type value - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match="Invalid tool choice, tool_choice=\\{'type': 'invalid'\\}\\.") as exc_info: validate_chat_completion_tool_choice({"type": "invalid"}) assert "Invalid tool choice" in str(exc_info.value) # Has type but missing function when type is "function" - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match="Invalid tool choice, tool_choice=\\{'type': 'function'\\}\\.") as exc_info: validate_chat_completion_tool_choice({"type": "function"}) assert "Invalid tool choice" in str(exc_info.value) def test_validate_tool_choice_invalid_type(): """Test that invalid types raise exceptions.""" - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match="\\. Expecting str, or dict\\. Please ensure") as exc_info: validate_chat_completion_tool_choice(123) assert "Got=" in str(exc_info.value) - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match="Invalid tool choice, tool_choice=\\[\\]\\. Got=\\.") as exc_info: validate_chat_completion_tool_choice([]) assert "Got=" in str(exc_info.value) diff --git a/tests/llm_responses_api_testing/base_responses_api.py b/tests/llm_responses_api_testing/base_responses_api.py index f5751aa79e8..d5057944ba7 100644 --- a/tests/llm_responses_api_testing/base_responses_api.py +++ b/tests/llm_responses_api_testing/base_responses_api.py @@ -28,6 +28,7 @@ from openai.types.responses.response_create_params import ( ResponseInputParam, ) from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler +import openai def validate_responses_api_response(response, final_chunk: bool = False): @@ -700,12 +701,12 @@ class BaseResponsesAPITest(ABC): base_completion_call_args = self.get_base_completion_call_args() if sync_mode: - with pytest.raises(Exception): + with pytest.raises(openai.APIError): litellm.cancel_responses( response_id="invalid_response_id_12345", **base_completion_call_args ) else: - with pytest.raises(Exception): + with pytest.raises(openai.APIError): await litellm.acancel_responses( response_id="invalid_response_id_12345", **base_completion_call_args ) diff --git a/tests/llm_responses_api_testing/test_openai_responses_api.py b/tests/llm_responses_api_testing/test_openai_responses_api.py index bd1517dbffb..d19fa09451c 100644 --- a/tests/llm_responses_api_testing/test_openai_responses_api.py +++ b/tests/llm_responses_api_testing/test_openai_responses_api.py @@ -1643,10 +1643,13 @@ async def test_openai_responses_api_token_limit_error(): model="gpt-5-mini", input=oversized_text, stream=True ) - with pytest.raises(litellm.APIError) as exc_info: + async def _drain(): async for event in response: print(event) + with pytest.raises(litellm.APIError) as exc_info: + await _drain() + assert exc_info.value.status_code == 400 assert "exceeds the context window" in str(exc_info.value) diff --git a/tests/llm_responses_api_testing/test_responses_hooks.py b/tests/llm_responses_api_testing/test_responses_hooks.py index 2344a62de4d..66dbb29dba5 100644 --- a/tests/llm_responses_api_testing/test_responses_hooks.py +++ b/tests/llm_responses_api_testing/test_responses_hooks.py @@ -295,7 +295,7 @@ async def test_responses_streaming_failure_triggers_failure_handlers(): call_type=CallTypes.responses.value, ) - with pytest.raises(ValueError): + with pytest.raises(ValueError, match="boom"): iterator._process_chunk('{"delta": "chunk"}') # allow failure callbacks to run diff --git a/tests/llm_translation/test_azure_openai.py b/tests/llm_translation/test_azure_openai.py index 4f12e12700d..eb5ba44c410 100644 --- a/tests/llm_translation/test_azure_openai.py +++ b/tests/llm_translation/test_azure_openai.py @@ -650,7 +650,7 @@ def test_azure_openai_responses_bridge(): mock_responses.assert_called_once() assert ( mock_responses.call_args.kwargs["model"] - == "test-azure-computer-use-preview" + == "azure/test-azure-computer-use-preview" ) assert mock_responses.call_args.kwargs["custom_llm_provider"] == "azure" diff --git a/tests/llm_translation/test_bedrock_completion.py b/tests/llm_translation/test_bedrock_completion.py index 3303fafafb0..9534bc8de3c 100644 --- a/tests/llm_translation/test_bedrock_completion.py +++ b/tests/llm_translation/test_bedrock_completion.py @@ -1890,7 +1890,7 @@ def test_bedrock_completion_test_4(modify_params): ] assert transformed_messages == expected_messages else: - with pytest.raises(Exception) as e: + with pytest.raises(Exception, match=r"litellm\.modify_params") as e: litellm.completion(**data) assert "litellm.modify_params" in str(e.value) diff --git a/tests/llm_translation/test_bedrock_moonshot.py b/tests/llm_translation/test_bedrock_moonshot.py index a9f4a86b3b6..a82d1c6f029 100644 --- a/tests/llm_translation/test_bedrock_moonshot.py +++ b/tests/llm_translation/test_bedrock_moonshot.py @@ -12,6 +12,7 @@ This test suite verifies: """ from base_llm_unit_tests import BaseLLMChatTest +import httpx import pytest import sys import os @@ -208,14 +209,6 @@ class TestBedrockMoonshotInvoke(BaseLLMChatTest): endpoint with the messages body. Iteration of the stream itself is not exercised here — moonshot streaming delegates to the OpenAI parser and is covered by the OpenAI test suite. - - Note: bedrock invoke streaming cannot be intercepted by patching - the caller-supplied client, because ``CustomStreamWrapper.fetch_sync_stream`` - at streaming_handler.py invokes the stored ``make_call`` partial with - ``client=litellm.module_level_client``, which overrides any client the - caller passed. Patch ``make_sync_call`` at its import site in - ``base_invoke_transformation`` so we observe the exact kwargs the - partial was built with at stream-wrapper construction time. """ from litellm.utils import CustomStreamWrapper @@ -225,7 +218,7 @@ class TestBedrockMoonshotInvoke(BaseLLMChatTest): captured.update(kwargs) # Return an empty iterator so the stream wrapper's iteration # doesn't try to parse real bytes. - return iter([]) + return iter([]), httpx.Headers() with patch( "litellm.llms.bedrock.chat.invoke_transformations." @@ -246,11 +239,6 @@ class TestBedrockMoonshotInvoke(BaseLLMChatTest): aws_region_name="us-west-2", ) assert isinstance(response, CustomStreamWrapper) - # Trigger fetch_sync_stream → make_call(...) → fake_make_sync_call. - try: - next(iter(response)) - except StopIteration: - pass assert captured, "make_sync_call was never invoked" assert captured["api_base"].endswith("/invoke-with-response-stream") diff --git a/tests/llm_translation/test_llm_response_utils/test_convert_dict_to_chat_completion.py b/tests/llm_translation/test_llm_response_utils/test_convert_dict_to_chat_completion.py index 5683d973ac9..8c7390d3d04 100644 --- a/tests/llm_translation/test_llm_response_utils/test_convert_dict_to_chat_completion.py +++ b/tests/llm_translation/test_llm_response_utils/test_convert_dict_to_chat_completion.py @@ -982,7 +982,7 @@ def test_convert_to_model_response_object_with_real_error(): }, } - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception) as exc_info: # noqa: PT011 # message rides on .message, str() is empty convert_to_model_response_object( model_response_object=ModelResponse(), response_object=response_object, @@ -1243,7 +1243,7 @@ def test_convert_to_model_response_object_with_error_code_only(): }, } - with pytest.raises(Exception): + with pytest.raises(Exception) as exc_info: # noqa: B017, PT011 # bare Exception, empty message, so status_code is the assertion convert_to_model_response_object( model_response_object=ModelResponse(), response_object=response_object, @@ -1255,6 +1255,8 @@ def test_convert_to_model_response_object_with_error_code_only(): convert_tool_call_to_json_mode=False, ) + assert exc_info.value.status_code == 500 + def test_model_prefix_preservation(): """ @@ -1421,7 +1423,7 @@ def test_error_message_includes_function_args(): "choices": [{"index": 0}], } - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='in convert_to_model_response_object') as exc_info: convert_to_model_response_object( model_response_object=ModelResponse(), response_object=response_object, @@ -2473,14 +2475,14 @@ class TestConvertToModelResponseObjectCompletion: assert "reasoning_content" not in (message.provider_specific_fields or {}) def test_response_none_raises(self): - with pytest.raises(Exception): + with pytest.raises(Exception, match="Invalid response object"): convert_to_model_response_object( response_object=None, model_response_object=ModelResponse(), ) def test_model_response_none_raises(self): - with pytest.raises(Exception): + with pytest.raises(Exception, match="Invalid response object"): convert_to_model_response_object( response_object={ "choices": [ diff --git a/tests/llm_translation/test_openai.py b/tests/llm_translation/test_openai.py index 10ed5f1ef68..405dbb0e6ec 100644 --- a/tests/llm_translation/test_openai.py +++ b/tests/llm_translation/test_openai.py @@ -1458,7 +1458,7 @@ def test_responses_gpt54_with_xhigh_reasoning(): # Stop execution right after request generation to avoid external API calls. mock_responses.side_effect = RuntimeError("stop_after_request_build") - with pytest.raises(Exception): + with pytest.raises(litellm.APIConnectionError): litellm.completion( model="openai/responses/gpt-5.4", messages=[{"role": "user", "content": "What is 2+2?"}], diff --git a/tests/llm_translation/test_prompt_factory.py b/tests/llm_translation/test_prompt_factory.py index ae215602e31..1b4c8a82cf4 100644 --- a/tests/llm_translation/test_prompt_factory.py +++ b/tests/llm_translation/test_prompt_factory.py @@ -1288,7 +1288,8 @@ def test_just_system_message(): model="anthropic.claude-3-sonnet-20240229-v1:0", llm_provider="bedrock", ) - assert "bedrock requires at least one non-system message" in str(e.value) + + assert "bedrock requires at least one non-system message" in str(e.value) def test_convert_generic_image_chunk_to_openai_image_obj(): @@ -1844,7 +1845,7 @@ def test_parse_tool_call_arguments_malformed_json(): parse_tool_call_arguments, ) - with pytest.raises(ValueError) as exc_info: + with pytest.raises(ValueError, match="Failed to parse tool call arguments for tool 'load_skill") as exc_info: parse_tool_call_arguments( '{"skill_name": "pptx', tool_name="load_skill", @@ -1876,7 +1877,7 @@ def test_convert_to_anthropic_tool_invoke_malformed_json(): } ] - with pytest.raises(ValueError) as exc_info: + with pytest.raises(ValueError, match="Failed to parse tool call arguments for tool 'bad_tool") as exc_info: convert_to_anthropic_tool_invoke(tool_calls) error_msg = str(exc_info.value) @@ -2022,7 +2023,7 @@ def test_parse_tool_call_arguments_still_raises_for_unrepairable(): parse_tool_call_arguments, ) - with pytest.raises(ValueError) as exc_info: + with pytest.raises(ValueError, match="Failed to parse tool call arguments for tool 'test_tool") as exc_info: parse_tool_call_arguments( '{"key": "unterminated', tool_name="test_tool", diff --git a/tests/llm_translation/test_together_ai.py b/tests/llm_translation/test_together_ai.py index 4ad0c90230d..387e61656ea 100644 --- a/tests/llm_translation/test_together_ai.py +++ b/tests/llm_translation/test_together_ai.py @@ -20,7 +20,7 @@ import pytest class TestTogetherAI(BaseLLMChatTest): def get_base_completion_call_args(self) -> dict: litellm.set_verbose = True - return {"model": "together_ai/Qwen/Qwen2.5-7B-Instruct-Turbo"} + return {"model": "together_ai/openai/gpt-oss-20b"} def test_tool_call_no_arguments(self, tool_call_no_arguments): """Test that tool calls with no arguments is translated correctly. Relevant issue: https://github.com/BerriAI/litellm/issues/6833""" diff --git a/tests/llm_translation/test_triton.py b/tests/llm_translation/test_triton.py index 21887e8d848..f4a26360a6c 100644 --- a/tests/llm_translation/test_triton.py +++ b/tests/llm_translation/test_triton.py @@ -45,7 +45,7 @@ def test_split_embedding_by_shape_fails_with_shape_value_error(): "data": [1, 2, 3, 4, 5, 6], } ] - with pytest.raises(ValueError): + with pytest.raises(ValueError, match='Shape must be of length'): TritonEmbeddingConfig.split_embedding_by_shape( data[0]["data"], data[0]["shape"] ) diff --git a/tests/llm_translation/test_unit_test_bedrock_invoke.py b/tests/llm_translation/test_unit_test_bedrock_invoke.py index 14f08c759c5..39f02263f03 100644 --- a/tests/llm_translation/test_unit_test_bedrock_invoke.py +++ b/tests/llm_translation/test_unit_test_bedrock_invoke.py @@ -59,7 +59,7 @@ def test_transform_request_invalid_provider(bedrock_transformer): """Test request transformation with invalid provider""" messages = [{"role": "user", "content": "Hello"}] - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Bedrock Invoke HTTPX: Unknown provider=None') as exc_info: bedrock_transformer.transform_request( model="invalid.model", messages=messages, diff --git a/tests/local_testing/test_aim_guardrails.py b/tests/local_testing/test_aim_guardrails.py index 2cb7f9cd357..5e5fb0d5459 100644 --- a/tests/local_testing/test_aim_guardrails.py +++ b/tests/local_testing/test_aim_guardrails.py @@ -101,26 +101,26 @@ async def test_block_callback(mode: str): ], } - with pytest.raises(ProxyException, match="Jailbreak detected") as exc_info: - with patch( - "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", - return_value=Response( - json={ - "analysis_result": { - "analysis_time_ms": 212, - "policy_drill_down": {}, - "session_entities": [], - }, - "required_action": { - "action_type": "block_action", - "detection_message": "Jailbreak detected", - "policy_name": "blocking policy", - }, + with patch( + "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", + return_value=Response( + json={ + "analysis_result": { + "analysis_time_ms": 212, + "policy_drill_down": {}, + "session_entities": [], }, - status_code=200, - request=Request(method="POST", url="http://aim"), - ), - ): + "required_action": { + "action_type": "block_action", + "detection_message": "Jailbreak detected", + "policy_name": "blocking policy", + }, + }, + status_code=200, + request=Request(method="POST", url="http://aim"), + ), + ): + async def _call_guardrail(): if mode == "pre_call": await aim_guardrail.async_pre_call_hook( data=data, @@ -135,6 +135,9 @@ async def test_block_callback(mode: str): call_type="completion", ) + with pytest.raises(ProxyException, match="Jailbreak detected") as exc_info: + await _call_guardrail() + exc = exc_info.value assert exc.code == "400" assert exc.type == "invalid_request_error" diff --git a/tests/local_testing/test_auth_utils.py b/tests/local_testing/test_auth_utils.py index 9aecb7e10e4..88e8c02a606 100644 --- a/tests/local_testing/test_auth_utils.py +++ b/tests/local_testing/test_auth_utils.py @@ -264,10 +264,6 @@ def test_get_end_user_id_from_request_body_backwards_compatibility(): ["gpt-3.5-turbo", "gpt-4o-mini-general-deployment"], ), ({"model": "gpt-3.5-turbo"}, "gpt-3.5-turbo"), - ( - {"model": "gpt-3.5-turbo, gpt-4o-mini-general-deployment"}, - ["gpt-3.5-turbo", "gpt-4o-mini-general-deployment"], - ), ], ) def test_get_model_from_request(request_data, expected_model): diff --git a/tests/local_testing/test_completion.py b/tests/local_testing/test_completion.py index 01fd35cb42d..5b0bff65959 100644 --- a/tests/local_testing/test_completion.py +++ b/tests/local_testing/test_completion.py @@ -67,7 +67,7 @@ def test_completion_custom_provider_model_name(): try: litellm.cache = None response = completion( - model="together_ai/Qwen/Qwen2.5-7B-Instruct-Turbo", + model="together_ai/openai/gpt-oss-20b", messages=messages, logger_fn=logger_fn, ) @@ -2817,7 +2817,7 @@ def test_customprompt_together_ai(): print(litellm.success_callback) print(litellm._async_success_callback) response = completion( - model="together_ai/Qwen/Qwen2.5-7B-Instruct-Turbo", + model="together_ai/openai/gpt-oss-20b", messages=messages, roles={ "system": { @@ -3657,7 +3657,7 @@ def test_completion_together_ai_stream(): messages = [{"content": user_message, "role": "user"}] try: response = completion( - model="together_ai/Qwen/Qwen2.5-7B-Instruct-Turbo", + model="together_ai/openai/gpt-oss-20b", messages=messages, stream=True, max_tokens=5, diff --git a/tests/local_testing/test_completion_cost.py b/tests/local_testing/test_completion_cost.py index e34d5c349c5..7dfcb55e29a 100644 --- a/tests/local_testing/test_completion_cost.py +++ b/tests/local_testing/test_completion_cost.py @@ -625,17 +625,9 @@ def test_vertex_ai_completion_cost(): print("calculated_input_cost: {}".format(calculated_input_cost)) -@pytest.mark.skip(reason="new test - WIP, working on fixing this") def test_vertex_ai_medlm_completion_cost(): """Test for medlm completion cost .""" - with pytest.raises(Exception) as e: - model = "vertex_ai/medlm-medium" - messages = [{"role": "user", "content": "Test MedLM completion cost."}] - predictive_cost = completion_cost( - model=model, messages=messages, custom_llm_provider="vertex_ai" - ) - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" litellm.model_cost = litellm.get_model_cost_map(url="") diff --git a/tests/local_testing/test_exceptions.py b/tests/local_testing/test_exceptions.py index 8c1df52e28e..8dd90cbfb37 100644 --- a/tests/local_testing/test_exceptions.py +++ b/tests/local_testing/test_exceptions.py @@ -1417,7 +1417,7 @@ async def test_exception_bubbling_up(sync_mode, stream_mode, model): import litellm litellm.set_verbose = True - with pytest.raises(Exception) as exc_info: + async def _call_with_bad_role(): if sync_mode: litellm.completion( model=model, @@ -1433,6 +1433,9 @@ async def test_exception_bubbling_up(sync_mode, stream_mode, model): sync_stream=sync_mode, ) + with pytest.raises(Exception, match='litellm\\.BadRequestError: OpenAIException - Invalid value') as exc_info: + await _call_with_bad_role() + assert exc_info.value.code == "invalid_value" assert exc_info.value.param is not None assert exc_info.value.type == "invalid_request_error" diff --git a/tests/local_testing/test_file_types.py b/tests/local_testing/test_file_types.py index db83ba0e74b..7fda81ebd45 100644 --- a/tests/local_testing/test_file_types.py +++ b/tests/local_testing/test_file_types.py @@ -23,13 +23,13 @@ class TestFileConsts: def test_get_file_extension_from_mime_type(self): assert get_file_extension_from_mime_type("audio/aac") == "aac" assert get_file_extension_from_mime_type("application/pdf") == "pdf" - with pytest.raises(ValueError): + with pytest.raises(ValueError, match='Unknown extension for mime type: application'): get_file_extension_from_mime_type("application/unknown") def test_get_file_type_from_extension(self): assert get_file_type_from_extension("aac") == FileType.AAC assert get_file_type_from_extension("pdf") == FileType.PDF - with pytest.raises(ValueError): + with pytest.raises(ValueError, match='Unknown file type for extension: unknown'): get_file_type_from_extension("unknown") def test_get_file_extension_for_file_type(self): diff --git a/tests/local_testing/test_function_calling.py b/tests/local_testing/test_function_calling.py index 4095962f91d..d6adde84400 100644 --- a/tests/local_testing/test_function_calling.py +++ b/tests/local_testing/test_function_calling.py @@ -357,14 +357,13 @@ def test_parallel_function_call_anthropic_error_msg( if expect_unsupported_params_error: with pytest.raises(litellm.UnsupportedParamsError) as e: - second_response = litellm.completion( + litellm.completion( model=model, messages=messages, temperature=0.2, seed=22, drop_params=True, - ) # get a new response from the model where it can see the function response - print("second response\n", second_response) + ) else: second_response = litellm.completion( model=model, diff --git a/tests/local_testing/test_get_llm_provider.py b/tests/local_testing/test_get_llm_provider.py index 4c3e13da17a..0e667b82a66 100644 --- a/tests/local_testing/test_get_llm_provider.py +++ b/tests/local_testing/test_get_llm_provider.py @@ -569,5 +569,5 @@ class TestClaudeModelPatternMatching: ) set_fallback_generalizations([]) - with pytest.raises(Exception): + with pytest.raises(litellm.BadRequestError): litellm.get_llm_provider(model="claude-opus-4-9") diff --git a/tests/local_testing/test_get_model_info.py b/tests/local_testing/test_get_model_info.py index 385be25fb07..cef05050ac9 100644 --- a/tests/local_testing/test_get_model_info.py +++ b/tests/local_testing/test_get_model_info.py @@ -134,7 +134,6 @@ def test_get_model_info_bedrock_region(): "ft:gpt-3.5-turbo:my-org:custom_suffix:id", "ft:gpt-4-0613:my-org:custom_suffix:id", "ft:davinci-002:my-org:custom_suffix:id", - "ft:gpt-4-0613:my-org:custom_suffix:id", "ft:babbage-002:my-org:custom_suffix:id", "gpt-35-turbo", "ada", diff --git a/tests/local_testing/test_mock_request.py b/tests/local_testing/test_mock_request.py index 710024b61b1..c9cd14633ba 100644 --- a/tests/local_testing/test_mock_request.py +++ b/tests/local_testing/test_mock_request.py @@ -128,13 +128,12 @@ def test_router_mock_request_with_mock_timeout(): ], ) with pytest.raises(litellm.Timeout): - response = router.completion( + router.completion( model="gpt-3.5-turbo", messages=[{"role": "user", "content": "Hey, I'm a mock request"}], timeout=3, mock_timeout=True, ) - print(response) end_time = time.time() assert end_time - start_time >= 3, f"Time taken: {end_time - start_time}" diff --git a/tests/local_testing/test_router_budget_limiter.py b/tests/local_testing/test_router_budget_limiter.py index 1a36e9de8f2..4ef99ec8c12 100644 --- a/tests/local_testing/test_router_budget_limiter.py +++ b/tests/local_testing/test_router_budget_limiter.py @@ -160,13 +160,11 @@ async def test_provider_budgets_e2e_test_expect_to_fail(): await asyncio.sleep(2.5) for _ in range(3): - with pytest.raises(Exception) as exc_info: - response = await router.acompletion( + with pytest.raises(Exception, match="Exceeded budget for provider") as exc_info: + await router.acompletion( messages=[{"role": "user", "content": "Hello, how are you?"}], model="anthropic/claude-sonnet-4-5-20250929", ) - print(response) - print("response.hidden_params", response._hidden_params) await asyncio.sleep(0.5) # Verify the error is related to budget exceeded @@ -596,13 +594,11 @@ async def test_deployment_budgets_e2e_test_expect_to_fail(): await asyncio.sleep(2.5) for _ in range(3): - with pytest.raises(Exception) as exc_info: - response = await router.acompletion( + with pytest.raises(Exception, match="Exceeded budget for deployment") as exc_info: + await router.acompletion( messages=[{"role": "user", "content": "Hello, how are you?"}], model="openai/gpt-4o-mini", ) - print(response) - print("response.hidden_params", response._hidden_params) await asyncio.sleep(0.5) # Verify the error is related to budget exceeded @@ -650,14 +646,12 @@ async def test_tag_budgets_e2e_test_expect_to_fail(): await asyncio.sleep(2.5) for _ in range(3): - with pytest.raises(Exception) as exc_info: - response = await router.acompletion( + with pytest.raises(Exception, match=f"Exceeded budget for tag='{TAG_NAME}'") as exc_info: + await router.acompletion( messages=[{"role": "user", "content": "Hello, how are you?"}], model="openai/gpt-4o-mini", metadata={"tags": [TAG_NAME]}, ) - print(response) - print("response.hidden_params", response._hidden_params) await asyncio.sleep(0.5) # Verify the error is related to budget exceeded diff --git a/tests/local_testing/test_router_fallbacks.py b/tests/local_testing/test_router_fallbacks.py index 7c09c978029..86dec406332 100644 --- a/tests/local_testing/test_router_fallbacks.py +++ b/tests/local_testing/test_router_fallbacks.py @@ -1416,7 +1416,7 @@ async def test_router_fallbacks_default_and_model_specific_fallbacks(sync_mode): default_fallbacks=["bad-model"], ) - with pytest.raises(Exception) as exc_info: + async def _call_bad_model(): if sync_mode: resp = router.completion( model="bad-model", @@ -1429,6 +1429,9 @@ async def test_router_fallbacks_default_and_model_specific_fallbacks(sync_mode): model="bad-model", messages=[{"role": "user", "content": "Hey, how's it going?"}], ) + + with pytest.raises(Exception, match='litellm\\.AuthenticationError: AuthenticationError') as exc_info: + await _call_bad_model() assert isinstance( exc_info.value, litellm.AuthenticationError ), f"Expected AuthenticationError, but got {type(exc_info.value).__name__}" diff --git a/tests/local_testing/test_router_max_parallel_requests.py b/tests/local_testing/test_router_max_parallel_requests.py index 1b81b9eb999..7bb40dd7a2f 100644 --- a/tests/local_testing/test_router_max_parallel_requests.py +++ b/tests/local_testing/test_router_max_parallel_requests.py @@ -205,9 +205,12 @@ async def test_max_parallel_requests_tpm_rate_limiting_base_case(): num_retries=0, ) - with pytest.raises(litellm.RateLimitError): + async def _exceed_limit(): for _ in range(2): await router.acompletion( model="gpt-4o-2024-08-06", messages=_messages, ) + + with pytest.raises(litellm.RateLimitError): + await _exceed_limit() diff --git a/tests/local_testing/test_streaming.py b/tests/local_testing/test_streaming.py index a4f564b227f..1fe9a1ab297 100644 --- a/tests/local_testing/test_streaming.py +++ b/tests/local_testing/test_streaming.py @@ -2926,11 +2926,14 @@ def test_unit_test_custom_stream_wrapper_repeating_chunk( print(f"expected_chunk_fail: {expected_chunk_fail}") if (loop_amount > litellm.REPEATED_STREAMING_CHUNK_LIMIT) and expected_chunk_fail: + def _drain(): + for chunk in response: + continue + with pytest.raises( (litellm.InternalServerError, litellm.exceptions.MidStreamFallbackError) ): - for chunk in response: - continue + _drain() else: for chunk in response: continue diff --git a/tests/local_testing/test_text_completion.py b/tests/local_testing/test_text_completion.py index b22988a468e..63cee71f999 100644 --- a/tests/local_testing/test_text_completion.py +++ b/tests/local_testing/test_text_completion.py @@ -4036,7 +4036,7 @@ def test_async_text_completion_together_ai(): async def test_get_response(): try: response = await litellm.atext_completion( - model="together_ai/Qwen/Qwen2.5-7B-Instruct-Turbo", + model="together_ai/openai/gpt-oss-20b", prompt="good morning", max_tokens=10, ) diff --git a/tests/logging_callback_tests/test_standard_logging_payload.py b/tests/logging_callback_tests/test_standard_logging_payload.py index d13cdf1337a..6a632c32fc2 100644 --- a/tests/logging_callback_tests/test_standard_logging_payload.py +++ b/tests/logging_callback_tests/test_standard_logging_payload.py @@ -293,7 +293,7 @@ def test_cleanup_timestamps(): assert all(isinstance(x, float) for x in result) # Test invalid input - with pytest.raises(ValueError): + with pytest.raises(ValueError, match="start_time is required, got=invalid of type "): StandardLoggingPayloadSetup.cleanup_timestamps( "invalid", end_float, completion_float ) diff --git a/tests/multi_instance_e2e_tests/test_update_team_e2e.py b/tests/multi_instance_e2e_tests/test_update_team_e2e.py index dfbfbd310ee..13091fd3df6 100644 --- a/tests/multi_instance_e2e_tests/test_update_team_e2e.py +++ b/tests/multi_instance_e2e_tests/test_update_team_e2e.py @@ -143,7 +143,7 @@ async def test_team_blocking_behavior_multi_instance(): assert team_info_4001["blocked"] is True, "Team should be blocked after update" # 8. Make a chat completion request on port 4000 with a new prompt; expect it to be blocked. - with pytest.raises(Exception) as excinfo: + with pytest.raises(Exception, match="(?i)blocked") as excinfo: await chat_completion_on_port( session, key=key, @@ -157,7 +157,7 @@ async def test_team_blocking_behavior_multi_instance(): ), f"Expected error indicating team blocked, got: {error_msg}" # 9. Make a chat completion request on port 4000 with a new prompt; expect it to be blocked. - with pytest.raises(Exception) as excinfo: + with pytest.raises(Exception, match="(?i)blocked") as excinfo: await chat_completion_on_port( session, key=key, @@ -171,7 +171,7 @@ async def test_team_blocking_behavior_multi_instance(): ), f"Expected error indicating team blocked, got: {error_msg}" # 9. Repeat the chat completion request with another new prompt; expect it to be blocked. - with pytest.raises(Exception) as excinfo_second: + with pytest.raises(Exception, match="(?i)blocked") as excinfo_second: await chat_completion_on_port( session, key=key, diff --git a/tests/ocr_tests/test_ocr_azure_document_intelligence.py b/tests/ocr_tests/test_ocr_azure_document_intelligence.py index 5736bd797e3..e6a2e5e5735 100644 --- a/tests/ocr_tests/test_ocr_azure_document_intelligence.py +++ b/tests/ocr_tests/test_ocr_azure_document_intelligence.py @@ -101,7 +101,7 @@ class TestAzureDocumentIntelligencePagesParam: cfg.map_ocr_params({"pages": [True, False]}, {}, "prebuilt-layout") def test_map_ocr_params_unsupported_type_raises(self, cfg): - with pytest.raises(ValueError): + with pytest.raises(ValueError, match='based, Mistral-style\\) or a string like'): cfg.map_ocr_params({"pages": 5}, {}, "prebuilt-layout") def test_get_complete_url_appends_pages_query(self, cfg): diff --git a/tests/openai_endpoints_tests/test_e2e_openai_responses_api.py b/tests/openai_endpoints_tests/test_e2e_openai_responses_api.py index 220a44f0792..be565972b94 100644 --- a/tests/openai_endpoints_tests/test_e2e_openai_responses_api.py +++ b/tests/openai_endpoints_tests/test_e2e_openai_responses_api.py @@ -1,5 +1,5 @@ import httpx -from openai import OpenAI, BadRequestError +from openai import OpenAI, BadRequestError, APIStatusError import pytest @@ -87,7 +87,7 @@ def test_basic_response(): print("DELETE response=", delete_response) # expect an error when getting the response again since it was deleted - with pytest.raises(Exception): + with pytest.raises(APIStatusError): get_response = client.responses.retrieve(response.id) @@ -195,6 +195,6 @@ def test_cancel_streaming_response(): def test_cancel_invalid_response_id(): client = get_test_client() - with pytest.raises(Exception): + with pytest.raises(APIStatusError): # Try to cancel a non-existent response ID client.responses.cancel("invalid_response_id_12345") diff --git a/tests/otel_tests/test_e2e_model_access.py b/tests/otel_tests/test_e2e_model_access.py index 5b5f2a89c8d..e5e93c0b179 100644 --- a/tests/otel_tests/test_e2e_model_access.py +++ b/tests/otel_tests/test_e2e_model_access.py @@ -3,6 +3,7 @@ import asyncio import aiohttp import json from httpx import AsyncClient +from openai import PermissionDeniedError from typing import Any, Optional, List, Literal @@ -134,7 +135,7 @@ async def test_model_access_update(): await mock_chat_completion(session=session, key=key, model="openai/gpt-5.5") # Should fail with gpt-5-mini - with pytest.raises(Exception) as exc_info: + with pytest.raises(PermissionDeniedError) as exc_info: await mock_chat_completion( session=session, key=key, model="openai/gpt-5-mini" ) @@ -157,7 +158,7 @@ async def test_model_access_update(): ) # Non-OpenAI model should still fail - with pytest.raises(Exception) as exc_info: + with pytest.raises(PermissionDeniedError) as exc_info: await mock_chat_completion( session=session, key=key, model="anthropic/claude-2" ) @@ -254,7 +255,7 @@ async def test_team_model_access_update(): await mock_chat_completion(session=session, key=key, model="openai/gpt-5.5") # Should fail with gpt-5-mini - with pytest.raises(Exception) as exc_info: + with pytest.raises(PermissionDeniedError) as exc_info: await mock_chat_completion( session=session, key=key, model="openai/gpt-5-mini" ) @@ -279,7 +280,7 @@ async def test_team_model_access_update(): ) # Non-OpenAI model should still fail - with pytest.raises(Exception) as exc_info: + with pytest.raises(PermissionDeniedError) as exc_info: await mock_chat_completion( session=session, key=key, model="anthropic/claude-2" ) diff --git a/tests/proxy_admin_ui_tests/test_access_group_team_sync.py b/tests/proxy_admin_ui_tests/test_access_group_team_sync.py index 629d77f20fc..f7092d3ec00 100644 --- a/tests/proxy_admin_ui_tests/test_access_group_team_sync.py +++ b/tests/proxy_admin_ui_tests/test_access_group_team_sync.py @@ -170,12 +170,15 @@ async def test_a_failed_mirror_takes_the_new_team_row_with_it(): async with _clean_db() as db: await _seed(db, {GROUPS[0]: [], GROUPS[1]: [OTHER_TEAM]}) - with pytest.raises(RuntimeError): + async def _blow_up_after_reconcile(): async with db.tx() as tx: await tx.litellm_teamtable.create(data={"team_id": TEAM, "access_group_ids": [GROUPS[0]]}) await reconcile_team_access_group_membership(tx, TEAM) raise RuntimeError("the cache handoff blew up") + with pytest.raises(RuntimeError): + await _blow_up_after_reconcile() + assert await _read(db) == {GROUPS[0]: [], GROUPS[1]: [OTHER_TEAM]} assert await db.litellm_teamtable.find_unique(where={"team_id": TEAM}) is None diff --git a/tests/proxy_admin_ui_tests/test_key_management.py b/tests/proxy_admin_ui_tests/test_key_management.py index 4c5a045509a..7e8494b77fc 100644 --- a/tests/proxy_admin_ui_tests/test_key_management.py +++ b/tests/proxy_admin_ui_tests/test_key_management.py @@ -1340,6 +1340,6 @@ async def test_team_model_alias(prisma_client, requested_model, should_pass): }, "Expected model aliases to be present" else: # Verify the key fails with non-aliased models - with pytest.raises(Exception) as exc_info: + with pytest.raises(ProxyException) as exc_info: await user_api_key_auth(request=request, api_key=f"Bearer {generated_key}") assert exc_info.value.type == ProxyErrorTypes.key_model_access_denied diff --git a/tests/proxy_admin_ui_tests/test_role_based_access.py b/tests/proxy_admin_ui_tests/test_role_based_access.py index 9398428bd67..f9506fb694b 100644 --- a/tests/proxy_admin_ui_tests/test_role_based_access.py +++ b/tests/proxy_admin_ui_tests/test_role_based_access.py @@ -9,7 +9,7 @@ from litellm._uuid import uuid from datetime import datetime from dotenv import load_dotenv -from fastapi import Request +from fastapi import HTTPException, Request from fastapi.routing import APIRoute load_dotenv() @@ -530,7 +530,7 @@ async def test_user_role_permissions(prisma_client, route, user_role, expected_r print(f"Auth passed as expected for {route} with role {user_role}") else: # Should raise an error - with pytest.raises(Exception) as exc_info: + with pytest.raises((ProxyException, HTTPException)) as exc_info: await user_api_key_auth(request=request, api_key=bearer_token) print(f"Auth failed as expected for {route} with role {user_role}") print(f"Error message: {str(exc_info.value)}") diff --git a/tests/proxy_unit_tests/test_auth_checks.py b/tests/proxy_unit_tests/test_auth_checks.py index e58e6c9694b..ef3cbd0ae95 100644 --- a/tests/proxy_unit_tests/test_auth_checks.py +++ b/tests/proxy_unit_tests/test_auth_checks.py @@ -173,7 +173,7 @@ async def test_can_key_call_model(model, expect_to_work): if expect_to_work: await can_key_call_model(**args) else: - with pytest.raises(Exception) as e: + with pytest.raises(Exception, match='key not allowed to access model\\. This key can only access') as e: await can_key_call_model(**args) print(e) @@ -242,8 +242,8 @@ async def test_can_team_call_model(model, expect_to_work): ) @pytest.mark.asyncio async def test_can_key_call_model_wildcard_access(key_models, model, expect_to_work): + from litellm.proxy._types import ProxyException from litellm.proxy.auth.auth_checks import can_key_call_model - from fastapi import HTTPException llm_model_list = [ { @@ -294,7 +294,7 @@ async def test_can_key_call_model_wildcard_access(key_models, model, expect_to_w llm_router=router, ) else: - with pytest.raises(Exception) as e: + with pytest.raises(ProxyException): await can_key_call_model( model=model, llm_model_list=llm_model_list, @@ -302,8 +302,6 @@ async def test_can_key_call_model_wildcard_access(key_models, model, expect_to_w llm_router=router, ) - print(e) - @pytest.mark.parametrize( "key_models, model, expect_to_work", @@ -330,6 +328,7 @@ async def test_wildcard_access_after_cost_map_reload(key_models, model, expect_t Fix: each reload now calls litellm.add_known_models(model_cost_map=new_map) with the fetched map passed explicitly to avoid any reference ambiguity. """ + from litellm.proxy._types import ProxyException from litellm.proxy.auth.auth_checks import can_key_call_model # Build a new cost map that includes the brand-new model — exactly what @@ -378,7 +377,7 @@ async def test_wildcard_access_after_cost_map_reload(key_models, model, expect_t llm_router=router, ) else: - with pytest.raises(Exception): + with pytest.raises(ProxyException): await can_key_call_model( model=model, llm_model_list=llm_model_list, @@ -959,7 +958,7 @@ async def test_can_key_call_model_with_aliases(model, alias_map, expect_to_work) llm_router=router, ) else: - with pytest.raises(Exception) as e: + with pytest.raises(Exception, match='key not allowed to access model\\. This key can only access') as e: await can_key_call_model( model=model, llm_model_list=llm_model_list, diff --git a/tests/proxy_unit_tests/test_jwt.py b/tests/proxy_unit_tests/test_jwt.py index beaa120dcb9..686d7021257 100644 --- a/tests/proxy_unit_tests/test_jwt.py +++ b/tests/proxy_unit_tests/test_jwt.py @@ -41,6 +41,7 @@ from litellm.proxy.auth.handle_jwt import JWTHandler, JWTAuthManager from litellm.proxy.management_endpoints.team_endpoints import new_team from litellm.proxy.proxy_server import chat_completion from typing import Literal, Optional +from litellm.proxy._types import ProxyException public_key = { "kty": "RSA", @@ -1045,11 +1046,8 @@ async def test_allow_access_by_email( assert result is not None # Adjust this based on your actual response check else: # Expect the call to fail - with pytest.raises( - Exception - ): # Replace with the actual exception raised on failure - resp = await user_api_key_auth(request=request, api_key=bearer_token) - print(resp) + with pytest.raises(ProxyException): + await user_api_key_auth(request=request, api_key=bearer_token) def test_get_public_key_from_jwk_url(): @@ -1585,7 +1583,7 @@ async def test_auth_jwt_mismatched_key_fails(monkeypatch): h = JWTHandler() with patch.object(h, "get_public_key", new=AsyncMock(return_value=rsa_jwk)): - with pytest.raises(Exception) as exc: + with pytest.raises(Exception, match='Validation fails: Expecting a PEM-formatted key\\.') as exc: await h.auth_jwt(token) assert "Validation fails" in str(exc.value) @@ -1828,7 +1826,7 @@ async def test_multi_issuer_jwt_unknown_issuer_without_global_jwks_rejected( kid="issuer-key", ) - with pytest.raises(Exception) as exc: + with pytest.raises(Exception, match='Missing JWT Public Key URL from environment\\.') as exc: await jwt_handler.auth_jwt(token=token) assert "Missing JWT Public Key URL" in str(exc.value) @@ -1859,7 +1857,7 @@ async def test_multi_issuer_jwt_rejects_wrong_audience(monkeypatch): kid="issuer-key", ) - with pytest.raises(Exception) as exc: + with pytest.raises(Exception, match="Validation fails: Audience doesn't match") as exc: await jwt_handler.auth_jwt(token=token) assert "Validation fails" in str(exc.value) @@ -1902,7 +1900,7 @@ async def test_multi_issuer_jwt_same_kid_does_not_cross_issuer_keys(monkeypatch) kid=shared_kid, ) - with pytest.raises(Exception) as exc: + with pytest.raises(Exception, match='Validation fails: Signature verification failed') as exc: await jwt_handler.auth_jwt(token=token) assert "Validation fails" in str(exc.value) @@ -1955,7 +1953,7 @@ def test_multi_issuer_jwt_requires_audience_unless_explicitly_disabled( issuer = "https://issuer.example.com" jwks_url = f"{issuer}/keys" - with pytest.raises(Exception) as exc: + with pytest.raises(Exception, match='must configure audience or set') as exc: LiteLLM_JWTAuth( issuers=[ { diff --git a/tests/proxy_unit_tests/test_proxy_config_unit_test.py b/tests/proxy_unit_tests/test_proxy_config_unit_test.py index e6b38f31b48..99b0dc4fd13 100644 --- a/tests/proxy_unit_tests/test_proxy_config_unit_test.py +++ b/tests/proxy_unit_tests/test_proxy_config_unit_test.py @@ -53,7 +53,7 @@ async def test_read_config_from_bad_file_path(): """ proxy_config_instance = ProxyConfig() config_path = "non-existent-file.yaml" - with pytest.raises(Exception): + with pytest.raises(Exception, match="Config file not found"): config = await proxy_config_instance.get_config(config_file_path=config_path) diff --git a/tests/proxy_unit_tests/test_proxy_server.py b/tests/proxy_unit_tests/test_proxy_server.py index bfbc92adc74..04bc80bf0d6 100644 --- a/tests/proxy_unit_tests/test_proxy_server.py +++ b/tests/proxy_unit_tests/test_proxy_server.py @@ -2412,12 +2412,14 @@ async def test_proxy_server_prisma_setup(): @pytest.mark.asyncio -async def test_proxy_server_prisma_setup_invalid_db(): +async def test_proxy_server_prisma_setup_invalid_db(monkeypatch): """ PROD TEST: Test that proxy server startup fails when it's unable to connect to the database Think 2-3 times before editing / deleting this test, it's important for PROD """ + import httpx + from litellm.proxy.proxy_server import ProxyStartupEvent from litellm.proxy.utils import ProxyLogging from litellm.caching import DualCache @@ -2425,24 +2427,14 @@ async def test_proxy_server_prisma_setup_invalid_db(): user_api_key_cache = DualCache() invalid_db_url = "postgresql://invalid:invalid@localhost:5432/nonexistent" - _old_db_url = os.getenv("DATABASE_URL") - os.environ["DATABASE_URL"] = invalid_db_url + monkeypatch.setenv("DATABASE_URL", invalid_db_url) - with pytest.raises(Exception) as exc_info: + with pytest.raises(httpx.ConnectError): await ProxyStartupEvent._setup_prisma_client( database_url=invalid_db_url, proxy_logging_obj=ProxyLogging(user_api_key_cache=user_api_key_cache), user_api_key_cache=user_api_key_cache, ) - print("GOT EXCEPTION=", exc_info) - - assert "httpx.ConnectError" in str(exc_info.value) - - # # Verify the error message indicates a database connection issue - # assert any(x in str(exc_info.value).lower() for x in ["database", "connection", "authentication"]) - - if _old_db_url: - os.environ["DATABASE_URL"] = _old_db_url @pytest.mark.asyncio diff --git a/tests/proxy_unit_tests/test_proxy_utils.py b/tests/proxy_unit_tests/test_proxy_utils.py index ad852c16905..de2a9282300 100644 --- a/tests/proxy_unit_tests/test_proxy_utils.py +++ b/tests/proxy_unit_tests/test_proxy_utils.py @@ -29,6 +29,7 @@ from litellm.proxy.litellm_pre_call_utils import ( _get_dynamic_logging_metadata, add_litellm_data_to_request, ) +from pydantic import ValidationError pytestmark = pytest.mark.xdist_group("proxy_heavy") @@ -1025,7 +1026,7 @@ def test_enforced_params_check( from litellm.proxy.litellm_pre_call_utils import _enforced_params_check if expected_error: - with pytest.raises(ValueError): + with pytest.raises(ValueError, match='in request body\\. This is a required param'): _enforced_params_check( request_body=request_body, general_settings=general_settings, @@ -1695,13 +1696,13 @@ def test_update_key_request_validation(): """ from litellm.proxy._types import UpdateKeyRequest - with pytest.raises(Exception): + with pytest.raises(ValidationError): UpdateKeyRequest( key="test_key", temp_budget_increase=100, ) - with pytest.raises(Exception): + with pytest.raises(ValidationError): UpdateKeyRequest( key="test_key", temp_budget_expiry="2024-01-20T00:00:00Z", @@ -1848,7 +1849,7 @@ async def test_end_user_transactions_reset(): mock_client.db.tx = AsyncMock(side_effect=Exception("DB Error")) # Call function - should raise error - with pytest.raises(Exception): + with pytest.raises(TypeError): await ProxyUpdateSpend.update_end_user_spend( n_retry_times=0, prisma_client=mock_client, @@ -1878,7 +1879,7 @@ async def test_spend_logs_cleanup_after_error(): original_logs = mock_client.spend_log_transactions.copy() # Call function - should raise error - with pytest.raises(Exception): + with pytest.raises(TypeError): await ProxyUpdateSpend.update_spend_logs( n_retry_times=0, prisma_client=mock_client, @@ -2625,7 +2626,7 @@ async def test_during_call_hook_parallel_execution_with_error(): try: litellm.callbacks = [FailingGuardrail()] - with pytest.raises(ValueError) as exc_info: + with pytest.raises(ValueError, match='Guardrail violation detected!') as exc_info: await proxy_logging.during_call_hook( data={ "model": "gpt-4", diff --git a/tests/proxy_unit_tests/test_skills_db.py b/tests/proxy_unit_tests/test_skills_db.py index 5f420bc314a..9548e78d6ed 100644 --- a/tests/proxy_unit_tests/test_skills_db.py +++ b/tests/proxy_unit_tests/test_skills_db.py @@ -26,6 +26,7 @@ from litellm.proxy import proxy_server from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.utils import PrismaClient, ProxyLogging from litellm.types.utils import LlmProviders +import openai proxy_logging_obj = ProxyLogging(user_api_key_cache=DualCache()) @@ -254,7 +255,7 @@ async def test_delete_skill_sdk(prisma_client): assert result.type == "skill_deleted" # Verify skill no longer exists - with pytest.raises(Exception): + with pytest.raises(openai.APIError): await aget_skill( skill_id=created_skill.id, custom_llm_provider=LlmProviders.LITELLM_PROXY.value, diff --git a/tests/proxy_unit_tests/test_update_spend.py b/tests/proxy_unit_tests/test_update_spend.py index 0d1d6dcf3c6..6b8973fbad2 100644 --- a/tests/proxy_unit_tests/test_update_spend.py +++ b/tests/proxy_unit_tests/test_update_spend.py @@ -166,7 +166,7 @@ async def test_update_spend_logs_non_connection_error(): prisma_client.db.litellm_spendlogs.create_many = create_many_mock # Execute and verify it raises immediately without retrying - with pytest.raises(ValueError) as exc_info: + with pytest.raises(ValueError, match='Unexpected database error') as exc_info: await update_spend(prisma_client, None, proxy_logging_obj) # Verify error message diff --git a/tests/proxy_unit_tests/test_user_api_key_auth.py b/tests/proxy_unit_tests/test_user_api_key_auth.py index ccf710c5708..58dbe3ad370 100644 --- a/tests/proxy_unit_tests/test_user_api_key_auth.py +++ b/tests/proxy_unit_tests/test_user_api_key_auth.py @@ -436,7 +436,7 @@ def test_ui_token_route_access(route, user_role, should_be_allowed): ) assert result is True else: - with pytest.raises(Exception): + with pytest.raises(Exception, match="Only proxy admin can be used to generate"): _is_api_route_allowed( route=route, request=request, diff --git a/tests/router_unit_tests/test_router_helper_utils.py b/tests/router_unit_tests/test_router_helper_utils.py index c3db9e67f9c..f81578dbd99 100644 --- a/tests/router_unit_tests/test_router_helper_utils.py +++ b/tests/router_unit_tests/test_router_helper_utils.py @@ -90,7 +90,7 @@ def test_routing_strategy_init_invalid_strategy(model_list): router = Router(model_list=model_list) # Test common mistake: "simple" instead of "simple-shuffle" - with pytest.raises(ValueError) as exc_info: + with pytest.raises(ValueError, match="usage-based-routing', 'provider-budget-routing'\\]\\. Check") as exc_info: router.routing_strategy_init( routing_strategy="simple", routing_strategy_args={} ) @@ -106,7 +106,7 @@ def test_routing_strategy_init_invalid_strategy(model_list): assert "Router SDK" in error_msg # Test completely invalid strategy - with pytest.raises(ValueError) as exc_info: + with pytest.raises(ValueError, match="usage-based-routing', 'provider-budget-routing'\\]\\. Check") as exc_info: router.routing_strategy_init( routing_strategy="not-a-real-strategy", routing_strategy_args={} ) @@ -423,10 +423,11 @@ def test_get_timeout(model_list): def test_handle_mock_testing_fallbacks(model_list, fallback_kwarg, expected_error): """Test if the '_handle_mock_testing_fallbacks' function is working correctly""" router = Router(model_list=model_list) + data = { + fallback_kwarg: True, + } + with pytest.raises(expected_error): - data = { - fallback_kwarg: True, - } router._handle_mock_testing_fallbacks( kwargs=data, ) @@ -435,10 +436,11 @@ def test_handle_mock_testing_fallbacks(model_list, fallback_kwarg, expected_erro def test_handle_mock_testing_rate_limit_error(model_list): """Test if the '_handle_mock_testing_rate_limit_error' function is working correctly""" router = Router(model_list=model_list) + data = { + "mock_testing_rate_limit_error": True, + } + with pytest.raises(litellm.RateLimitError): - data = { - "mock_testing_rate_limit_error": True, - } router._handle_mock_testing_rate_limit_error( kwargs=data, ) diff --git a/tests/store_model_in_db_tests/test_mcp_servers.py b/tests/store_model_in_db_tests/test_mcp_servers.py index e9c26221580..735d5d71ad3 100644 --- a/tests/store_model_in_db_tests/test_mcp_servers.py +++ b/tests/store_model_in_db_tests/test_mcp_servers.py @@ -471,7 +471,7 @@ def test_validate_mcp_server_name_direct(): validate_mcp_server_name("valid name") # Test that invalid names with hyphens raise exceptions - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match="Server name cannot contain '-'\\. Use an alternative") as exc_info: validate_mcp_server_name("invalid-name") assert "cannot contain" in str(exc_info.value) diff --git a/tests/test_litellm/a2a_protocol/test_a2a_exception_mapping_utils.py b/tests/test_litellm/a2a_protocol/test_a2a_exception_mapping_utils.py index 06191d1a370..c31d50960b1 100644 --- a/tests/test_litellm/a2a_protocol/test_a2a_exception_mapping_utils.py +++ b/tests/test_litellm/a2a_protocol/test_a2a_exception_mapping_utils.py @@ -171,9 +171,12 @@ async def test_stream_with_retry_raises_after_localhost_retries_exhausted(): api_base="https://agent.example", agent_name="test-agent", ) + async def _drain(): + async for _chunk in stream: + pytest.fail("expected retry exhaustion to raise before yielding") + with pytest.raises( RuntimeError, match="no response received after retry attempts", ): - async for _chunk in stream: - pytest.fail("expected retry exhaustion to raise before yielding") + await _drain() diff --git a/tests/test_litellm/caching/test_qdrant_semantic_cache.py b/tests/test_litellm/caching/test_qdrant_semantic_cache.py index 852bed4a9df..a5fbaf151ca 100644 --- a/tests/test_litellm/caching/test_qdrant_semantic_cache.py +++ b/tests/test_litellm/caching/test_qdrant_semantic_cache.py @@ -966,3 +966,67 @@ async def test_qdrant_async_embedding_explicit_limit_beats_deployment_limit(monk sent_input = router.aembedding.call_args.kwargs["input"] assert _token_count("sem-embed", sent_input) == 3 + + +@pytest.mark.asyncio +async def test_qdrant_async_embedding_call_is_bounded(monkeypatch): + from litellm.caching.qdrant_semantic_cache import QdrantSemanticCache + + cache = QdrantSemanticCache.__new__(QdrantSemanticCache) + cache.embedding_model = "sem-embed" + cache.embedding_max_input_tokens = None + cache.embedding_timeout = 1.5 + + router = MagicMock() + router.get_configured_token_limits.return_value = (None, None) + router.aembedding = AsyncMock(return_value={"data": [{"embedding": [0.1, 0.2]}]}) + monkeypatch.setitem( + sys.modules, + "litellm.proxy.proxy_server", + _router_proxy_module(router, "sem-embed"), + ) + + await cache._get_async_embedding("What is the capital of France?") + + assert router.aembedding.call_args.kwargs["timeout"] == 1.5 + assert router.aembedding.call_args.kwargs["num_retries"] == 0 + + +@pytest.mark.asyncio +async def test_qdrant_async_embedding_gives_up_on_unresponsive_endpoint(monkeypatch): + import asyncio + import time + + from litellm.caching.qdrant_semantic_cache import QdrantSemanticCache + + cache = QdrantSemanticCache.__new__(QdrantSemanticCache) + cache.embedding_model = "sem-embed" + cache.embedding_max_input_tokens = None + cache.embedding_timeout = 0.05 + + async def never_responds(**kwargs): + await asyncio.sleep(3) + return {"data": [{"embedding": [0.1, 0.2]}]} + + router = MagicMock() + router.get_configured_token_limits.return_value = (None, None) + router.aembedding = never_responds + monkeypatch.setitem( + sys.modules, + "litellm.proxy.proxy_server", + _router_proxy_module(router, "sem-embed"), + ) + + started = time.monotonic() + with pytest.raises(asyncio.TimeoutError): + await cache._get_async_embedding("What is the capital of France?") + assert time.monotonic() - started < 1.0 + + +def test_qdrant_semantic_cache_defaults_embedding_timeout(): + from litellm.caching.qdrant_semantic_cache import QdrantSemanticCache + from litellm.constants import SEMANTIC_CACHE_EMBEDDING_TIMEOUT_SECONDS + + cache = QdrantSemanticCache.__new__(QdrantSemanticCache) + assert cache.embedding_timeout == SEMANTIC_CACHE_EMBEDDING_TIMEOUT_SECONDS + assert SEMANTIC_CACHE_EMBEDDING_TIMEOUT_SECONDS < 60 diff --git a/tests/test_litellm/caching/test_redis_cache.py b/tests/test_litellm/caching/test_redis_cache.py index 59200719197..6a76decd5b1 100644 --- a/tests/test_litellm/caching/test_redis_cache.py +++ b/tests/test_litellm/caching/test_redis_cache.py @@ -520,13 +520,15 @@ async def test_circuit_breaker_covers_lua_script_execution(redis_no_ping): counted toward taking Redis out of the pool and kept paying a full socket timeout each, which is the traffic the outage hurts most. """ + from redis.exceptions import ConnectionError as RedisConnectionError + from litellm.constants import REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD cache = RedisCache(host="127.0.0.1", port=_closed_port(), socket_timeout=0.5) run_script = cache.async_register_script("return 1") for _ in range(REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD): - with pytest.raises(Exception): + with pytest.raises(RedisConnectionError): await run_script(keys=["lit4930"], args=[1]) with pytest.raises(Exception, match="circuit breaker is open"): @@ -603,8 +605,11 @@ async def test_only_connectivity_failures_open_the_breaker(error, opens_breaker) async def failing_call(): raise raised - for _ in range(breaker.failure_threshold + 1): - with pytest.raises(Exception): + for _ in range(breaker.failure_threshold): + with pytest.raises(type(raised)): await _run_under_circuit_breaker(breaker, "op", failing_call) + with pytest.raises(Exception, match="circuit breaker is open" if opens_breaker else "boom"): + await _run_under_circuit_breaker(breaker, "op", failing_call) + assert breaker.is_open() is opens_breaker diff --git a/tests/test_litellm/caching/test_redis_semantic_cache.py b/tests/test_litellm/caching/test_redis_semantic_cache.py index 9fd333cf87c..66271579d31 100644 --- a/tests/test_litellm/caching/test_redis_semantic_cache.py +++ b/tests/test_litellm/caching/test_redis_semantic_cache.py @@ -310,11 +310,12 @@ def test_redis_semantic_cache_reraises_unexpected_isolated_index_error(monkeypat monkeypatch.setenv("REDIS_PORT", "6379") monkeypatch.setenv("REDIS_PASSWORD", "test_password") + cache = RedisSemanticCache( + similarity_threshold=0.8, + index_name="existing_index", + ) + with pytest.raises(ValueError, match="connection failed"): - cache = RedisSemanticCache( - similarity_threshold=0.8, - index_name="existing_index", - ) _ = cache.llmcache @@ -1329,3 +1330,157 @@ def test_redis_llmcache_setter_supported(): sentinel = MagicMock() cache.llmcache = sentinel assert cache.llmcache is sentinel + + +def _router_proxy_module(router, model_name): + import types + + fake_proxy = types.ModuleType("litellm.proxy.proxy_server") + fake_proxy.llm_router = router + fake_proxy.llm_model_list = [{"model_name": model_name}] + return fake_proxy + + +def test_redis_sync_embedding_call_is_bounded(monkeypatch): + import sys + + from litellm.caching.redis_semantic_cache import RedisSemanticCache + + cache = RedisSemanticCache.__new__(RedisSemanticCache) + cache.embedding_model = "sem-embed" + cache.embedding_timeout = 1.5 + + router = MagicMock() + router.get_configured_token_limits.return_value = (None, None) + router.embedding = MagicMock(return_value={"data": [{"embedding": [0.5, 0.6]}]}) + monkeypatch.setitem( + sys.modules, + "litellm.proxy.proxy_server", + _router_proxy_module(router, "sem-embed"), + ) + + assert cache._get_embedding("hello") == [0.5, 0.6] + assert router.embedding.call_args.kwargs["timeout"] == 1.5 + assert router.embedding.call_args.kwargs["num_retries"] == 0 + + +@pytest.mark.asyncio +async def test_redis_async_embedding_call_is_bounded(monkeypatch): + import sys + + from litellm.caching.redis_semantic_cache import RedisSemanticCache + + cache = RedisSemanticCache.__new__(RedisSemanticCache) + cache.embedding_model = "sem-embed" + cache.embedding_timeout = 1.5 + + router = MagicMock() + router.get_configured_token_limits.return_value = (None, None) + router.aembedding = AsyncMock(return_value={"data": [{"embedding": [0.5, 0.6]}]}) + monkeypatch.setitem( + sys.modules, + "litellm.proxy.proxy_server", + _router_proxy_module(router, "sem-embed"), + ) + + assert await cache._get_async_embedding("hello") == [0.5, 0.6] + assert router.aembedding.call_args.kwargs["timeout"] == 1.5 + assert router.aembedding.call_args.kwargs["num_retries"] == 0 + + +@pytest.mark.asyncio +async def test_redis_async_embedding_gives_up_on_unresponsive_endpoint(monkeypatch): + import asyncio + import sys + import time + + from litellm.caching.redis_semantic_cache import RedisSemanticCache + + cache = RedisSemanticCache.__new__(RedisSemanticCache) + cache.embedding_model = "sem-embed" + cache.embedding_timeout = 0.05 + + async def never_responds(**kwargs): + await asyncio.sleep(3) + return {"data": [{"embedding": [0.1, 0.2]}]} + + router = MagicMock() + router.get_configured_token_limits.return_value = (None, None) + router.aembedding = never_responds + monkeypatch.setitem( + sys.modules, + "litellm.proxy.proxy_server", + _router_proxy_module(router, "sem-embed"), + ) + + started = time.monotonic() + with pytest.raises(ValueError, match="Failed to generate embedding"): + await cache._get_async_embedding("hello") + assert time.monotonic() - started < 1.0 + + +@pytest.mark.asyncio +async def test_redis_async_get_cache_fails_open_when_embedding_hangs(monkeypatch): + import asyncio + import sys + import time + + from litellm.caching.redis_semantic_cache import RedisSemanticCache + + cache = RedisSemanticCache.__new__(RedisSemanticCache) + cache.embedding_model = "sem-embed" + cache.embedding_timeout = 0.05 + cache.similarity_threshold = 0.8 + cache.distance_threshold = 0.2 + cache.llmcache = MagicMock() + + async def never_responds(**kwargs): + await asyncio.sleep(3) + return {"data": [{"embedding": [0.1, 0.2]}]} + + router = MagicMock() + router.get_configured_token_limits.return_value = (None, None) + router.aembedding = never_responds + monkeypatch.setitem( + sys.modules, + "litellm.proxy.proxy_server", + _router_proxy_module(router, "sem-embed"), + ) + + metadata = {} + started = time.monotonic() + result = await cache.async_get_cache( + key="test_key", + messages=[{"role": "user", "content": "What is the capital of France?"}], + metadata=metadata, + ) + elapsed = time.monotonic() - started + + assert result is None + assert metadata["semantic-similarity"] == 0.0 + assert elapsed < 1.0 + cache.llmcache.acheck.assert_not_called() + + +def test_cache_forwards_semantic_cache_embedding_timeout(): + from litellm.caching.caching import Cache + from litellm.types.caching import LiteLLMCacheType + + with patch("litellm.caching.caching.RedisSemanticCache") as backend: + Cache( + type=LiteLLMCacheType.REDIS_SEMANTIC, + similarity_threshold=0.8, + redis_url="redis://localhost:6379", + semantic_cache_embedding_timeout=2.5, + ) + + assert backend.call_args.kwargs["embedding_timeout"] == 2.5 + + +def test_redis_semantic_cache_defaults_embedding_timeout(): + from litellm.caching.redis_semantic_cache import RedisSemanticCache + from litellm.constants import SEMANTIC_CACHE_EMBEDDING_TIMEOUT_SECONDS + + cache = RedisSemanticCache.__new__(RedisSemanticCache) + assert cache.embedding_timeout == SEMANTIC_CACHE_EMBEDDING_TIMEOUT_SECONDS + assert SEMANTIC_CACHE_EMBEDDING_TIMEOUT_SECONDS < 60 diff --git a/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py b/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py index 5508931b35d..382b41807d4 100644 --- a/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py +++ b/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py @@ -3,7 +3,7 @@ import json import os import sys import unittest -from typing import List, Optional, Tuple +from typing import TYPE_CHECKING, List, Literal, Optional, Tuple from unittest.mock import ANY, MagicMock, Mock, patch import httpx @@ -17,6 +17,13 @@ from litellm.completion_extras.litellm_responses_transformation.transformation i LiteLLMResponsesTransformationHandler, ) +if TYPE_CHECKING: + from openai.types.responses import ResponseOutputItem + from openai.types.responses.response_reasoning_item import ResponseReasoningItem + + from litellm.types.llms.openai import ResponsesAPIResponse + from litellm.types.utils import ModelResponse + def test_convert_chat_completion_messages_to_responses_api_image_input(): from litellm.completion_extras.litellm_responses_transformation.transformation import ( @@ -3485,3 +3492,278 @@ async def test_acompletion_bridge_normalizes_tool_choice_on_the_wire( post_kwargs = mock_post.call_args.kwargs request_body = post_kwargs["json"] if "json" in post_kwargs else json.loads(post_kwargs["data"]) assert request_body["tool_choice"] == expected_wire_tool_choice + + +def _make_incomplete_responses_api_response( + incomplete_reason: Optional[str], + output: "List[ResponseOutputItem]", + status: Literal["completed", "incomplete"] = "incomplete", + empty_incomplete_details: bool = False, +) -> "ResponsesAPIResponse": + from litellm.types.llms.openai import ( + InputTokensDetails, + OutputTokensDetails, + ResponseAPIUsage, + ResponsesAPIResponse, + ) + + return ResponsesAPIResponse( + id="resp_incomplete", + created_at=1760144904, + error=None, + incomplete_details=( + {"reason": incomplete_reason} + if incomplete_reason is not None or empty_incomplete_details + else None + ), + instructions=None, + metadata={}, + model="gpt-5.6-sol", + object="response", + output=output, + parallel_tool_calls=True, + temperature=1.0, + tool_choice="auto", + tools=[], + top_p=1.0, + max_output_tokens=16, + previous_response_id=None, + reasoning={"effort": "high", "summary": None}, + status=status, + text={"format": {"type": "text"}, "verbosity": "medium"}, + truncation="disabled", + usage=ResponseAPIUsage( + input_tokens=37, + input_tokens_details=InputTokensDetails( + audio_tokens=None, cached_tokens=0, text_tokens=None + ), + output_tokens=16, + output_tokens_details=OutputTokensDetails( + reasoning_tokens=16, text_tokens=None + ), + total_tokens=53, + cost=None, + ), + user=None, + store=True, + background=False, + billing={"payer": "developer"}, + max_tool_calls=None, + prompt_cache_key=None, + safety_identifier=None, + service_tier="default", + top_logprobs=0, + ) + + +def _make_reasoning_only_output_item() -> "ResponseReasoningItem": + from openai.types.responses.response_reasoning_item import ResponseReasoningItem + + return ResponseReasoningItem( + id="rs_incomplete", + summary=[], + type="reasoning", + content=None, + encrypted_content="enc_abc", + status=None, + ) + + +def _call_transform_response( + handler: LiteLLMResponsesTransformationHandler, + raw_response: "ResponsesAPIResponse", +) -> "ModelResponse": + logging_obj = Mock() + logging_obj.model_call_details = {} + return handler.transform_response( + model="gpt-5.6-sol", + raw_response=raw_response, + model_response=_make_empty_model_response(), + logging_obj=logging_obj, + request_data={"model": "gpt-5.6-sol"}, + messages=[{"role": "user", "content": "compute something hard"}], + optional_params={}, + litellm_params={}, + encoding=Mock(), + ) + + +def test_transform_response_incomplete_reasoning_only_returns_empty_length_choice(): + handler = LiteLLMResponsesTransformationHandler() + raw_response = _make_incomplete_responses_api_response( + "max_output_tokens", [_make_reasoning_only_output_item()] + ) + + result = _call_transform_response(handler, raw_response) + + assert len(result.choices) == 1 + choice = result.choices[0] + assert choice.finish_reason == "length" + assert choice.index == 0 + assert choice.message.role == "assistant" + assert choice.message.content == "" + assert choice.message.reasoning_items[0]["encrypted_content"] == "enc_abc" + assert result.usage.prompt_tokens == 37 + assert result.usage.completion_tokens == 16 + assert result.usage.total_tokens == 53 + assert result.usage.completion_tokens_details.reasoning_tokens == 16 + + +def test_transform_response_incomplete_content_filter_maps_finish_reason(): + handler = LiteLLMResponsesTransformationHandler() + raw_response = _make_incomplete_responses_api_response( + "content_filter", [_make_reasoning_only_output_item()] + ) + + result = _call_transform_response(handler, raw_response) + + assert len(result.choices) == 1 + assert result.choices[0].finish_reason == "content_filter" + assert result.choices[0].message.content == "" + + +def test_transform_response_zero_choices_not_incomplete_still_raises(): + handler = LiteLLMResponsesTransformationHandler() + raw_response = _make_empty_responses_api_response() + + with pytest.raises(ValueError, match="Unknown items"): + _call_transform_response(handler, raw_response) + + +def test_transform_response_completed_with_reasonless_incomplete_details_keeps_stop(): + from openai.types.responses import ResponseOutputMessage, ResponseOutputText + + handler = LiteLLMResponsesTransformationHandler() + output_message = ResponseOutputMessage( + id="msg_complete", + content=[ + ResponseOutputText( + annotations=[], text="full answer", type="output_text", logprobs=[] + ) + ], + role="assistant", + status="completed", + type="message", + ) + raw_response = _make_incomplete_responses_api_response( + None, [output_message], status="completed", empty_incomplete_details=True + ) + + result = _call_transform_response(handler, raw_response) + + assert len(result.choices) == 1 + assert result.choices[0].finish_reason == "stop" + assert result.choices[0].message.content == "full answer" + + +def test_transform_response_incomplete_partial_text_overrides_finish_reason_to_length(): + from openai.types.responses import ResponseOutputMessage, ResponseOutputText + + handler = LiteLLMResponsesTransformationHandler() + output_message = ResponseOutputMessage( + id="msg_partial", + content=[ + ResponseOutputText( + annotations=[], text="partial answer", type="output_text", logprobs=[] + ) + ], + role="assistant", + status="incomplete", + type="message", + ) + raw_response = _make_incomplete_responses_api_response( + "max_output_tokens", [_make_reasoning_only_output_item(), output_message] + ) + + result = _call_transform_response(handler, raw_response) + + assert len(result.choices) == 1 + choice = result.choices[0] + assert choice.finish_reason == "length" + assert choice.message.content == "partial answer" + + +def test_response_incomplete_stream_event_emits_length_and_usage(): + from litellm.completion_extras.litellm_responses_transformation.transformation import ( + OpenAiResponsesToChatCompletionStreamIterator, + ) + + iterator = OpenAiResponsesToChatCompletionStreamIterator( + streaming_response=None, sync_stream=True + ) + + chunk = { + "type": "response.incomplete", + "response": { + "id": "resp_123", + "status": "incomplete", + "incomplete_details": {"reason": "max_output_tokens"}, + "output": [ + { + "type": "reasoning", + "id": "rs_1", + "encrypted_content": "enc_abc", + "summary": [], + } + ], + "usage": { + "input_tokens": 37, + "output_tokens": 16, + "output_tokens_details": {"reasoning_tokens": 16}, + "total_tokens": 53, + }, + }, + } + + result = iterator.chunk_parser(chunk) + + assert len(result.choices) == 1 + assert result.choices[0].finish_reason == "length" + assert result.choices[0].delta.reasoning_items[0]["encrypted_content"] == "enc_abc" + assert result.usage is not None + assert result.usage.prompt_tokens == 37 + assert result.usage.completion_tokens == 16 + assert result.usage.total_tokens == 53 + + +def test_response_incomplete_stream_event_content_filter_maps_finish_reason(): + from litellm.completion_extras.litellm_responses_transformation.transformation import ( + OpenAiResponsesToChatCompletionStreamIterator, + ) + + iterator = OpenAiResponsesToChatCompletionStreamIterator( + streaming_response=None, sync_stream=True + ) + + chunk = { + "type": "response.incomplete", + "response": { + "id": "resp_123", + "status": "incomplete", + "incomplete_details": {"reason": "content_filter"}, + "output": [], + }, + } + + result = iterator.chunk_parser(chunk) + + assert result.choices[0].finish_reason == "content_filter" + + +def test_response_incomplete_stream_event_without_details_defaults_to_length(): + from litellm.completion_extras.litellm_responses_transformation.transformation import ( + OpenAiResponsesToChatCompletionStreamIterator, + ) + + iterator = OpenAiResponsesToChatCompletionStreamIterator( + streaming_response=None, sync_stream=True + ) + + chunk = { + "type": "response.incomplete", + "response": {"id": "resp_123", "status": "incomplete", "output": []}, + } + + result = iterator.chunk_parser(chunk) + + assert result.choices[0].finish_reason == "length" diff --git a/tests/test_litellm/containers/test_container_api.py b/tests/test_litellm/containers/test_container_api.py index 4032c072594..de6fd1bc8ce 100644 --- a/tests/test_litellm/containers/test_container_api.py +++ b/tests/test_litellm/containers/test_container_api.py @@ -384,7 +384,7 @@ class TestContainerAPI: "container_create_handler", side_effect=Exception("API Error"), ): - with pytest.raises(Exception): + with pytest.raises(litellm.APIConnectionError): create_container( name="Error Test Container", custom_llm_provider="openai" ) diff --git a/tests/test_litellm/enterprise/enterprise_callbacks/send_emails/test_sendgrid_email.py b/tests/test_litellm/enterprise/enterprise_callbacks/send_emails/test_sendgrid_email.py index 40439a78a49..5fe4b217e4f 100644 --- a/tests/test_litellm/enterprise/enterprise_callbacks/send_emails/test_sendgrid_email.py +++ b/tests/test_litellm/enterprise/enterprise_callbacks/send_emails/test_sendgrid_email.py @@ -104,7 +104,7 @@ async def test_send_email_missing_api_key(): try: logger = SendGridEmailLogger() - with pytest.raises(ValueError): + with pytest.raises(ValueError, match='SENDGRID_API_KEY is not set'): await logger.send_email( from_email="test@example.com", to_email=["recipient@example.com"], diff --git a/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py b/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py index 6cc31f991a3..fcd03e77aa2 100644 --- a/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py +++ b/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py @@ -471,7 +471,7 @@ async def test_afile_content_error_reports_unified_id_not_provider_uri(): mock_router.get_deployment_credentials_with_provider = MagicMock(return_value=None) mock_router.afile_content = AsyncMock(side_effect=Exception("deployment failed")) - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='LiteLLM Managed File object with') as exc_info: await managed_files.afile_content( file_id=unified_file_id, litellm_parent_otel_span=None, diff --git a/tests/test_litellm/experimental_mcp_client/test_mcp_client.py b/tests/test_litellm/experimental_mcp_client/test_mcp_client.py index 6c3c852395b..1ddb2cc1c8d 100644 --- a/tests/test_litellm/experimental_mcp_client/test_mcp_client.py +++ b/tests/test_litellm/experimental_mcp_client/test_mcp_client.py @@ -75,11 +75,10 @@ class TestMCPClient: # Test missing stdio_config client = MCPClient(transport_type=MCPTransport.stdio) + async def _noop(session): + return None + with pytest.raises(ValueError, match="stdio_config is required for stdio transport"): - - async def _noop(session): - return None - await client.run_with_session(_noop) @pytest.mark.asyncio diff --git a/tests/test_litellm/integrations/bitbucket/test_bitbucket_integration.py b/tests/test_litellm/integrations/bitbucket/test_bitbucket_integration.py index 46cd1d6e765..142be536f6b 100644 --- a/tests/test_litellm/integrations/bitbucket/test_bitbucket_integration.py +++ b/tests/test_litellm/integrations/bitbucket/test_bitbucket_integration.py @@ -88,39 +88,44 @@ def test_bitbucket_prompt_manager_error_handling(mock_client_class): "access_token": "test-token", } + manager = BitBucketPromptManager(config, prompt_id="test_prompt") + with pytest.raises( Exception, match="Failed to load prompt 'test_prompt' from BitBucket" ): - manager = BitBucketPromptManager(config, prompt_id="test_prompt") - _ = manager.prompt_manager # This triggers the error + _ = manager.prompt_manager def test_bitbucket_prompt_manager_config_validation(): """Test BitBucketPromptManager configuration validation.""" # Test missing required fields - validation happens when prompt_manager is accessed - with pytest.raises( - ValueError, match="workspace, repository, and access_token are required" - ): - manager = BitBucketPromptManager({}) - _ = manager.prompt_manager # This triggers validation + manager = BitBucketPromptManager({}) with pytest.raises( ValueError, match="workspace, repository, and access_token are required" ): - manager = BitBucketPromptManager({"workspace": "test"}) - _ = manager.prompt_manager # This triggers validation + _ = manager.prompt_manager + + manager = BitBucketPromptManager({"workspace": "test"}) with pytest.raises( ValueError, match="workspace, repository, and access_token are required" ): - manager = BitBucketPromptManager({"repository": "test"}) - _ = manager.prompt_manager # This triggers validation + _ = manager.prompt_manager + + manager = BitBucketPromptManager({"repository": "test"}) with pytest.raises( ValueError, match="workspace, repository, and access_token are required" ): - manager = BitBucketPromptManager({"access_token": "test"}) - _ = manager.prompt_manager # This triggers validation + _ = manager.prompt_manager + + manager = BitBucketPromptManager({"access_token": "test"}) + + with pytest.raises( + ValueError, match="workspace, repository, and access_token are required" + ): + _ = manager.prompt_manager @patch("litellm.integrations.bitbucket.bitbucket_prompt_manager.BitBucketClient") diff --git a/tests/test_litellm/integrations/cloudzero/test_cloudzero_database.py b/tests/test_litellm/integrations/cloudzero/test_cloudzero_database.py index 89a5028011c..7f930f90247 100644 --- a/tests/test_litellm/integrations/cloudzero/test_cloudzero_database.py +++ b/tests/test_litellm/integrations/cloudzero/test_cloudzero_database.py @@ -51,7 +51,7 @@ async def test_get_usage_data_rejects_invalid_limit(monkeypatch: pytest.MonkeyPa """limit must coerce to int or raise ValueError before hitting the DB.""" db, query_mock = _setup_db(monkeypatch, []) - with pytest.raises(ValueError): + with pytest.raises(ValueError, match='limit must be an integer'): await db.get_usage_data(limit="invalid") assert query_mock.await_count == 0 diff --git a/tests/test_litellm/integrations/cloudzero/test_cz_stream_api.py b/tests/test_litellm/integrations/cloudzero/test_cz_stream_api.py index 440ce39e021..a715116e5ee 100644 --- a/tests/test_litellm/integrations/cloudzero/test_cz_stream_api.py +++ b/tests/test_litellm/integrations/cloudzero/test_cz_stream_api.py @@ -108,7 +108,7 @@ class TestCloudZeroStreamer: """Test _parse_and_convert_timestamp method with invalid timestamp.""" streamer = CloudZeroStreamer("test-key", "test-connection") - with pytest.raises(ValueError): + with pytest.raises(ValueError, match="Could not parse timestamp 'invalid-timestamp': Invalid"): streamer._parse_and_convert_timestamp("invalid-timestamp") def test_prepare_batch_payload(self): diff --git a/tests/test_litellm/integrations/focus/test_focus_database.py b/tests/test_litellm/integrations/focus/test_focus_database.py index d77af2dd170..5c13665f1f1 100644 --- a/tests/test_litellm/integrations/focus/test_focus_database.py +++ b/tests/test_litellm/integrations/focus/test_focus_database.py @@ -68,7 +68,7 @@ async def test_should_accept_string_timestamps(monkeypatch: pytest.MonkeyPatch): async def test_should_reject_invalid_limit(monkeypatch: pytest.MonkeyPatch): db, query_mock = _setup_db(monkeypatch, []) - with pytest.raises(ValueError): + with pytest.raises(ValueError, match='limit must be an integer'): await db.get_usage_data(limit="invalid") assert query_mock.await_count == 0 diff --git a/tests/test_litellm/integrations/focus/test_s3_destination.py b/tests/test_litellm/integrations/focus/test_s3_destination.py index f915b2c56a3..8e54b561f82 100644 --- a/tests/test_litellm/integrations/focus/test_s3_destination.py +++ b/tests/test_litellm/integrations/focus/test_s3_destination.py @@ -20,7 +20,7 @@ def _window(freq: str = "hourly", hour: int = 5) -> FocusTimeWindow: def test_should_require_bucket_name(): - with pytest.raises(ValueError): + with pytest.raises(ValueError, match='bucket_name must be provided for S'): FocusS3Destination(prefix="focus", config={}) diff --git a/tests/test_litellm/integrations/gitlab/test_gitlab_client.py b/tests/test_litellm/integrations/gitlab/test_gitlab_client.py index 4556950cd3e..529868ca06a 100644 --- a/tests/test_litellm/integrations/gitlab/test_gitlab_client.py +++ b/tests/test_litellm/integrations/gitlab/test_gitlab_client.py @@ -95,9 +95,9 @@ def enc_project(p): # how client encodes project in urls # Constructor / config tests # ----------------------------- def test_init_requires_project_and_token(): - with pytest.raises(ValueError): + with pytest.raises(ValueError, match='project and access_token are required'): GitLabClient({"project": "p"}) - with pytest.raises(ValueError): + with pytest.raises(ValueError, match='project and access_token are required'): GitLabClient({"access_token": "t"}) @@ -127,7 +127,7 @@ def test_set_ref_updates_effective_ref(): c = make_client(branch="main") c.set_ref("feature/x") assert c.ref == "feature/x" - with pytest.raises(ValueError): + with pytest.raises(ValueError, match='ref must be a non-empty string'): c.set_ref("") @@ -193,12 +193,12 @@ def test_get_file_content_permission_errors_are_mapped(): raw_url = f"https://gitlab.example.com/api/v4/projects/{enc_project('group/sub/repo')}/repository/files/secure%2Ffile.prompt/raw?ref=main" # raise_for_status will be called, so return 403 response (not an exception from transport) c.http_handler.routes[raw_url] = FakeResponse(status_code=403) - with pytest.raises(Exception) as ei: + with pytest.raises(Exception, match="Check your GitLab permissions for project 'group") as ei: c.get_file_content("secure/file.prompt") assert "Access denied" in str(ei.value) c.http_handler.routes[raw_url] = FakeResponse(status_code=401) - with pytest.raises(Exception) as ei2: + with pytest.raises(Exception, match='Authentication failed\\. Check your GitLab token and') as ei2: c.get_file_content("secure/file.prompt") assert "Authentication failed" in str(ei2.value) diff --git a/tests/test_litellm/integrations/levo/test_levo.py b/tests/test_litellm/integrations/levo/test_levo.py index 98b0327dbf2..903be644671 100644 --- a/tests/test_litellm/integrations/levo/test_levo.py +++ b/tests/test_litellm/integrations/levo/test_levo.py @@ -198,7 +198,7 @@ class TestLevoIntegration(unittest.TestCase): """Test health check returns unhealthy status when required vars are missing.""" # Try to create logger without required env vars # This should fail during config, but we can test health check logic - with pytest.raises(ValueError): + with pytest.raises(ValueError, match='LEVOAI_API_KEY environment variable is required for Levo'): LevoLogger.get_levo_config() @patch.dict( diff --git a/tests/test_litellm/integrations/otel/test_otel_v2_metrics.py b/tests/test_litellm/integrations/otel/test_otel_v2_metrics.py index e1b8e4b5721..b810ffdc6be 100644 --- a/tests/test_litellm/integrations/otel/test_otel_v2_metrics.py +++ b/tests/test_litellm/integrations/otel/test_otel_v2_metrics.py @@ -554,7 +554,7 @@ def test_token_type_rejected_from_either_list(attributes, monkeypatch): recorder rather than silently ignored, so the misconfig is caught at all.""" recorder = _recorder(monkeypatch, attributes) kwargs, response_obj, start, end = _build_call() - with pytest.raises(ValueError) as exc_info: + with pytest.raises(ValueError, match='otel\\.attributes: gen_ai\\.token\\.type is a structural') as exc_info: recorder.record(kwargs, response_obj, start, end) # The dedicated discriminator guard, not the generic unknown-name path: assert # the specific reason so dropping that guard (and falling through to "unknown diff --git a/tests/test_litellm/integrations/test_langfuse.py b/tests/test_litellm/integrations/test_langfuse.py index 6e57a36c5b6..3c7dd51bff8 100644 --- a/tests/test_litellm/integrations/test_langfuse.py +++ b/tests/test_litellm/integrations/test_langfuse.py @@ -1163,7 +1163,7 @@ def test_max_langfuse_clients_limit(): assert litellm.initialized_langfuse_clients == 2 # Third client should fail with exception - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Max langfuse clients reached') as exc_info: logger3 = LangFuseLogger( langfuse_public_key="test_key_3", langfuse_secret="test_secret_3", diff --git a/tests/test_litellm/interactions/test_agents_main_and_utils.py b/tests/test_litellm/interactions/test_agents_main_and_utils.py index 7c0183d20c6..f5523cf1cf8 100644 --- a/tests/test_litellm/interactions/test_agents_main_and_utils.py +++ b/tests/test_litellm/interactions/test_agents_main_and_utils.py @@ -314,7 +314,7 @@ class TestAsyncErrorWrapping: handler.create_agent.side_effect = RuntimeError("kaboom") with patch(_HANDLER_PATH, handler): - with pytest.raises(Exception): + with pytest.raises(litellm.APIConnectionError): await acreate(name="waverunner", api_key="AIza") @pytest.mark.asyncio @@ -323,7 +323,7 @@ class TestAsyncErrorWrapping: handler.get_agent.side_effect = RuntimeError("kaboom") with patch(_HANDLER_PATH, handler): - with pytest.raises(Exception): + with pytest.raises(litellm.APIConnectionError): await aget(name="waverunner", api_key="AIza") @pytest.mark.asyncio @@ -332,7 +332,7 @@ class TestAsyncErrorWrapping: handler.list_agents.side_effect = RuntimeError("kaboom") with patch(_HANDLER_PATH, handler): - with pytest.raises(Exception): + with pytest.raises(litellm.APIConnectionError): await alist(api_key="AIza") @pytest.mark.asyncio @@ -341,7 +341,7 @@ class TestAsyncErrorWrapping: handler.delete_agent.side_effect = RuntimeError("kaboom") with patch(_HANDLER_PATH, handler): - with pytest.raises(Exception): + with pytest.raises(litellm.APIConnectionError): await adelete(name="waverunner", api_key="AIza") @pytest.mark.asyncio @@ -350,5 +350,5 @@ class TestAsyncErrorWrapping: handler.list_agent_versions.side_effect = RuntimeError("kaboom") with patch(_HANDLER_PATH, handler): - with pytest.raises(Exception): + with pytest.raises(litellm.APIConnectionError): await alist_versions(name="waverunner", api_key="AIza") diff --git a/tests/test_litellm/interactions/test_google_interactions_integration.py b/tests/test_litellm/interactions/test_google_interactions_integration.py index 41f0fa0d7fb..49cd978c683 100644 --- a/tests/test_litellm/interactions/test_google_interactions_integration.py +++ b/tests/test_litellm/interactions/test_google_interactions_integration.py @@ -18,6 +18,7 @@ sys.path.insert(0, os.path.abspath("../../..")) import litellm import litellm.interactions as interactions +import openai # Test API key - should be set in environment GEMINI_API_KEY = os.getenv("GEMINI_API_KEY") @@ -258,7 +259,7 @@ class TestGoogleInteractionsErrorHandling: def test_invalid_model(self, api_key): """Test error handling for invalid model.""" - with pytest.raises(Exception): + with pytest.raises(openai.APIError): interactions.create( model="gemini/invalid-model-name-xyz", input="Hello", @@ -267,7 +268,7 @@ class TestGoogleInteractionsErrorHandling: def test_missing_model_and_agent(self, api_key): """Test error when neither model nor agent is provided.""" - with pytest.raises(Exception): # Can be ValueError or APIConnectionError + with pytest.raises((ValueError, litellm.APIConnectionError)): interactions.create( input="Hello", api_key=api_key, diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py index afe8e1d37a2..f66056a54e2 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py @@ -1056,6 +1056,53 @@ def test_generic_cost_per_token_gpt56_terra_cache_costs_by_tier_and_context( assert prompt_cost == pytest.approx(expected_prompt_cost) +@pytest.mark.parametrize("model", ["gpt-5.6-cyber", "daybreak-red-latest"]) +@pytest.mark.parametrize( + "prompt_tokens,input_rate,cache_write_rate,cache_read_rate,output_rate", + [ + (100000, 1.25e-5, 1.5625e-5, 1.25e-6, 7.5e-5), + (300000, 2.5e-5, 3.125e-5, 2.5e-6, 1.125e-4), + ], +) +def test_generic_cost_per_token_gpt56_cyber( + model, + prompt_tokens, + input_rate, + cache_write_rate, + cache_read_rate, + output_rate, + monkeypatch, +): + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) + + cached_tokens = 50000 + cache_write_tokens = 40000 + text_tokens = prompt_tokens - cached_tokens - cache_write_tokens + completion_tokens = 1000 + usage = Usage( + prompt_tokens=prompt_tokens, + completion_tokens=completion_tokens, + total_tokens=prompt_tokens + completion_tokens, + prompt_tokens_details=PromptTokensDetailsWrapper( + cached_tokens=cached_tokens, cache_write_tokens=cache_write_tokens + ), + ) + + prompt_cost, completion_cost = generic_cost_per_token( + model=model, + usage=usage, + custom_llm_provider="openai", + ) + + assert prompt_cost == pytest.approx( + text_tokens * input_rate + + cached_tokens * cache_read_rate + + cache_write_tokens * cache_write_rate + ) + assert completion_cost == pytest.approx(completion_tokens * output_rate) + + @pytest.mark.parametrize( "model,input_cost,output_cost,cache_read_cost", [ diff --git a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py index fffbc884782..08d8c17cc2e 100644 --- a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py +++ b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py @@ -1169,7 +1169,7 @@ def test_bedrock_image_processor_content_type_fallback_failure(): # Test with URL without recognizable extension image_url = "https://example.com/unknown-file" - with pytest.raises(ValueError) as excinfo: + with pytest.raises(ValueError, match='Unable to determine content type from URL: https') as excinfo: BedrockImageProcessor._post_call_image_processing(mock_response, image_url) assert "Unable to determine content type" in str(excinfo.value) diff --git a/tests/test_litellm/litellm_core_utils/test_fallback_generalizations.py b/tests/test_litellm/litellm_core_utils/test_fallback_generalizations.py index 0414836fa79..c6aac4d3991 100644 --- a/tests/test_litellm/litellm_core_utils/test_fallback_generalizations.py +++ b/tests/test_litellm/litellm_core_utils/test_fallback_generalizations.py @@ -288,7 +288,7 @@ def test_capability_info_backfills_requested_provider(restore_generalizations): def test_routing_only_match_does_not_resolve_model_info(restore_generalizations): restore_generalizations([{"name": "route", "pattern": r"^ceeco-", "model_info": {"litellm_provider": "openai"}}]) litellm.get_model_info.cache_clear() - with pytest.raises(Exception): + with pytest.raises(Exception, match="This model isn't mapped yet"): litellm.get_model_info("ceeco-fast-1", custom_llm_provider="openai") @@ -470,7 +470,7 @@ def test_shipped_adaptive_rule_requires_claude_prefix(shipped_cost_map): model = "openai/team-sonnet-5-1-alias" assert model not in litellm.model_cost assert match_capability_generalizations("team-sonnet-5-1-alias") is None - with pytest.raises(Exception): + with pytest.raises(Exception, match="This model isn't mapped yet"): litellm.get_model_info(model) @@ -496,7 +496,7 @@ def test_shipped_rules_lose_to_exact_entries_across_cost_ladder_variants(shipped from litellm.types.utils import ModelResponse, Usage assert "claude-haiku-4-5-20251001" in litellm.model_cost - with pytest.raises(Exception): + with pytest.raises(Exception, match="This model isn't mapped yet"): litellm.get_model_info("claude-haiku-4-5-20251001", custom_llm_provider="bedrock") entry = litellm.model_cost["us.anthropic.claude-haiku-4-5-20251001-v1:0"] diff --git a/tests/test_litellm/litellm_core_utils/test_initialize_dynamic_callback_params.py b/tests/test_litellm/litellm_core_utils/test_initialize_dynamic_callback_params.py index 0dca4f3a1b1..956f86a9292 100644 --- a/tests/test_litellm/litellm_core_utils/test_initialize_dynamic_callback_params.py +++ b/tests/test_litellm/litellm_core_utils/test_initialize_dynamic_callback_params.py @@ -110,7 +110,7 @@ def test_top_level_kwargs_overrides_metadata_slots(): def test_env_reference_at_top_level_raises_with_guidance(): kwargs = {"langfuse_public_key": "os.environ/LANGFUSE_PUBLIC_KEY"} - with pytest.raises(ValueError) as exc_info: + with pytest.raises(ValueError, match="Callback param 'langfuse_public_key' \\(from request body\\)") as exc_info: initialize_standard_callback_dynamic_params(kwargs) message = str(exc_info.value) @@ -127,7 +127,7 @@ def test_env_reference_in_metadata_raises_with_guidance(): } } - with pytest.raises(ValueError) as exc_info: + with pytest.raises(ValueError, match="Callback param 'langsmith_api_key' \\(from metadata\\) contains") as exc_info: initialize_standard_callback_dynamic_params(kwargs) message = str(exc_info.value) diff --git a/tests/test_litellm/litellm_core_utils/test_llm_judge.py b/tests/test_litellm/litellm_core_utils/test_llm_judge.py index 5c092caa7c3..a0a2311914b 100644 --- a/tests/test_litellm/litellm_core_utils/test_llm_judge.py +++ b/tests/test_litellm/litellm_core_utils/test_llm_judge.py @@ -27,7 +27,7 @@ def test_parse_json_verdict_tolerates_fences_and_prose(raw, expected): def test_parse_json_verdict_rejects_non_object(): - with pytest.raises(ValueError): + with pytest.raises(ValueError, match='judge response is not a JSON object'): parse_json_verdict('["not", "an", "object"]') with pytest.raises((json.JSONDecodeError, ValueError)): parse_json_verdict("no json here at all") diff --git a/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py b/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py index dff54515098..ccf353b1b6c 100644 --- a/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py +++ b/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py @@ -62,7 +62,7 @@ def test_realtime_streaming_store_message(): # Test 3: Invalid message format invalid_msg = "invalid json" - with pytest.raises(Exception): + with pytest.raises(json.JSONDecodeError): streaming.store_message(invalid_msg) # Test 4: Message type not in logged events diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py index 05b44fffbc5..fbdfcac1adc 100644 --- a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py +++ b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py @@ -982,7 +982,7 @@ async def test_bedrock_validation_error_raises_directly(logging_obj: Logging): make_call=_raise_400, ) - with pytest.raises(Exception) as excinfo: + with pytest.raises(Exception, match='litellm\\.BadRequestError: BedrockException') as excinfo: await response.__anext__() assert not isinstance(excinfo.value, MidStreamFallbackError) assert getattr(excinfo.value, "status_code", None) == 400 @@ -2143,10 +2143,13 @@ def test_raise_on_model_repetition( chunks = _build_chunks(chunks_pattern, len(chunks_pattern)) if should_raise: - with pytest.raises(litellm.InternalServerError) as exc_info: + def _feed(): for chunk in chunks: wrapper.chunks.append(chunk) wrapper.raise_on_model_repetition() + + with pytest.raises(litellm.InternalServerError) as exc_info: + _feed() assert "repeating the same chunk" in str(exc_info.value) else: for chunk in chunks: @@ -2719,7 +2722,7 @@ def test_dispatch_text_completion_codestral_requires_string( is a programming error and must surface loudly.""" initialized_custom_stream_wrapper.custom_llm_provider = "text-completion-codestral" - with pytest.raises(ValueError): + with pytest.raises(ValueError, match="chunk is not a string: \\{'not': 'a string'\\}"): _run_dispatch(initialized_custom_stream_wrapper, {"not": "a string"}) @@ -3388,6 +3391,128 @@ def test_record_partial_usage_for_failure_noop_without_chunks(): assert "combined_usage_object" not in logging_obj.model_call_details +def _wrapper_with_partial_chunks( + chunk_model: str, + usage: Optional[Usage] = None, + model: str = "gpt-4o-mini", + custom_llm_provider: str = "openai", +) -> tuple: + logging_obj = Logging( + model=model, + messages=[{"role": "user", "content": "Tell me a long story"}], + stream=True, + call_type="completion", + start_time=time.time(), + litellm_call_id="partial-usage-alias", + function_id="1245", + ) + logging_obj.model_call_details["custom_llm_provider"] = custom_llm_provider + logging_obj.optional_params = {} + wrapper = CustomStreamWrapper( + completion_stream=None, + model=model, + logging_obj=logging_obj, + custom_llm_provider=custom_llm_provider, + ) + wrapper.chunks = [ + ModelResponseStream( + id="chatcmpl-partial-alias-1", + created=1742056047, + model=chunk_model, + object="chat.completion.chunk", + choices=[ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta( + content="The Roman Empire began when", role="assistant" + ), + ) + ], + usage=usage, + ) + ] + return wrapper, logging_obj + + +def test_record_partial_usage_for_failure_prices_alias_restamped_chunks_at_real_model(): + wrapper, logging_obj = _wrapper_with_partial_chunks( + chunk_model="bedrock-claude-opus-5", + usage=Usage(prompt_tokens=40, completion_tokens=5, total_tokens=45), + model="us.anthropic.claude-opus-5", + custom_llm_provider="bedrock", + ) + assert "bedrock/bedrock-claude-opus-5" not in litellm.model_cost + + wrapper._record_partial_usage_for_failure() + + stashed = logging_obj.model_call_details["combined_usage_object"] + assert stashed.completion_tokens == 5 + rates = litellm.model_cost["us.anthropic.claude-opus-5"] + expected = 40 * rates["input_cost_per_token"] + 5 * rates["output_cost_per_token"] + assert logging_obj.model_call_details["response_cost"] == pytest.approx(expected) + + +def test_record_partial_usage_for_failure_counts_prompt_tokens_from_request_messages(): + wrapper, logging_obj = _wrapper_with_partial_chunks(chunk_model="my-public-alias") + + wrapper._record_partial_usage_for_failure() + + stashed = logging_obj.model_call_details["combined_usage_object"] + assert stashed.prompt_tokens > 0 + + +def test_record_partial_usage_for_failure_backfills_missing_cache_fields(): + wrapper, logging_obj = _wrapper_with_partial_chunks(chunk_model="gpt-4o-mini") + + wrapper._record_partial_usage_for_failure() + + stashed = logging_obj.model_call_details["combined_usage_object"] + assert stashed.cache_creation_input_tokens == 0 + assert stashed.cache_read_input_tokens == 0 + assert stashed.prompt_tokens_details is not None + assert stashed.prompt_tokens_details.cached_tokens == 0 + + +def test_record_partial_usage_for_failure_carries_up_openai_style_cached_tokens(): + recovered = Usage( + prompt_tokens=1000, + completion_tokens=10, + total_tokens=1010, + prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=500), + ) + wrapper, logging_obj = _wrapper_with_partial_chunks( + chunk_model="gpt-4o-mini", usage=recovered + ) + + wrapper._record_partial_usage_for_failure() + + stashed = logging_obj.model_call_details["combined_usage_object"] + assert stashed.cache_read_input_tokens == 500 + assert stashed.cache_creation_input_tokens == 0 + + +def test_record_partial_usage_for_failure_keeps_cache_values_recovered_from_chunks(): + recovered = Usage( + prompt_tokens=40, + completion_tokens=5, + total_tokens=45, + cache_read_input_tokens=7, + cache_creation_input_tokens=3, + ) + wrapper, logging_obj = _wrapper_with_partial_chunks( + chunk_model="gpt-4o-mini", usage=recovered + ) + + wrapper._record_partial_usage_for_failure() + + stashed = logging_obj.model_call_details["combined_usage_object"] + assert stashed.cache_read_input_tokens == 7 + assert stashed.cache_creation_input_tokens == 3 + assert stashed.prompt_tokens_details is not None + assert stashed.prompt_tokens_details.cached_tokens == 7 + + @pytest.mark.parametrize("sync_mode", [True, False]) @pytest.mark.asyncio async def test_stream_chunk_builder_raise_at_end_of_stream_still_recovers_usage( @@ -3616,10 +3741,13 @@ async def test_transport_read_error_before_finish_reason_raises(logging_obj: Log ) received = [] - with pytest.raises(MidStreamFallbackError): + async def _drain(): async for chunk in response: received.append(chunk) + with pytest.raises(MidStreamFallbackError): + await _drain() + fabricated_finish_reasons = [ chunk.choices[0].finish_reason for chunk in received @@ -4176,7 +4304,7 @@ async def test_stream_wrapper_anext_max_duration_timeout_restores_consumer_corre wrapper._stream_created_time = time.time() - 10 - with pytest.raises(Exception): + with pytest.raises(litellm.Timeout): await wrapper.__anext__() assert trace_id_var.get() == "outer-trace-max-duration" @@ -4323,7 +4451,9 @@ def test_handle_stream_fallback_error_restores_context_only_after_exception_mapp monkeypatch.setattr("litellm.litellm_core_utils.streaming_handler.exception_type", fake_exception_type) - with pytest.raises(Exception): + from litellm.exceptions import MidStreamFallbackError + + with pytest.raises(MidStreamFallbackError): wrapper._handle_stream_fallback_error(RuntimeError("boom")) # The mapper ran while the stream's own ids were still active. diff --git a/tests/test_litellm/litellm_core_utils/test_token_counter.py b/tests/test_litellm/litellm_core_utils/test_token_counter.py index 3c33ee13c3f..eec4b307c87 100644 --- a/tests/test_litellm/litellm_core_utils/test_token_counter.py +++ b/tests/test_litellm/litellm_core_utils/test_token_counter.py @@ -763,24 +763,6 @@ class TestTokenizerSelection(unittest.TestCase): ], } ], - [ - { - "role": "user", - "content": [ - { - "type": "text", - "text": "These are some sample images from a movie. Based on these images, what do you think the tone of the movie is?", - }, - { - "type": "text", - "image_url": { - "url": "https://gratisography.com/wp-content/uploads/2024/11/gratisography-augmented-reality-800x525.jpg", - "detail": "high", - }, - }, - ], - } - ], ], ) def test_bad_input_token_counter(model, messages): @@ -1174,7 +1156,7 @@ def test_count_content_list_rejects_unknown_type(): """ from litellm.litellm_core_utils.token_counter import _count_content_list - with pytest.raises(ValueError) as exc_info: + with pytest.raises(ValueError, match='Error getting number of tokens from content list: Invalid') as exc_info: _count_content_list( count_function=len, content_list=[{"type": "totally_unknown_block"}], diff --git a/tests/test_litellm/litellm_core_utils/test_url_utils.py b/tests/test_litellm/litellm_core_utils/test_url_utils.py index cef09f3f2b0..751b548adcd 100644 --- a/tests/test_litellm/litellm_core_utils/test_url_utils.py +++ b/tests/test_litellm/litellm_core_utils/test_url_utils.py @@ -100,12 +100,12 @@ class TestEncodeUrlPathSegment: @pytest.mark.parametrize("value", ["", ".", "..", None]) def test_rejects_empty_and_dot_segments(self, value): - with pytest.raises(ValueError): + with pytest.raises(ValueError, match="resource_id (is required|cannot be a dot path segment)"): encode_url_path_segment(value, field_name="resource_id") @pytest.mark.parametrize("value", ["../model", "model/../other", "/model"]) def test_rejects_dot_segments_in_multi_segment_paths(self, value): - with pytest.raises(ValueError): + with pytest.raises(ValueError, match="model (is required|cannot be a dot path segment)"): encode_url_path_segments(value, field_name="model") diff --git a/tests/test_litellm/llms/aiml/image_generation/test_aiml_image_generation_transformation.py b/tests/test_litellm/llms/aiml/image_generation/test_aiml_image_generation_transformation.py index 7485f2121df..dd74379a883 100644 --- a/tests/test_litellm/llms/aiml/image_generation/test_aiml_image_generation_transformation.py +++ b/tests/test_litellm/llms/aiml/image_generation/test_aiml_image_generation_transformation.py @@ -113,7 +113,7 @@ def test_flux_style_request_still_remaps_to_legacy_fields(): def test_openai_style_unsupported_param_raises_without_drop_params(): - with pytest.raises(ValueError): + with pytest.raises(ValueError, match='Supported parameters are'): AimlImageGenerationConfig().map_openai_params( non_default_params={"image_size": {"width": 1024, "height": 1024}}, optional_params={}, diff --git a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py index 391bd8566a2..d6aa384e03d 100644 --- a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py +++ b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py @@ -1986,10 +1986,11 @@ def test_effort_validation(): ) assert result["output_config"]["effort"] == effort + optional_params = {"output_config": {"effort": "invalid"}} + with pytest.raises( litellm.exceptions.BadRequestError, match="Invalid effort value" ): - optional_params = {"output_config": {"effort": "invalid"}} config.transform_request( model="claude-opus-4-5-20251101", messages=messages, @@ -2043,11 +2044,12 @@ def test_max_effort_rejected_for_opus_45(): messages = [{"role": "user", "content": "Test"}] + optional_params = {"output_config": {"effort": "max"}} + with pytest.raises( litellm.exceptions.BadRequestError, match="effort='max' is not supported by this model", ): - optional_params = {"output_config": {"effort": "max"}} config.transform_request( model="claude-opus-4-5-20251101", messages=messages, diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_mcp_handler.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_mcp_handler.py index 060c3e459d0..f3cb2956aeb 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_mcp_handler.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_mcp_handler.py @@ -61,7 +61,7 @@ def test_anthropic_messages_handler_skips_the_gateway_on_recursion(): "litellm.llms.anthropic.experimental_pass_through.messages.mcp_handler.anthropic_messages_with_mcp", new=AsyncMock(return_value={"routed": True}), ) as routed: - with pytest.raises(Exception): + with pytest.raises(ValueError, match='anthropic_messages_handler is not implemented for sync calls'): anthropic_messages_handler( max_tokens=100, messages=[{"role": "user", "content": "hi"}], @@ -80,7 +80,7 @@ def test_anthropic_messages_handler_leaves_native_tools_alone(): "litellm.llms.anthropic.experimental_pass_through.messages.mcp_handler.anthropic_messages_with_mcp", new=AsyncMock(return_value={"routed": True}), ) as routed: - with pytest.raises(Exception): + with pytest.raises(ValueError, match='anthropic_messages_handler is not implemented for sync calls'): anthropic_messages_handler( max_tokens=100, messages=[{"role": "user", "content": "hi"}], diff --git a/tests/test_litellm/llms/azure/videos/test_azure_video_transformation.py b/tests/test_litellm/llms/azure/videos/test_azure_video_transformation.py index d4f7a75895e..97c9e590d08 100644 --- a/tests/test_litellm/llms/azure/videos/test_azure_video_transformation.py +++ b/tests/test_litellm/llms/azure/videos/test_azure_video_transformation.py @@ -19,6 +19,7 @@ from litellm.types.videos.main import ( VideoCreateOptionalRequestParams, ) from litellm.types.router import GenericLiteLLMParams +from pydantic import ValidationError class TestAzureVideoConfig: @@ -299,7 +300,7 @@ class TestAzureVideoConfig: logging_obj = MagicMock() # Test that error responses raise exceptions - with pytest.raises(Exception): + with pytest.raises(ValidationError): self.config.transform_video_create_response( model=self.model, raw_response=mock_response, logging_obj=logging_obj ) diff --git a/tests/test_litellm/llms/azure_ai/rerank/test_azure_ai_rerank_transformation.py b/tests/test_litellm/llms/azure_ai/rerank/test_azure_ai_rerank_transformation.py index ffabce6e00c..602cbf68f3f 100644 --- a/tests/test_litellm/llms/azure_ai/rerank/test_azure_ai_rerank_transformation.py +++ b/tests/test_litellm/llms/azure_ai/rerank/test_azure_ai_rerank_transformation.py @@ -16,7 +16,7 @@ class TestAzureAIRerankConfigGetCompleteUrl: self.model = "azure_ai/cohere-rerank-v3-english" def test_api_base_required(self): - with pytest.raises(ValueError) as exc_info: + with pytest.raises(ValueError, match='Azure AI API Base is required\\. api_base=None\\. Set in') as exc_info: self.config.get_complete_url(api_base=None, model=self.model) assert "api_base=None" in str(exc_info.value) @@ -31,7 +31,7 @@ class TestAzureAIRerankConfigGetCompleteUrl: ], ) def test_api_base_requires_scheme(self, api_base): - with pytest.raises(ValueError) as exc_info: + with pytest.raises(ValueError, match='Azure AI API Base must be an absolute URL including scheme') as exc_info: self.config.get_complete_url(api_base=api_base, model=self.model) error_message = str(exc_info.value).lower() diff --git a/tests/test_litellm/llms/base_llm/search/test_base_search_transformation.py b/tests/test_litellm/llms/base_llm/search/test_base_search_transformation.py index b93ffdb0b44..e4402bbec49 100644 --- a/tests/test_litellm/llms/base_llm/search/test_base_search_transformation.py +++ b/tests/test_litellm/llms/base_llm/search/test_base_search_transformation.py @@ -270,7 +270,7 @@ async def test_asearch_does_not_leak_server_key_to_caller_api_base( new_callable=AsyncMock, ) as mock_get, ): - with pytest.raises(Exception): + with pytest.raises(litellm.APIConnectionError): await litellm.asearch( query="secrets", search_provider="serper", @@ -319,7 +319,7 @@ async def test_query_param_key_not_leaked_with_dummy_caller_key( "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.get", fake_get, ): - with pytest.raises(Exception): + with pytest.raises(litellm.APIConnectionError): await litellm.asearch( query="secrets", search_provider=provider, diff --git a/tests/test_litellm/llms/bedrock/chat/test_invoke_handler.py b/tests/test_litellm/llms/bedrock/chat/test_invoke_handler.py index ee50b9db015..e8964910c69 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_invoke_handler.py +++ b/tests/test_litellm/llms/bedrock/chat/test_invoke_handler.py @@ -2,17 +2,20 @@ import os import sys from unittest.mock import AsyncMock, MagicMock +import httpx import pytest sys.path.insert( 0, os.path.abspath("../../../../..") ) # Adds the parent directory to the system path +import litellm from litellm.llms.bedrock.chat.invoke_handler import ( AWSEventStreamDecoder, make_call, make_sync_call, ) +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler def test_transform_thinking_blocks_with_redacted_content(): @@ -293,3 +296,50 @@ def test_make_sync_call_honors_explicit_stream_chunk_size(): response.iter_bytes.assert_called_once_with(chunk_size=2048) + +def test_invoke_streaming_forwards_bedrock_response_headers(): + response = MagicMock() + response.status_code = 200 + response.iter_bytes = MagicMock(return_value=iter([])) + response.headers = httpx.Headers({"x-amzn-requestid": "req-789"}) + client = HTTPHandler() + client.post = MagicMock(return_value=response) + + stream = litellm.completion( + model="bedrock/invoke/anthropic.claude-haiku-4-5-20251001-v1:0", + messages=[{"role": "user", "content": "hi"}], + stream=True, + client=client, + aws_access_key_id="fake", + aws_secret_access_key="fake", + aws_region_name="us-east-1", + ) + + assert stream._hidden_params["additional_headers"]["llm_provider-x-amzn-requestid"] == "req-789" + + +@pytest.mark.asyncio +async def test_async_invoke_streaming_forwards_bedrock_response_headers(): + async def _no_bytes(chunk_size=None): + return + yield b"" + + response = MagicMock() + response.status_code = 200 + response.aiter_bytes = _no_bytes + response.headers = httpx.Headers({"x-amzn-requestid": "req-987"}) + client = AsyncHTTPHandler() + client.post = AsyncMock(return_value=response) + + stream = await litellm.acompletion( + model="bedrock/invoke/anthropic.claude-haiku-4-5-20251001-v1:0", + messages=[{"role": "user", "content": "hi"}], + stream=True, + client=client, + aws_access_key_id="fake", + aws_secret_access_key="fake", + aws_region_name="us-east-1", + ) + + assert stream._hidden_params["additional_headers"]["llm_provider-x-amzn-requestid"] == "req-987" + diff --git a/tests/test_litellm/llms/bedrock/test_base_aws_llm.py b/tests/test_litellm/llms/bedrock/test_base_aws_llm.py index cfe9930e76e..b9f8283b78e 100644 --- a/tests/test_litellm/llms/bedrock/test_base_aws_llm.py +++ b/tests/test_litellm/llms/bedrock/test_base_aws_llm.py @@ -1944,7 +1944,7 @@ def test_role_assumption_access_denied_raises_when_different_role(): with patch.object( base_aws_llm, "_is_already_running_as_role", return_value=False ): - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='An error occurred \\(AccessDenied\\) when calling the') as exc_info: base_aws_llm._auth_with_aws_role( aws_access_key_id=None, aws_secret_access_key=None, @@ -1969,7 +1969,7 @@ def test_role_assumption_non_access_denied_error_propagated(): ) with patch("boto3.client", return_value=mock_sts_client): - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='An error occurred \\(MalformedPolicyDocument\\) when calling') as exc_info: base_aws_llm._auth_with_aws_role( aws_access_key_id=None, aws_secret_access_key=None, diff --git a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py index 8281f3387d9..28c8e5c7ed6 100644 --- a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py +++ b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py @@ -109,7 +109,7 @@ class TestBedrockMantleResponsesURL: monkeypatch.delenv("BEDROCK_MANTLE_API_BASE", raising=False) monkeypatch.delenv("AWS_REGION", raising=False) cfg = BedrockMantleResponsesAPIConfig() - with pytest.raises(ValueError): + with pytest.raises(ValueError, match="api\\.aws\\.attacker\\.example/'\\. Region names must contain only"): cfg.get_complete_url( api_base=None, litellm_params={ @@ -1418,7 +1418,7 @@ class TestBedrockMantleResponsesSigV4: signer.get_credentials = MagicMock(side_effect=NoCredentialsError()) cfg = BedrockMantleResponsesAPIConfig(aws_signer=signer) - with pytest.raises(ValueError) as exc: + with pytest.raises(ValueError, match='Bedrock Mantle auth failed: no Bearer token and no usable') as exc: cfg.sign_request( headers={}, optional_params={"aws_region_name": "us-east-2"}, @@ -1448,7 +1448,7 @@ class TestBedrockMantleResponsesSigV4: signer.get_credentials = MagicMock(side_effect=cred_error) cfg = BedrockMantleResponsesAPIConfig(aws_signer=signer) - with pytest.raises(ValueError) as exc: + with pytest.raises(ValueError, match='Bedrock Mantle auth failed: no Bearer token and no usable') as exc: cfg.sign_request( headers={}, optional_params={"aws_region_name": "us-east-2"}, diff --git a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_transformation.py b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_transformation.py index 275fb460b9f..07910b0b56f 100644 --- a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_transformation.py +++ b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_transformation.py @@ -107,7 +107,7 @@ class TestBedrockMantleConfig: monkeypatch.delenv("BEDROCK_MANTLE_API_BASE", raising=False) monkeypatch.delenv("AWS_REGION", raising=False) cfg = BedrockMantleChatConfig() - with pytest.raises(ValueError): + with pytest.raises(ValueError, match="api\\.aws\\.attacker\\.example/'\\. Region names must contain only"): cfg._get_openai_compatible_provider_info( None, None, @@ -416,7 +416,7 @@ class TestBedrockMantleChatAuth: signer.get_credentials = MagicMock(side_effect=NoCredentialsError()) cfg = BedrockMantleChatConfig(aws_signer=signer) - with pytest.raises(ValueError) as exc: + with pytest.raises(ValueError, match='Bedrock Mantle auth failed: no Bearer token and no usable') as exc: cfg.sign_request( headers={}, optional_params={"aws_region_name": "us-east-2"}, diff --git a/tests/test_litellm/llms/bytez/chat/test_bytez_chat_transformation.py b/tests/test_litellm/llms/bytez/chat/test_bytez_chat_transformation.py index 2f8cc5484ba..94b8c51dd52 100644 --- a/tests/test_litellm/llms/bytez/chat/test_bytez_chat_transformation.py +++ b/tests/test_litellm/llms/bytez/chat/test_bytez_chat_transformation.py @@ -35,11 +35,10 @@ class TestBytezChatConfig: assert result["user-agent"] == f"litellm/{version}" def test_missing_api_key(self): - with pytest.raises(Exception) as excinfo: - config = BytezChatConfig() - - headers = {} + config = BytezChatConfig() + headers = {} + with pytest.raises(Exception, match='Missing api_key, make sure you pass in your api key') as excinfo: config.validate_environment( headers=headers, model=TEST_MODEL, diff --git a/tests/test_litellm/llms/chat/test_converse_handler.py b/tests/test_litellm/llms/chat/test_converse_handler.py index 2a3db5982ef..6f8a2788c38 100644 --- a/tests/test_litellm/llms/chat/test_converse_handler.py +++ b/tests/test_litellm/llms/chat/test_converse_handler.py @@ -1,14 +1,16 @@ +import json import os import sys -from unittest.mock import MagicMock +from unittest.mock import AsyncMock, MagicMock +import httpx import pytest import litellm from litellm.llms.bedrock.chat import BedrockConverseLLM from litellm.llms.bedrock.chat.converse_handler import make_sync_call from litellm.llms.bedrock.common_utils import _get_all_bedrock_regions -from litellm.llms.custom_httpx.http_handler import HTTPHandler +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler sys.path.insert( 0, os.path.abspath("../../../../..") @@ -202,6 +204,104 @@ def test_make_sync_call_honors_explicit_stream_chunk_size(): response.iter_bytes.assert_called_once_with(chunk_size=2048) +def _converse_response_body() -> dict: + return { + "output": {"message": {"role": "assistant", "content": [{"text": "hi"}]}}, + "stopReason": "end_turn", + "usage": {"inputTokens": 1, "outputTokens": 1, "totalTokens": 2}, + } + + +def test_converse_completion_forwards_bedrock_response_headers(): + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.json = MagicMock(return_value=_converse_response_body()) + mock_response.text = json.dumps(_converse_response_body()) + mock_response.headers = httpx.Headers({"x-amzn-requestid": "req-123"}) + client = HTTPHandler() + client.post = MagicMock(return_value=mock_response) + + response = litellm.completion( + model="bedrock/converse/anthropic.claude-haiku-4-5-20251001-v1:0", + messages=[{"role": "user", "content": "hi"}], + client=client, + aws_access_key_id="fake", + aws_secret_access_key="fake", + aws_region_name="us-east-1", + ) + + assert response._hidden_params["additional_headers"]["llm_provider-x-amzn-requestid"] == "req-123" + + +def test_converse_streaming_forwards_bedrock_response_headers(): + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.iter_bytes = MagicMock(return_value=iter([])) + mock_response.headers = httpx.Headers({"x-amzn-requestid": "req-456"}) + client = HTTPHandler() + client.post = MagicMock(return_value=mock_response) + + response = litellm.completion( + model="bedrock/converse/anthropic.claude-haiku-4-5-20251001-v1:0", + messages=[{"role": "user", "content": "hi"}], + stream=True, + client=client, + aws_access_key_id="fake", + aws_secret_access_key="fake", + aws_region_name="us-east-1", + ) + + assert response._hidden_params["additional_headers"]["llm_provider-x-amzn-requestid"] == "req-456" + + +@pytest.mark.asyncio +async def test_async_converse_completion_forwards_bedrock_response_headers(): + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.json = MagicMock(return_value=_converse_response_body()) + mock_response.text = json.dumps(_converse_response_body()) + mock_response.headers = httpx.Headers({"x-amzn-requestid": "req-abc"}) + client = AsyncHTTPHandler() + client.post = AsyncMock(return_value=mock_response) + + response = await litellm.acompletion( + model="bedrock/converse/anthropic.claude-haiku-4-5-20251001-v1:0", + messages=[{"role": "user", "content": "hi"}], + client=client, + aws_access_key_id="fake", + aws_secret_access_key="fake", + aws_region_name="us-east-1", + ) + + assert response._hidden_params["additional_headers"]["llm_provider-x-amzn-requestid"] == "req-abc" + + +@pytest.mark.asyncio +async def test_async_converse_streaming_forwards_bedrock_response_headers(): + async def _no_bytes(chunk_size=None): + return + yield b"" + + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.aiter_bytes = _no_bytes + mock_response.headers = httpx.Headers({"x-amzn-requestid": "req-def"}) + client = AsyncHTTPHandler() + client.post = AsyncMock(return_value=mock_response) + + response = await litellm.acompletion( + model="bedrock/converse/anthropic.claude-haiku-4-5-20251001-v1:0", + messages=[{"role": "user", "content": "hi"}], + stream=True, + client=client, + aws_access_key_id="fake", + aws_secret_access_key="fake", + aws_region_name="us-east-1", + ) + + assert response._hidden_params["additional_headers"]["llm_provider-x-amzn-requestid"] == "req-def" + + def test_completion_plumbs_stream_chunk_size_through_converse(): iter_bytes_spy = _stream_completion_with_spied_iter_bytes( model="bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0" diff --git a/tests/test_litellm/llms/custom_httpx/test_aiohttp_transport.py b/tests/test_litellm/llms/custom_httpx/test_aiohttp_transport.py index b0b092a541f..2dc7fbfd62a 100644 --- a/tests/test_litellm/llms/custom_httpx/test_aiohttp_transport.py +++ b/tests/test_litellm/llms/custom_httpx/test_aiohttp_transport.py @@ -127,10 +127,13 @@ async def test_client_payload_error_mid_stream_raises_read_error(): stream = AiohttpResponseStream(mock_response) # type: ignore received_chunks = [] - with pytest.raises(httpx.ReadError): + async def _drain(): async for chunk in stream: received_chunks.append(chunk) + with pytest.raises(httpx.ReadError): + await _drain() + assert received_chunks == [b"chunk1"] assert mock_response.closed is True @@ -151,10 +154,13 @@ async def test_client_payload_error_before_first_chunk_raises_read_error(): stream = AiohttpResponseStream(mock_response) # type: ignore received_chunks = [] - with pytest.raises(httpx.ReadError): + async def _drain(): async for chunk in stream: received_chunks.append(chunk) + with pytest.raises(httpx.ReadError): + await _drain() + assert received_chunks == [] assert mock_response.closed is True @@ -171,10 +177,13 @@ async def test_connection_closed_runtime_error_raises_read_error(): stream = AiohttpResponseStream(mock_response) # type: ignore received_chunks = [] - with pytest.raises(httpx.ReadError): + async def _drain(): async for chunk in stream: received_chunks.append(chunk) + with pytest.raises(httpx.ReadError): + await _drain() + assert received_chunks == [b"data1"] assert mock_response.closed is True @@ -209,10 +218,13 @@ async def test_transfer_encoding_error_raises_read_error(): stream = AiohttpResponseStream(mock_response) # type: ignore received_chunks = [] - with pytest.raises(httpx.ReadError): + async def _drain(): async for chunk in stream: received_chunks.append(chunk) + with pytest.raises(httpx.ReadError): + await _drain() + assert received_chunks == [b"data1"] assert mock_response.closed is True @@ -254,10 +266,13 @@ async def test_timeout_exception_gets_mapped(): received_chunks = [] # This should raise httpx.TimeoutException (mapped from aiohttp.ServerTimeoutError) - with pytest.raises(httpx.TimeoutException): + async def _drain(): async for chunk in stream: received_chunks.append(chunk) + with pytest.raises(httpx.TimeoutException): + await _drain() + # Should have received the first chunk before the error assert received_chunks == [b"chunk1"] @@ -1077,7 +1092,7 @@ async def test_session_closed_retry_does_not_close_concurrent_replacement(): raise StopAsyncIteration("stop after retry dispatch") with patch.object(transport, "_make_aiohttp_request", side_effect=fake_make_request): - with pytest.raises(Exception): + with pytest.raises(StopAsyncIteration): await transport.handle_async_request(httpx.Request("GET", "http://example.com")) try: diff --git a/tests/test_litellm/llms/custom_httpx/test_credential_leak_prevention.py b/tests/test_litellm/llms/custom_httpx/test_credential_leak_prevention.py index 0a3bf403bf8..bd9db87a765 100644 --- a/tests/test_litellm/llms/custom_httpx/test_credential_leak_prevention.py +++ b/tests/test_litellm/llms/custom_httpx/test_credential_leak_prevention.py @@ -287,10 +287,11 @@ class TestHTTPHandlerErrorPaths: "send", side_effect=_make_httpx_status_error(url="https://api.test.com?key=SECRET"), ): + kwargs = {"url": "https://api.test.com?key=SECRET"} + if method != "delete": + kwargs["data"] = {"test": 1} + with pytest.raises(MaskedHTTPStatusError) as exc_info: - kwargs = {"url": "https://api.test.com?key=SECRET"} - if method != "delete": - kwargs["data"] = {"test": 1} getattr(sync_handler, method)(**kwargs) assert "SECRET" not in str(exc_info.value.request.url) @@ -304,10 +305,11 @@ class TestHTTPHandlerErrorPaths: new_callable=AsyncMock, side_effect=_make_httpx_status_error(url="https://api.test.com?key=SECRET"), ): + kwargs = {"url": "https://api.test.com?key=SECRET"} + if method != "delete": + kwargs["data"] = {"test": 1} + with pytest.raises(MaskedHTTPStatusError) as exc_info: - kwargs = {"url": "https://api.test.com?key=SECRET"} - if method != "delete": - kwargs["data"] = {"test": 1} await getattr(async_handler, method)(**kwargs) assert "SECRET" not in str(exc_info.value.request.url) diff --git a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py index c568b82ebba..c87abbd8bc4 100644 --- a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py @@ -2369,3 +2369,79 @@ async def test_async_anthropic_messages_handler_carries_deployment_vertex_locati unconfigured_deployment = await logging_obj_after_handler(GenericLiteLLMParams()) assert "vertex_location" not in unconfigured_deployment.litellm_params + + +_GENERIC_STREAM_SSE = ( + b'data: {"id":"chatcmpl-1","object":"chat.completion.chunk","created":1,' + b'"model":"test-model","choices":[{"index":0,"delta":{"content":"hi"},' + b'"finish_reason":null}]}\n\n' + b"data: [DONE]\n\n" +) + + +def _generic_stream_upstream_response() -> httpx.Response: + return httpx.Response( + 200, + headers={ + "x-request-id": "generic-req-123", + "x-ratelimit-remaining-requests": "42", + }, + content=_GENERIC_STREAM_SSE, + request=httpx.Request("POST", "https://fake-vllm.test/v1/chat/completions"), + ) + + +def test_generic_http_handler_sync_streaming_forwards_provider_response_headers(): + """ + Regression test for the generic BaseLLMHTTPHandler streaming path used by + ~30 providers (deepseek, groq, hosted_vllm, databricks, openrouter, ...). + + The sync `completion()` streaming branch builds the CustomStreamWrapper from + `make_sync_call`, which returns the upstream response headers alongside the + stream. Those headers must reach the caller as `llm_provider-*` entries in + `_hidden_params["additional_headers"]`, which is what the proxy merges into + the client-facing response headers. + """ + mock_client = Mock(spec=HTTPHandler) + mock_client.post = Mock(return_value=_generic_stream_upstream_response()) + + response = litellm.completion( + model="hosted_vllm/test-model", + messages=[{"role": "user", "content": "Hello"}], + api_base="https://fake-vllm.test/v1", + api_key="sk-test", + stream=True, + client=mock_client, + ) + + additional_headers = response._hidden_params["additional_headers"] + assert additional_headers["llm_provider-x-request-id"] == "generic-req-123" + assert additional_headers["llm_provider-x-ratelimit-remaining-requests"] == "42" + + assert "".join([chunk.choices[0].delta.content or "" for chunk in response]) == "hi" + + +@pytest.mark.asyncio +async def test_generic_http_handler_async_streaming_forwards_provider_response_headers(): + """ + Companion to the sync test above for `acompletion_stream_function`, which + builds its CustomStreamWrapper from `make_async_call_stream_helper`. + """ + mock_client = AsyncMock(spec=AsyncHTTPHandler) + mock_client.post = AsyncMock(return_value=_generic_stream_upstream_response()) + + response = await litellm.acompletion( + model="hosted_vllm/test-model", + messages=[{"role": "user", "content": "Hello"}], + api_base="https://fake-vllm.test/v1", + api_key="sk-test", + stream=True, + client=mock_client, + ) + + additional_headers = response._hidden_params["additional_headers"] + assert additional_headers["llm_provider-x-request-id"] == "generic-req-123" + assert additional_headers["llm_provider-x-ratelimit-remaining-requests"] == "42" + + collected = [chunk async for chunk in response] + assert "".join([chunk.choices[0].delta.content or "" for chunk in collected]) == "hi" diff --git a/tests/test_litellm/llms/deepinfra/test_deepinfra_rerank_transformation.py b/tests/test_litellm/llms/deepinfra/test_deepinfra_rerank_transformation.py index a5411078cf7..ae3c166e7aa 100644 --- a/tests/test_litellm/llms/deepinfra/test_deepinfra_rerank_transformation.py +++ b/tests/test_litellm/llms/deepinfra/test_deepinfra_rerank_transformation.py @@ -258,7 +258,7 @@ class TestDeepinfraRerankTransform: status_code = 401 headers = {"content-type": "application/json"} - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Authentication failed') as exc_info: self.config.get_error_class(error_message, status_code, headers) # The method should raise a BaseLLMException @@ -271,7 +271,7 @@ class TestDeepinfraRerankTransform: status_code = 404 headers = {"content-type": "application/json"} - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Model not found') as exc_info: self.config.get_error_class(error_message, status_code, headers) # Should extract the nested error message @@ -284,7 +284,7 @@ class TestDeepinfraRerankTransform: status_code = 503 headers = {"content-type": "application/json"} - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Service unavailable') as exc_info: self.config.get_error_class(error_message, status_code, headers) # Should extract the string detail @@ -296,7 +296,7 @@ class TestDeepinfraRerankTransform: status_code = 500 headers = {"content-type": "application/json"} - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Invalid JSON error message') as exc_info: self.config.get_error_class(error_message, status_code, headers) # Should use the original error message when JSON parsing fails diff --git a/tests/test_litellm/llms/fal_ai/image_generation/test_fal_ai_gpt_image_2_transformation.py b/tests/test_litellm/llms/fal_ai/image_generation/test_fal_ai_gpt_image_2_transformation.py index 3c8cf9f9e0a..1a527230f1b 100644 --- a/tests/test_litellm/llms/fal_ai/image_generation/test_fal_ai_gpt_image_2_transformation.py +++ b/tests/test_litellm/llms/fal_ai/image_generation/test_fal_ai_gpt_image_2_transformation.py @@ -128,19 +128,23 @@ def test_transform_image_generation_request(): @pytest.mark.parametrize( - "model", + ("model", "expected_cost_for_two_images"), [ - "openai/gpt-image-2", - "gpt-image-2", - "openai/gpt-image-2/edit", + ("openai/gpt-image-2", 0.29), + ("gpt-image-2", 0.29), + ("openai/gpt-image-2/edit", 0.302), ], ) -def test_cost_calculator_uses_registry_price(model, monkeypatch: pytest.MonkeyPatch): +def test_cost_calculator_uses_registry_price( + model, expected_cost_for_two_images, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) + litellm.get_model_info.cache_clear() response = ImageResponse( data=[ ImageObject(url="https://v3b.fal.media/files/b/one.png"), ImageObject(url="https://v3b.fal.media/files/b/two.png"), ] ) - assert cost_calculator(model=model, image_response=response) == pytest.approx(0.29) + assert cost_calculator(model=model, image_response=response) == pytest.approx(expected_cost_for_two_images) diff --git a/tests/test_litellm/llms/fal_ai/image_generation/test_fal_ai_nano_banana_transformation.py b/tests/test_litellm/llms/fal_ai/image_generation/test_fal_ai_nano_banana_transformation.py index 593593bfa73..c0f74eff51b 100644 --- a/tests/test_litellm/llms/fal_ai/image_generation/test_fal_ai_nano_banana_transformation.py +++ b/tests/test_litellm/llms/fal_ai/image_generation/test_fal_ai_nano_banana_transformation.py @@ -113,7 +113,7 @@ def test_response_format_is_ignored(): def test_unsupported_param_raises_without_drop_params(): - with pytest.raises(ValueError): + with pytest.raises(ValueError, match="Supported parameters are \\['n', 'response_format', 'size'\\]\\."): FalAINanoBananaConfig().map_openai_params( non_default_params={"style": "vivid"}, optional_params={}, diff --git a/tests/test_litellm/llms/fal_ai/test_cost_calculator.py b/tests/test_litellm/llms/fal_ai/test_cost_calculator.py new file mode 100644 index 00000000000..f167aceaa95 --- /dev/null +++ b/tests/test_litellm/llms/fal_ai/test_cost_calculator.py @@ -0,0 +1,156 @@ +import pytest + +import litellm +from litellm.litellm_core_utils.llm_cost_calc.utils import CostCalculatorUtils +from litellm.llms.fal_ai.cost_calculator import cost_calculator +from litellm.types.utils import ImageObject, ImageResponse + + +@pytest.fixture(autouse=True) +def _use_local_model_cost_map(monkeypatch): + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) + litellm.get_model_info.cache_clear() + yield + litellm.get_model_info.cache_clear() + + +def _image_response(num_images: int = 1) -> ImageResponse: + return ImageResponse(data=[ImageObject(url="https://example.com/img.png") for _ in range(num_images)]) + + +def test_high_quality_1024x1024_uses_keyed_price(): + cost = cost_calculator( + model="openai/gpt-image-2", + image_response=_image_response(), + optional_params={"quality": "high", "image_size": {"width": 1024, "height": 1024}}, + ) + assert cost == pytest.approx(0.211) + + +def test_alias_model_uses_keyed_price(): + cost = cost_calculator( + model="gpt-image-2", + image_response=_image_response(), + optional_params={"quality": "high", "image_size": {"width": 1024, "height": 1024}}, + ) + assert cost == pytest.approx(0.211) + + +def test_provider_prefixed_model_uses_keyed_price(): + cost = cost_calculator( + model="fal_ai/openai/gpt-image-2", + image_response=_image_response(), + optional_params={"quality": "high", "image_size": {"width": 1024, "height": 1024}}, + ) + assert cost == pytest.approx(0.211) + + +def test_provider_prefixed_edit_model_uses_keyed_edit_price(): + cost = cost_calculator( + model="fal_ai/openai/gpt-image-2/edit", + image_response=_image_response(), + optional_params={"quality": "high", "image_size": {"width": 1024, "height": 1024}}, + ) + assert cost == pytest.approx(0.219) + + +def test_default_request_priced_at_default_size_and_quality(): + cost = cost_calculator( + model="openai/gpt-image-2", + image_response=_image_response(), + optional_params={}, + ) + assert cost == pytest.approx(0.145) + + +def test_auto_quality_priced_as_high(): + cost = cost_calculator( + model="openai/gpt-image-2", + image_response=_image_response(), + optional_params={"quality": "auto", "image_size": {"width": 1024, "height": 1024}}, + ) + assert cost == pytest.approx(0.211) + + +def test_low_quality_4k_uses_keyed_price(): + cost = cost_calculator( + model="openai/gpt-image-2", + image_response=_image_response(), + optional_params={"quality": "low", "image_size": {"width": 3840, "height": 2160}}, + ) + assert cost == pytest.approx(0.012) + + +def test_named_fal_size_uses_keyed_price(): + cost = cost_calculator( + model="openai/gpt-image-2", + image_response=_image_response(), + optional_params={"quality": "high", "image_size": "square_hd"}, + ) + assert cost == pytest.approx(0.211) + + +def test_edit_model_uses_keyed_edit_price(): + cost = cost_calculator( + model="openai/gpt-image-2/edit", + image_response=_image_response(), + optional_params={"quality": "high", "image_size": {"width": 1024, "height": 1024}}, + ) + assert cost == pytest.approx(0.219) + + +def test_edit_model_without_size_falls_back_to_flat_price(): + cost = cost_calculator( + model="openai/gpt-image-2/edit", + image_response=_image_response(), + optional_params={"quality": "high"}, + ) + assert cost == pytest.approx(0.151) + + +def test_missing_optional_params_falls_back_to_flat_price(): + cost = cost_calculator( + model="openai/gpt-image-2", + image_response=_image_response(), + optional_params=None, + ) + assert cost == pytest.approx(0.145) + + +def test_unlisted_size_falls_back_to_flat_price(): + cost = cost_calculator( + model="openai/gpt-image-2", + image_response=_image_response(), + optional_params={"quality": "high", "image_size": {"width": 999, "height": 999}}, + ) + assert cost == pytest.approx(0.145) + + +def test_keyed_price_multiplies_per_image(): + cost = cost_calculator( + model="openai/gpt-image-2", + image_response=_image_response(num_images=2), + optional_params={"quality": "high", "image_size": {"width": 1024, "height": 1024}}, + ) + assert cost == pytest.approx(0.422) + + +def test_route_image_generation_passes_optional_params_to_fal(): + cost = CostCalculatorUtils.route_image_generation_cost_calculator( + model="openai/gpt-image-2", + completion_response=_image_response(), + custom_llm_provider="fal_ai", + optional_params={"quality": "high", "image_size": {"width": 1024, "height": 1024}}, + ) + assert cost == pytest.approx(0.211) + + +def test_route_image_generation_with_provider_prefixed_model_uses_keyed_price(): + cost = CostCalculatorUtils.route_image_generation_cost_calculator( + model="fal_ai/openai/gpt-image-2", + completion_response=_image_response(), + custom_llm_provider="fal_ai", + optional_params={"quality": "high", "image_size": {"width": 1024, "height": 1024}}, + ) + assert cost == pytest.approx(0.211) diff --git a/tests/test_litellm/llms/featherless_ai/chat/test_featherless_chat_transformation.py b/tests/test_litellm/llms/featherless_ai/chat/test_featherless_chat_transformation.py index 4dc467575a0..bf40abd7016 100644 --- a/tests/test_litellm/llms/featherless_ai/chat/test_featherless_chat_transformation.py +++ b/tests/test_litellm/llms/featherless_ai/chat/test_featherless_chat_transformation.py @@ -44,7 +44,7 @@ class TestFeatherlessAIConfig: """Test error handling when API key is missing""" config = FeatherlessAIConfig() - with pytest.raises(ValueError) as excinfo: + with pytest.raises(ValueError, match='Missing Featherless AI API Key') as excinfo: config.validate_environment( headers={}, model="featherless-ai/Qwerky-72B", @@ -112,7 +112,7 @@ class TestFeatherlessAIConfig: "tool_choice": {"type": "function", "function": {"name": "get_weather"}} } optional_params = {} - with pytest.raises(Exception) as excinfo: + with pytest.raises(Exception, match="litellm\\.UnsupportedParamsError: Featherless AI doesn't") as excinfo: config.map_openai_params( non_default_params=non_default_params, optional_params=optional_params, @@ -138,7 +138,7 @@ class TestFeatherlessAIConfig: assert "tools" not in result # Test with tools and drop_params=False - with pytest.raises(Exception) as excinfo: + with pytest.raises(Exception, match="litellm\\.UnsupportedParamsError: Featherless AI doesn't") as excinfo: config.map_openai_params( non_default_params=non_default_params, optional_params=optional_params, diff --git a/tests/test_litellm/llms/fireworks_ai/rerank/test_fireworks_ai_rerank_transformation.py b/tests/test_litellm/llms/fireworks_ai/rerank/test_fireworks_ai_rerank_transformation.py index 30bf5860dee..521ea4f8263 100644 --- a/tests/test_litellm/llms/fireworks_ai/rerank/test_fireworks_ai_rerank_transformation.py +++ b/tests/test_litellm/llms/fireworks_ai/rerank/test_fireworks_ai_rerank_transformation.py @@ -301,7 +301,7 @@ class TestFireworksAIRerankTransform: mock_logging = MagicMock() model_response = RerankResponse() - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Failed to parse response: Invalid JSON: line') as exc_info: self.config.transform_rerank_response( model=self.model, raw_response=mock_response, diff --git a/tests/test_litellm/llms/gemini/image_edit/test_gemini_image_edit_transformation.py b/tests/test_litellm/llms/gemini/image_edit/test_gemini_image_edit_transformation.py index 9b57e1991de..bd9b7006e58 100644 --- a/tests/test_litellm/llms/gemini/image_edit/test_gemini_image_edit_transformation.py +++ b/tests/test_litellm/llms/gemini/image_edit/test_gemini_image_edit_transformation.py @@ -244,7 +244,7 @@ class TestGeminiImageEditTransformation: def test_transform_image_edit_request_without_image_raises(self) -> None: optional_params = {} - with pytest.raises(ValueError): + with pytest.raises(ValueError, match='Gemini image edit requires at least one image\\.'): self.config.transform_image_edit_request( model=self.model, prompt=self.prompt, diff --git a/tests/test_litellm/llms/gemini/test_gemini_client_setup.py b/tests/test_litellm/llms/gemini/test_gemini_client_setup.py index 51c6fedf5b8..48b010aca48 100644 --- a/tests/test_litellm/llms/gemini/test_gemini_client_setup.py +++ b/tests/test_litellm/llms/gemini/test_gemini_client_setup.py @@ -28,7 +28,7 @@ def test_gemini_completion_no_api_key(): del os.environ[key] # Test without mock_response to ensure actual API key validation - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='in _complete_vertex_ai_beta') as exc_info: completion( model="gemini/gemini-1.5-flash", messages=[{"role": "user", "content": "Test message"}], @@ -60,7 +60,7 @@ def test_gemini_completion_no_api_key_with_mock(): with patch("litellm.get_secret") as mock_get_secret: mock_get_secret.return_value = None - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='in _complete_vertex_ai_beta') as exc_info: completion( model="gemini/gemini-1.5-flash", messages=[{"role": "user", "content": "Test message"}], diff --git a/tests/test_litellm/llms/github_copilot/test_github_copilot_transformation.py b/tests/test_litellm/llms/github_copilot/test_github_copilot_transformation.py index f69ba7df938..52f1a6a99b8 100644 --- a/tests/test_litellm/llms/github_copilot/test_github_copilot_transformation.py +++ b/tests/test_litellm/llms/github_copilot/test_github_copilot_transformation.py @@ -866,7 +866,7 @@ class TestGithubCopilotTransformResponse: ) model_response = ModelResponse() - with pytest.raises(Exception): + with pytest.raises(json.JSONDecodeError): config.transform_response( model="github_copilot/claude-opus-4.7", raw_response=raw_response, diff --git a/tests/test_litellm/llms/hosted_vllm/test_hosted_vllm_rerank_transformation.py b/tests/test_litellm/llms/hosted_vllm/test_hosted_vllm_rerank_transformation.py index 6425e815db0..e6e6aa946d5 100644 --- a/tests/test_litellm/llms/hosted_vllm/test_hosted_vllm_rerank_transformation.py +++ b/tests/test_litellm/llms/hosted_vllm/test_hosted_vllm_rerank_transformation.py @@ -109,7 +109,7 @@ class TestHostedVLLMRerankTransform: ) assert url2 == "https://api.example.com/rerank" # Raises if api_base is None - with pytest.raises(ValueError): + with pytest.raises(ValueError, match='api_base must be provided for Hosted VLLM rerank'): self.config.get_complete_url(None, self.model) def test_transform_response(self): diff --git a/tests/test_litellm/llms/huggingface/rerank/test_huggingface_rerank_transformation.py b/tests/test_litellm/llms/huggingface/rerank/test_huggingface_rerank_transformation.py index b7ae8aa5fb1..9d6b7290eb6 100644 --- a/tests/test_litellm/llms/huggingface/rerank/test_huggingface_rerank_transformation.py +++ b/tests/test_litellm/llms/huggingface/rerank/test_huggingface_rerank_transformation.py @@ -232,7 +232,7 @@ def test_huggingface_rerank_error_handling(mock_post): mock_response.text = "Unauthorized" mock_post.return_value = mock_response - with pytest.raises(Exception): + with pytest.raises(litellm.APIConnectionError): litellm.rerank( model="huggingface/BAAI/bge-reranker-base", query="hello", diff --git a/tests/test_litellm/llms/langflow/chat/test_langflow_chat_transformation.py b/tests/test_litellm/llms/langflow/chat/test_langflow_chat_transformation.py index c03919a0659..0c241add77b 100644 --- a/tests/test_litellm/llms/langflow/chat/test_langflow_chat_transformation.py +++ b/tests/test_litellm/llms/langflow/chat/test_langflow_chat_transformation.py @@ -46,7 +46,7 @@ def test_langflow_config_get_complete_url(): def test_langflow_config_get_complete_url_requires_api_base(): config = LangFlowConfig() - with pytest.raises(ValueError): + with pytest.raises(ValueError, match='api_base is required for LangFlow\\. Set it via'): config.get_complete_url( api_base=None, api_key=None, @@ -233,7 +233,7 @@ def test_langflow_extra_body_cannot_inject_tweaks_into_run_payload(): return resp with patch.object(HTTPHandler, "post", side_effect=fake_post): - with pytest.raises(Exception): + with pytest.raises(litellm.APIConnectionError): litellm.completion( model="langflow/my-flow", messages=[{"role": "user", "content": "hello"}], diff --git a/tests/test_litellm/llms/modelscope/image_generation/test_modelscope_image_gen_transformation.py b/tests/test_litellm/llms/modelscope/image_generation/test_modelscope_image_gen_transformation.py index 7f00f53c451..fbcec3d4d2e 100644 --- a/tests/test_litellm/llms/modelscope/image_generation/test_modelscope_image_gen_transformation.py +++ b/tests/test_litellm/llms/modelscope/image_generation/test_modelscope_image_gen_transformation.py @@ -154,7 +154,7 @@ class TestModelScopeImageGenerationTransformation: mock_get_secret.return_value = None headers = {} - with pytest.raises(ValueError) as exc_info: + with pytest.raises(ValueError, match='MODELSCOPE_API_KEY is not set\\. Please set it via') as exc_info: self.config.validate_environment( headers=headers, model=self.model, @@ -367,7 +367,7 @@ class TestModelScopeImageGenerationTransformation: model_response = ImageResponse(data=[]) - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='litellm\\.BadRequestError: ModelScope error: Invalid prompt') as exc_info: self.config.transform_image_generation_response( model=self.model, raw_response=mock_response, @@ -393,7 +393,7 @@ class TestModelScopeImageGenerationTransformation: model_response = ImageResponse(data=[]) - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='litellm\\.InternalServerError: Error parsing ModelScope') as exc_info: self.config.transform_image_generation_response( model=self.model, raw_response=mock_response, diff --git a/tests/test_litellm/llms/novita/chat/test_novita_chat_transformation.py b/tests/test_litellm/llms/novita/chat/test_novita_chat_transformation.py index 7a00b361252..ade5e4176e8 100644 --- a/tests/test_litellm/llms/novita/chat/test_novita_chat_transformation.py +++ b/tests/test_litellm/llms/novita/chat/test_novita_chat_transformation.py @@ -47,7 +47,7 @@ class TestNovitaConfig: """Test error handling when API key is missing""" config = NovitaConfig() - with pytest.raises(ValueError) as excinfo: + with pytest.raises(ValueError, match='Missing Novita AI API Key - A call is being made to novita') as excinfo: config.validate_environment( headers={}, model="novita/meta-llama/llama-3.3-70b-instruct", diff --git a/tests/test_litellm/llms/oci/chat/test_oci_chat_transformation.py b/tests/test_litellm/llms/oci/chat/test_oci_chat_transformation.py index 0f0033cae36..5aa96a66d2d 100644 --- a/tests/test_litellm/llms/oci/chat/test_oci_chat_transformation.py +++ b/tests/test_litellm/llms/oci/chat/test_oci_chat_transformation.py @@ -95,10 +95,10 @@ class TestOCIChatConfig: modified_params = params.copy() del modified_params[key] - with pytest.raises(Exception) as excinfo: - config = OCIChatConfig() - headers = {} + config = OCIChatConfig() + headers = {} + with pytest.raises(Exception, match='Missing required parameters: oci_user, oci_fingerprint') as excinfo: config.validate_environment( headers=headers, model=TEST_MODEL, @@ -272,7 +272,7 @@ class TestOCIChatConfig: "oci_serving_mode": "INVALID_MODE", } - with pytest.raises(Exception) as excinfo: + with pytest.raises(Exception, match="kwarg `oci_serving_mode` must be either 'ON_DEMAND' or") as excinfo: config.transform_request( model=TEST_MODEL_NAME, messages=TEST_MESSAGES, # type: ignore @@ -892,7 +892,7 @@ class TestOCISignerSupport: optional_params = {"oci_signer": MockSigner(), "method": "INVALID"} - with pytest.raises(ValueError) as excinfo: + with pytest.raises(ValueError, match='Unsupported HTTP method: INVALID') as excinfo: config.sign_request( headers={}, optional_params=optional_params, @@ -1604,7 +1604,7 @@ class TestOCIKeyNormalization: # We can't fully test signing without a real key, but we can verify # the error message indicates the key was processed (not a type error) - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='why-can-t-i-import-my-pem-file for more details\\.') as exc_info: sign_with_manual_credentials( headers={}, optional_params=optional_params, @@ -1630,7 +1630,7 @@ class TestOCIKeyNormalization: "oci_key": crlf_pem, } - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='why-can-t-i-import-my-pem-file for more details\\.') as exc_info: sign_with_manual_credentials( headers={}, optional_params=optional_params, @@ -1692,7 +1692,7 @@ class TestOCIValidateEnvironment: def test_missing_required_credentials_raises_error(self, config): """Test that missing required credentials raise an error.""" - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Missing required parameters: oci_user, oci_fingerprint') as exc_info: config.validate_environment( headers={}, model="oci/xai.grok-3", @@ -1875,7 +1875,7 @@ class TestOCIImageUrlTransformation: } ] - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Prop `image_url` must be a string or an object with a `url`') as exc_info: adapt_messages_to_generic_oci_standard(messages) assert "image_url" in str(exc_info.value) @@ -1899,7 +1899,7 @@ class TestOCIImageUrlTransformation: } ] - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Prop `image_url` must be a string or an object with a `url`') as exc_info: adapt_messages_to_generic_oci_standard(messages) assert "image_url" in str(exc_info.value) diff --git a/tests/test_litellm/llms/openai/test_openai_common_utils.py b/tests/test_litellm/llms/openai/test_openai_common_utils.py index a28e133700e..bfd681cc06e 100644 --- a/tests/test_litellm/llms/openai/test_openai_common_utils.py +++ b/tests/test_litellm/llms/openai/test_openai_common_utils.py @@ -375,20 +375,26 @@ async def test_async_streaming_output_limit_400_maps_to_length_truncated_stream( @pytest.mark.parametrize("provider", ["openai", "azure"]) @pytest.mark.parametrize("stream", [False, True]) def test_sync_genuine_bad_request_still_raises(provider, stream): - with pytest.raises(litellm.BadRequestError): + def _call_and_drain(): result = litellm.completion( **_completion_kwargs(provider, _sync_client_raising(provider, GENUINE_400_MESSAGE), stream=stream) ) list(result) + with pytest.raises(litellm.BadRequestError): + _call_and_drain() + @pytest.mark.parametrize("provider", ["openai", "azure"]) @pytest.mark.parametrize("stream", [False, True]) @pytest.mark.asyncio async def test_async_genuine_bad_request_still_raises(provider, stream): - with pytest.raises(litellm.BadRequestError): + async def _call_and_drain(): result = await litellm.acompletion( **_completion_kwargs(provider, _async_client_raising(provider, GENUINE_400_MESSAGE), stream=stream) ) async for _ in result: pass + + with pytest.raises(litellm.BadRequestError): + await _call_and_drain() diff --git a/tests/test_litellm/llms/openai_like/test_cognition_provider.py b/tests/test_litellm/llms/openai_like/test_cognition_provider.py new file mode 100644 index 00000000000..5c71b60e08a --- /dev/null +++ b/tests/test_litellm/llms/openai_like/test_cognition_provider.py @@ -0,0 +1,217 @@ +""" +Tests for the Cognition provider identity. + +Cognition serves an OpenAI-compatible /v1/chat/completions surface, but it must resolve to its +own `cognition` provider so OpenAI-specific pricing and provider-level reporting never apply to +its traffic. +""" + +import json +from pathlib import Path + +import pytest + +import litellm + + +class TestCognitionProviderIdentity: + def test_cognition_is_a_registered_provider(self): + from litellm import LlmProviders + + assert LlmProviders.COGNITION.value == "cognition" + assert "cognition" in litellm.provider_list + + def test_cognition_json_config(self): + from litellm.llms.openai_like.json_loader import JSONProviderRegistry + + cognition = JSONProviderRegistry.get("cognition") + assert cognition is not None + assert cognition.base_url == "https://api.cognition.ai/v1" + assert cognition.api_key_env == "COGNITION_API_KEY" + assert cognition.api_base_env == "COGNITION_API_BASE" + + def test_cognition_in_openai_compatible_providers(self): + from litellm.constants import openai_compatible_providers + + assert "cognition" in openai_compatible_providers + + def test_prefixed_model_resolves_to_cognition_not_openai(self): + from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider + + model, provider, _, api_base = get_llm_provider( + model="cognition/swe-1.7", + custom_llm_provider=None, + api_base=None, + api_key=None, + ) + + assert model == "swe-1.7" + assert provider == "cognition" + assert api_base == "https://api.cognition.ai/v1" + + def test_explicit_api_base_and_key_win(self): + from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider + + _, provider, api_key, api_base = get_llm_provider( + model="cognition/swe-1.7", + custom_llm_provider=None, + api_base="https://cognition.internal.example/v1", + api_key="sk-test", + ) + + assert provider == "cognition" + assert api_base == "https://cognition.internal.example/v1" + assert api_key == "sk-test" + + def test_api_base_autodetects_cognition(self, monkeypatch: pytest.MonkeyPatch): + from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider + + monkeypatch.setenv("COGNITION_API_KEY", "sk-cognition-env") + + _, provider, api_key, api_base = get_llm_provider( + model="swe-1.7", + custom_llm_provider=None, + api_base="https://api.cognition.ai/v1", + api_key=None, + ) + + assert provider == "cognition" + assert api_base == "https://api.cognition.ai/v1" + assert api_key == "sk-cognition-env" + + def test_autodetected_api_base_keeps_the_caller_api_key(self, monkeypatch: pytest.MonkeyPatch): + from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider + + monkeypatch.setenv("COGNITION_API_KEY", "sk-cognition-env") + + _, provider, api_key, _ = get_llm_provider( + model="swe-1.7", + custom_llm_provider=None, + api_base="https://api.cognition.ai/v1", + api_key="sk-cognition-caller", + ) + + assert provider == "cognition" + assert api_key == "sk-cognition-caller" + + def test_env_api_key_is_read_from_cognition_variable(self, monkeypatch: pytest.MonkeyPatch): + from litellm.llms.openai_like.dynamic_config import create_config_class + from litellm.llms.openai_like.json_loader import JSONProviderRegistry + + monkeypatch.delenv("OPENAI_API_KEY", raising=False) + monkeypatch.setenv("COGNITION_API_KEY", "sk-cognition-env") + + provider = JSONProviderRegistry.get("cognition") + assert provider is not None + + api_base, api_key = create_config_class(provider)()._get_openai_compatible_provider_info(None, None) + assert api_base == "https://api.cognition.ai/v1" + assert api_key == "sk-cognition-env" + + +class TestCognitionCostTracking: + @pytest.mark.parametrize( + "model, input_cost, output_cost, cache_read_cost", + [ + ("cognition/swe-1.6", 5e-07, 2.5e-06, 2e-07), + ("cognition/swe-1.7", 5e-07, 2.5e-06, 2e-07), + ("cognition/swe-1.7-lightning", 2.5e-06, 1.25e-05, 1e-06), + ], + ) + def test_cost_map_entries(self, model: str, input_cost: float, output_cost: float, cache_read_cost: float): + info = litellm.get_model_info(model=model) + + assert info["litellm_provider"] == "cognition" + assert info["mode"] == "chat" + assert info["input_cost_per_token"] == input_cost + assert info["output_cost_per_token"] == output_cost + assert info["cache_read_input_token_cost"] == cache_read_cost + + @pytest.mark.parametrize( + "model, expected_prompt_cost, expected_completion_cost", + [ + ("cognition/swe-1.7", 0.5, 2.5), + ("cognition/swe-1.7-lightning", 2.5, 12.5), + ], + ) + def test_cost_differs_from_openai_pricing( + self, model: str, expected_prompt_cost: float, expected_completion_cost: float + ): + """A cognition-prefixed model must never be priced off an OpenAI cost entry.""" + from litellm.cost_calculator import cost_per_token + + prompt_cost, completion_cost = cost_per_token( + model=model, + prompt_tokens=1_000_000, + completion_tokens=1_000_000, + custom_llm_provider="cognition", + ) + + assert prompt_cost == pytest.approx(expected_prompt_cost) + assert completion_cost == pytest.approx(expected_completion_cost) + + def test_lightning_is_five_times_the_standard_tier(self): + standard = litellm.get_model_info(model="cognition/swe-1.7") + lightning = litellm.get_model_info(model="cognition/swe-1.7-lightning") + + assert lightning["input_cost_per_token"] == pytest.approx(standard["input_cost_per_token"] * 5) + assert lightning["output_cost_per_token"] == pytest.approx(standard["output_cost_per_token"] * 5) + + def test_supported_endpoints_matrix(self): + matrix = json.loads((Path(litellm.__file__).parent / "provider_endpoints_support_backup.json").read_text()) + + endpoints = matrix["providers"]["cognition"]["endpoints"] + assert endpoints["chat_completions"] is True + assert endpoints["messages"] is True + assert endpoints["responses"] is True + assert endpoints["embeddings"] is False + + +class TestCognitionRouting: + @pytest.mark.asyncio + async def test_router_spend_is_attributed_to_cognition_pricing(self): + """Routed traffic is costed off the cognition entry, not an OpenAI one.""" + from litellm import Router + + router = Router( + model_list=[ + { + "model_name": "swe", + "litellm_params": {"model": "cognition/swe-1.7", "api_key": "sk-test"}, + } + ] + ) + + response = await router.acompletion( + model="swe", + messages=[{"role": "user", "content": "hi"}], + mock_response="hello from swe", + ) + + usage = response.usage + expected = usage.prompt_tokens * 5e-07 + usage.completion_tokens * 2.5e-06 + assert response._hidden_params["response_cost"] == pytest.approx(expected) + + @pytest.mark.asyncio + async def test_router_spend_uses_the_lightning_entry_for_lightning(self): + """The Lightning tier is its own model, costed off its own entry.""" + from litellm import Router + + router = Router( + model_list=[ + { + "model_name": "swe-lightning", + "litellm_params": {"model": "cognition/swe-1.7-lightning", "api_key": "sk-test"}, + } + ] + ) + + response = await router.acompletion( + model="swe-lightning", + messages=[{"role": "user", "content": "hi"}], + mock_response="hello from swe lightning", + ) + + usage = response.usage + expected = usage.prompt_tokens * 2.5e-06 + usage.completion_tokens * 1.25e-05 + assert response._hidden_params["response_cost"] == pytest.approx(expected) diff --git a/tests/test_litellm/llms/pg_vector/vector_stores/test_pg_vector_transformation.py b/tests/test_litellm/llms/pg_vector/vector_stores/test_pg_vector_transformation.py index 56953a574d6..1d44b2bc278 100644 --- a/tests/test_litellm/llms/pg_vector/vector_stores/test_pg_vector_transformation.py +++ b/tests/test_litellm/llms/pg_vector/vector_stores/test_pg_vector_transformation.py @@ -42,7 +42,7 @@ class TestPGVectorStoreConfig: litellm_params = GenericLiteLLMParams() headers = {} - with pytest.raises(ValueError) as exc_info: + with pytest.raises(ValueError, match='PG Vector API key is required\\. Set PG_VECTOR_API_KEY') as exc_info: config.validate_environment(headers, litellm_params) assert "PG Vector API key is required" in str(exc_info.value) @@ -84,7 +84,7 @@ class TestPGVectorStoreConfig: config = PGVectorStoreConfig() litellm_params = {} - with pytest.raises(ValueError) as exc_info: + with pytest.raises(ValueError, match='PG Vector API base URL is required\\. Set') as exc_info: config.get_complete_url(None, litellm_params) assert "PG Vector API base URL is required" in str(exc_info.value) diff --git a/tests/test_litellm/llms/recraft/image_edit/test_recraft_image_edit_transformation.py b/tests/test_litellm/llms/recraft/image_edit/test_recraft_image_edit_transformation.py index 0acabd05805..47811321133 100644 --- a/tests/test_litellm/llms/recraft/image_edit/test_recraft_image_edit_transformation.py +++ b/tests/test_litellm/llms/recraft/image_edit/test_recraft_image_edit_transformation.py @@ -167,7 +167,7 @@ class TestRecraftImageEditTransformation: mock_response.status_code = 500 mock_response.headers = {} - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Error transforming image edit response: Invalid JSON: line') as exc_info: self.config.transform_image_edit_response( model=self.model, raw_response=mock_response, diff --git a/tests/test_litellm/llms/recraft/image_generation/test_recraft_image_gen_transformation.py b/tests/test_litellm/llms/recraft/image_generation/test_recraft_image_gen_transformation.py index 70311201969..ccc72dde7b8 100644 --- a/tests/test_litellm/llms/recraft/image_generation/test_recraft_image_gen_transformation.py +++ b/tests/test_litellm/llms/recraft/image_generation/test_recraft_image_gen_transformation.py @@ -64,7 +64,7 @@ class TestRecraftImageGenerationTransformation: non_default_params = {"n": 2, "unsupported_param": "value"} optional_params = {} - with pytest.raises(ValueError) as exc_info: + with pytest.raises(ValueError, match='Supported parameters are') as exc_info: self.config.map_openai_params( non_default_params=non_default_params, optional_params=optional_params, @@ -171,7 +171,7 @@ class TestRecraftImageGenerationTransformation: mock_get_secret.return_value = None headers = {} - with pytest.raises(ValueError) as exc_info: + with pytest.raises(ValueError, match='RECRAFT_API_KEY is not set') as exc_info: self.config.validate_environment( headers=headers, model=self.model, @@ -248,7 +248,7 @@ class TestRecraftImageGenerationTransformation: model_response = ImageResponse(data=[]) - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Error transforming image generation response: Invalid JSON') as exc_info: self.config.transform_image_generation_response( model=self.model, raw_response=mock_response, diff --git a/tests/test_litellm/llms/sagemaker/test_sagemaker_chat_transformation.py b/tests/test_litellm/llms/sagemaker/test_sagemaker_chat_transformation.py index 54e6f95c795..da6caca4f05 100644 --- a/tests/test_litellm/llms/sagemaker/test_sagemaker_chat_transformation.py +++ b/tests/test_litellm/llms/sagemaker/test_sagemaker_chat_transformation.py @@ -19,6 +19,8 @@ from unittest.mock import MagicMock import httpx import pytest +import litellm +from litellm.llms.custom_httpx.http_handler import HTTPHandler from litellm.llms.sagemaker.chat.transformation import SagemakerChatConfig @@ -233,3 +235,85 @@ def test_decoder_reassembles_frames_across_arbitrary_byte_boundaries(split_size) ] assert texts == [f"token{i} " for i in range(len(frames))] + + +_INFERENCE_COMPONENT_HEADER = "X-Amzn-SageMaker-Inference-Component" + +_STUB_COMPLETION_RESPONSE = { + "id": "chatcmpl-test", + "object": "chat.completion", + "created": 1700000000, + "model": "served-model", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "hi"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, +} + + +class _RequestCapturingHTTPHandler(HTTPHandler): + """Injected transport that records exactly what sagemaker_chat put on the wire.""" + + def __init__(self) -> None: + super().__init__() + self.request_headers: dict[str, str] = {} + self.request_body: dict = {} + + def post(self, url: str, headers=None, data=None, **kwargs) -> httpx.Response: + self.request_headers = dict(headers or {}) + self.request_body = json.loads(data) + return httpx.Response(200, json=_STUB_COMPLETION_RESPONSE, request=httpx.Request("POST", url)) + + +def _invoke_sagemaker_chat(monkeypatch, **extra_params) -> _RequestCapturingHTTPHandler: + """Drive one sagemaker_chat completion against an injected transport. + + A Bedrock API key short-circuits SigV4 inside `BaseAWSLLM._sign_request`, which would hide + whether the inference-component header is really covered by the signature, so it is cleared. + """ + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + client = _RequestCapturingHTTPHandler() + litellm.completion( + model="sagemaker_chat/my-endpoint", + messages=[{"role": "user", "content": "hi"}], + aws_access_key_id="AKIATESTTESTTESTTEST", + aws_secret_access_key="test-secret-key", + aws_region_name="us-east-1", + client=client, + **extra_params, + ) + return client + + +def test_model_id_is_sent_as_a_signed_inference_component_header(monkeypatch): + """`model_id` names an inference component and must reach SageMaker as a signed header. + + Endpoints backed by inference components reject any request without + `X-Amzn-SageMaker-Inference-Component` with HTTP 400 INFERENCE_COMPONENT_NAME_MISSING, so the + header has to be built before `sign_request` runs and end up inside SignedHeaders. + """ + client = _invoke_sagemaker_chat(monkeypatch, model_id="my-inference-component") + + assert client.request_headers[_INFERENCE_COMPONENT_HEADER] == "my-inference-component" + assert "x-amzn-sagemaker-inference-component" in client.request_headers["Authorization"] + + +def test_no_inference_component_header_when_model_id_is_unset(monkeypatch): + """Plain endpoints must not receive the header at all, not even an empty one.""" + client = _invoke_sagemaker_chat(monkeypatch) + + assert not any(name.lower() == _INFERENCE_COMPONENT_HEADER.lower() for name in client.request_headers) + + +def test_hf_model_name_becomes_the_body_model(monkeypatch): + """`hf_model_name` names the served model, and containers that validate the body's `model` + 404 on the endpoint name, so it has to replace it rather than ride along as an extra field.""" + client = _invoke_sagemaker_chat(monkeypatch, hf_model_name="org/served-model") + + assert client.request_body["model"] == "org/served-model" + assert "hf_model_name" not in client.request_body + + +def test_body_model_stays_the_endpoint_name_when_hf_model_name_is_unset(monkeypatch): + """Without `hf_model_name` the body must keep the model it has today.""" + client = _invoke_sagemaker_chat(monkeypatch) + + assert client.request_body["model"] == "my-endpoint" diff --git a/tests/test_litellm/llms/stability/image_generation/test_stability_image_generation.py b/tests/test_litellm/llms/stability/image_generation/test_stability_image_generation.py index 6f1a04e78d3..c5b3c8fbdc5 100644 --- a/tests/test_litellm/llms/stability/image_generation/test_stability_image_generation.py +++ b/tests/test_litellm/llms/stability/image_generation/test_stability_image_generation.py @@ -83,7 +83,7 @@ class TestStabilityImageGenerationConfig: non_default_params = {"unsupported_param": "value"} optional_params = {} - with pytest.raises(ValueError) as exc_info: + with pytest.raises(ValueError, match="Supported parameters are \\['n', 'size',") as exc_info: self.config.map_openai_params( non_default_params=non_default_params, optional_params=optional_params, @@ -168,7 +168,7 @@ class TestStabilityImageGenerationConfig: def test_validate_environment_raises_without_api_key(self): """Test that validate_environment raises error without API key""" - with pytest.raises(ValueError) as exc_info: + with pytest.raises(ValueError, match='STABILITY_API_KEY is not set\\. Please set it via') as exc_info: self.config.validate_environment( headers={}, model="stability/sd3", @@ -251,7 +251,7 @@ class TestStabilityImageGenerationConfig: model_response = ImageResponse(data=[]) mock_logging = MagicMock() - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Content was filtered by Stability AI safety systems') as exc_info: self.config.transform_image_generation_response( model="stability/sd3", raw_response=mock_response, diff --git a/tests/test_litellm/llms/tinyfish/test_tinyfish_search.py b/tests/test_litellm/llms/tinyfish/test_tinyfish_search.py index 2dcccb8ea7e..69afbb416aa 100644 --- a/tests/test_litellm/llms/tinyfish/test_tinyfish_search.py +++ b/tests/test_litellm/llms/tinyfish/test_tinyfish_search.py @@ -697,7 +697,7 @@ class TestErrorHandling: } } mock_response = _make_mock_response(body, status_code=400) - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='TinyFish Search: query is required\\. See https') as exc_info: config.transform_search_response( raw_response=mock_response, logging_obj=None ) @@ -713,7 +713,7 @@ class TestErrorHandling: mock_response = _make_mock_response( body, status_code=429, headers={"Retry-After": "60"} ) - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='TinyFish Search: rate limit exceeded\\. See https') as exc_info: config.transform_search_response( raw_response=mock_response, logging_obj=None ) @@ -728,7 +728,7 @@ class TestErrorHandling: config = TinyfishSearchConfig() body = {"errors": [{"code": "10000", "message": "Internal"}]} mock_response = _make_mock_response(body, status_code=502) - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='TinyFish Search') as exc_info: config.transform_search_response( raw_response=mock_response, logging_obj=None ) @@ -742,7 +742,7 @@ class TestErrorHandling: mock_response = _make_mock_response( json_data=None, status_code=502, text="Bad Gateway" ) - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='TinyFish Search: Bad Gateway<') as exc_info: config.transform_search_response( raw_response=mock_response, logging_obj=None ) @@ -756,7 +756,7 @@ class TestErrorHandling: mock_response = _make_mock_response( json_data=None, status_code=200, text="not json" ) - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='TinyFish Search: Expected JSON response, got: not json\\.') as exc_info: config.transform_search_response( raw_response=mock_response, logging_obj=None ) @@ -785,7 +785,7 @@ class TestErrorHandling: # check TinyFish's schema, not their own input. config = TinyfishSearchConfig() mock_response = _make_mock_response({"query": "x"}) # no `results` key - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='validation error for SearchResponse') as exc_info: config.transform_search_response( raw_response=mock_response, logging_obj=None ) diff --git a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_integration.py b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_integration.py index 272565990bd..8f9acafa49d 100644 --- a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_integration.py +++ b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_integration.py @@ -159,7 +159,7 @@ class TestVertexAIFilesIntegration: # This test ensures the type annotations and error messages include vertex_ai # Test that calling with unsupported provider raises appropriate error - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match="unsupported_provider' is not a valid LlmProviders") as exc_info: litellm.file_content( file_id="test-file-id", custom_llm_provider="unsupported_provider", # This should fail diff --git a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_streaming.py b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_streaming.py index 957fc7dbcf4..7383513fb96 100644 --- a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_streaming.py +++ b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_streaming.py @@ -40,6 +40,7 @@ from litellm.llms.vertex_ai.files.transformation import ( _openai_batch_jsonl_entry_to_vertex_rows, ) from litellm.types.llms.openai import CreateFileRequest +from litellm.llms.vertex_ai.common_utils import VertexAIError def _upload_stream(transformed) -> BaseFileUploadStream: @@ -561,7 +562,7 @@ class TestStreamingMediaUpload: async def test_failed_upload_raises(self): raw = _make_openai_jsonl_bytes(80) - with pytest.raises(Exception): + with pytest.raises(VertexAIError): await self._run(raw, status=403) async def test_request_timeout_is_forwarded(self): diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py index 8ee8186f6bb..8c1de12e7d9 100644 --- a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py @@ -1,3 +1,7 @@ +import base64 + +import pytest + from litellm.litellm_core_utils.prompt_templates.factory import ( convert_to_gemini_tool_call_result, ) @@ -784,6 +788,367 @@ def test_dummy_signature_with_function_call_mode(): assert gemini_parts[0]["thoughtSignature"] == expected_dummy +def _parallel_tool_calls(*signatures): + return [ + { + "id": f"call_{idx}", + "type": "function", + "function": { + "name": f"tool_{idx}", + "arguments": '{"location": "Paris"}', + **( + {"provider_specific_fields": {"thought_signature": signature}} + if signature is not None + else {} + ), + }, + "index": idx, + } + for idx, signature in enumerate(signatures) + ] + + +def _parallel_tool_calls_signed_via_id(*signatures): + """Parallel tool calls in the shape LiteLLM actually hands back to clients. + + The signature rides in the tool call id behind __thought__, which is what an + OpenAI-format client echoes back on the next turn. + """ + from litellm.litellm_core_utils.prompt_templates.factory import ( + _encode_tool_call_id_with_signature, + ) + + return [ + { + "id": _encode_tool_call_id_with_signature(f"call_{idx}", signature), + "type": "function", + "function": {"name": f"tool_{idx}", "arguments": '{"location": "Paris"}'}, + "index": idx, + } + for idx, signature in enumerate(signatures) + ] + + +REAL_THOUGHT_SIGNATURE = "Co4CAdHtim/rWgXbz2Ghp4tShzLeMASrPw6JJyYIC3cbVyZnKzU3uv8/wVzyS2sKRPL2m8QQHHXbNQhEEz500G7n" +PLACEHOLDER_SIGNATURE = base64.b64encode(b"skip_thought_signature_validator").decode( + "utf-8" +) + + +def test_dummy_signature_only_on_first_parallel_tool_call(): + """Google documents the placeholder as a last resort that degrades quality, so an unsigned + parallel turn replayed to gemini-3 gets a budget of exactly one.""" + from litellm.litellm_core_utils.prompt_templates.factory import ( + convert_to_gemini_tool_call_invoke, + ) + + gemini_parts = convert_to_gemini_tool_call_invoke( + { + "role": "assistant", + "content": None, + "tool_calls": _parallel_tool_calls(None, None, None), + }, + model="gemini-3-pro-preview", + ) + + assert len(gemini_parts) == 3 + assert gemini_parts[0]["thoughtSignature"] == PLACEHOLDER_SIGNATURE + assert "thoughtSignature" not in gemini_parts[1] + assert "thoughtSignature" not in gemini_parts[2] + + +def test_real_signature_on_first_parallel_tool_call_leaves_siblings_empty(): + """Gemini signs only the first of N parallel function calls, so a faithful replay has + nothing to attach to the siblings.""" + from litellm.litellm_core_utils.prompt_templates.factory import ( + convert_to_gemini_tool_call_invoke, + ) + + gemini_parts = convert_to_gemini_tool_call_invoke( + { + "role": "assistant", + "content": None, + "tool_calls": _parallel_tool_calls(REAL_THOUGHT_SIGNATURE, None, None), + }, + model="gemini-3-pro-preview", + ) + + assert len(gemini_parts) == 3 + assert gemini_parts[0]["thoughtSignature"] == REAL_THOUGHT_SIGNATURE + assert "thoughtSignature" not in gemini_parts[1] + assert "thoughtSignature" not in gemini_parts[2] + + +def test_real_signature_on_later_parallel_tool_call_is_preserved(): + """Clients may reorder or drop calls, so a signature that lands on a non-first call is + still the model's own and must survive the round trip.""" + from litellm.litellm_core_utils.prompt_templates.factory import ( + convert_to_gemini_tool_call_invoke, + ) + + gemini_parts = convert_to_gemini_tool_call_invoke( + { + "role": "assistant", + "content": None, + "tool_calls": _parallel_tool_calls(None, REAL_THOUGHT_SIGNATURE), + }, + model="gemini-3-pro-preview", + ) + + assert len(gemini_parts) == 2 + assert gemini_parts[0]["thoughtSignature"] == PLACEHOLDER_SIGNATURE + assert gemini_parts[1]["thoughtSignature"] == REAL_THOUGHT_SIGNATURE + + +def test_no_signatures_on_parallel_tool_calls_for_gemini_2_5(): + """Non-gemini-3 models never get a placeholder signature, on any call.""" + from litellm.litellm_core_utils.prompt_templates.factory import ( + convert_to_gemini_tool_call_invoke, + ) + + gemini_parts = convert_to_gemini_tool_call_invoke( + { + "role": "assistant", + "content": None, + "tool_calls": _parallel_tool_calls(None, None), + }, + model="gemini-2.5-flash", + ) + + assert len(gemini_parts) == 2 + assert all("thoughtSignature" not in part for part in gemini_parts) + + +def test_signature_embedded_in_tool_call_id_only_on_first_parallel_call(): + """The production shape: the signature arrives inside the first call's id, siblings have bare ids.""" + from litellm.litellm_core_utils.prompt_templates.factory import ( + convert_to_gemini_tool_call_invoke, + ) + + gemini_parts = convert_to_gemini_tool_call_invoke( + { + "role": "assistant", + "content": None, + "tool_calls": _parallel_tool_calls_signed_via_id( + REAL_THOUGHT_SIGNATURE, None, None + ), + }, + model="gemini-3-pro-preview", + ) + + assert len(gemini_parts) == 3 + assert gemini_parts[0]["thoughtSignature"] == REAL_THOUGHT_SIGNATURE + assert "thoughtSignature" not in gemini_parts[1] + assert "thoughtSignature" not in gemini_parts[2] + + +def test_tool_level_provider_specific_fields_signature_leaves_siblings_empty(): + """A signature on the tool call itself, rather than on its function, behaves the same way.""" + from litellm.litellm_core_utils.prompt_templates.factory import ( + convert_to_gemini_tool_call_invoke, + ) + + tool_calls = _parallel_tool_calls(None, None) + tool_calls[0]["provider_specific_fields"] = { + "thought_signature": REAL_THOUGHT_SIGNATURE + } + + gemini_parts = convert_to_gemini_tool_call_invoke( + {"role": "assistant", "content": None, "tool_calls": tool_calls}, + model="gemini-3-pro-preview", + ) + + assert len(gemini_parts) == 2 + assert gemini_parts[0]["thoughtSignature"] == REAL_THOUGHT_SIGNATURE + assert "thoughtSignature" not in gemini_parts[1] + + +def test_placeholder_lands_on_first_emitted_part_not_first_tool_call_entry(): + """A non-function entry (e.g. an OpenAI custom tool call) emits no part, so it must not + consume the one placeholder slot and leave the real first function call bare.""" + from litellm.litellm_core_utils.prompt_templates.factory import ( + convert_to_gemini_tool_call_invoke, + ) + + tool_calls = [ + {"id": "call_custom", "type": "custom", "custom": {"name": "noop", "input": ""}} + ] + _parallel_tool_calls(None, None) + + gemini_parts = convert_to_gemini_tool_call_invoke( + {"role": "assistant", "content": None, "tool_calls": tool_calls}, + model="gemini-3-pro-preview", + ) + + assert len(gemini_parts) == 2 + assert gemini_parts[0]["thoughtSignature"] == PLACEHOLDER_SIGNATURE + assert "thoughtSignature" not in gemini_parts[1] + + +def test_no_placeholder_when_model_is_unknown(): + """Without a model there is nothing to prove the target needs a placeholder, so none is added.""" + from litellm.litellm_core_utils.prompt_templates.factory import ( + convert_to_gemini_tool_call_invoke, + ) + + gemini_parts = convert_to_gemini_tool_call_invoke( + { + "role": "assistant", + "content": None, + "tool_calls": _parallel_tool_calls(None, None), + }, + ) + + assert len(gemini_parts) == 2 + assert all("thoughtSignature" not in part for part in gemini_parts) + + +def test_real_signature_forwarded_to_gemini_2_5_without_placeholder_siblings(): + """Older models still receive a real signature that a client replays, and still get no placeholder.""" + from litellm.litellm_core_utils.prompt_templates.factory import ( + convert_to_gemini_tool_call_invoke, + ) + + gemini_parts = convert_to_gemini_tool_call_invoke( + { + "role": "assistant", + "content": None, + "tool_calls": _parallel_tool_calls(REAL_THOUGHT_SIGNATURE, None), + }, + model="gemini-2.5-flash", + ) + + assert len(gemini_parts) == 2 + assert gemini_parts[0]["thoughtSignature"] == REAL_THOUGHT_SIGNATURE + assert "thoughtSignature" not in gemini_parts[1] + + +def test_parallel_tool_call_history_replayed_through_full_message_conversion(): + """End to end through the message-history converter, the path a real /chat/completions replay takes.""" + from litellm.llms.vertex_ai.gemini.transformation import ( + _gemini_convert_messages_with_history, + ) + + messages = [ + {"role": "user", "content": "Weather in Paris, London and Tokyo?"}, + { + "role": "assistant", + "content": None, + "tool_calls": _parallel_tool_calls_signed_via_id( + REAL_THOUGHT_SIGNATURE, None, None + ), + }, + ] + + contents = _gemini_convert_messages_with_history( + messages=messages, model="gemini-3-pro-preview" + ) + + model_parts = contents[1]["parts"] + assert len(model_parts) == 3 + assert model_parts[0]["thoughtSignature"] == REAL_THOUGHT_SIGNATURE + assert "thoughtSignature" not in model_parts[1] + assert "thoughtSignature" not in model_parts[2] + + +@pytest.mark.parametrize( + "model", + ["gemini-3.5-flash", "vertex_ai/gemini-3.5-flash", "gemini/gemini-3.5-flash"], +) +def test_natively_signed_parallel_turn_never_carries_a_placeholder(model): + """A native gemini-3.5 parallel turn replays with zero skip_thought_signature_validator parts. + + Fabricating the placeholder alongside a real signature is what produced empty text responses + on gemini-3.5 parallel function calling, so the whole payload has to stay placeholder-free. + """ + import json + + from litellm.llms.vertex_ai.gemini.transformation import ( + _gemini_convert_messages_with_history, + ) + + messages = [ + {"role": "user", "content": "Weather in Paris, London and Tokyo?"}, + { + "role": "assistant", + "content": None, + "tool_calls": _parallel_tool_calls_signed_via_id( + REAL_THOUGHT_SIGNATURE, None, None + ), + }, + ] + + contents = _gemini_convert_messages_with_history(messages=messages, model=model) + + model_parts = contents[1]["parts"] + assert len(model_parts) == 3 + assert model_parts[0]["thoughtSignature"] == REAL_THOUGHT_SIGNATURE + assert "thoughtSignature" not in model_parts[1] + assert "thoughtSignature" not in model_parts[2] + assert PLACEHOLDER_SIGNATURE not in json.dumps(contents) + + +@pytest.mark.parametrize( + "model", + [ + "gemini-3-pro-preview", + "gemini-3-flash-preview", + "gemini-3.1-pro-preview", + "gemini-3.5-flash", + "gemini-3.6-flash", + "gemini-3.7-flash", + "vertex_ai/gemini-3.5-flash", + "vertex_ai/gemini-3.7-flash", + "gemini/gemini-3.5-flash", + "gemini/gemini-3.7-flash", + ], +) +def test_placeholder_scoped_to_first_call_across_gemini_3_variants(model): + """The gemini-3 gate is a substring match, so every family member and prefix form has to + land on the same one-placeholder budget rather than only the versions we happened to try.""" + from litellm.litellm_core_utils.prompt_templates.factory import ( + convert_to_gemini_tool_call_invoke, + ) + + gemini_parts = convert_to_gemini_tool_call_invoke( + { + "role": "assistant", + "content": None, + "tool_calls": _parallel_tool_calls(None, None, None), + }, + model=model, + ) + + assert len(gemini_parts) == 3 + assert gemini_parts[0]["thoughtSignature"] == PLACEHOLDER_SIGNATURE + assert "thoughtSignature" not in gemini_parts[1] + assert "thoughtSignature" not in gemini_parts[2] + + +def test_signed_text_part_survives_alongside_unsigned_parallel_tool_calls(): + """Text-part and function-call signatures are collected by separate code paths, so scoping the + placeholder must not disturb a real signature that arrived on the text part.""" + from litellm.llms.vertex_ai.gemini.transformation import ( + _gemini_convert_messages_with_history, + ) + + msg = { + "role": "assistant", + "content": "Checking all three cities.", + "provider_specific_fields": {"thought_signatures": ["real_25_signature"]}, + "tool_calls": _parallel_tool_calls(None, None, None), + } + + parts = _gemini_convert_messages_with_history( + messages=[msg], model="gemini-3-pro-preview" + )[0]["parts"] + + assert parts[0]["text"] == "Checking all three cities." + assert parts[0]["thoughtSignature"] == "real_25_signature" + assert parts[1]["thoughtSignature"] == PLACEHOLDER_SIGNATURE + assert "thoughtSignature" not in parts[2] + assert "thoughtSignature" not in parts[3] + + # Tests for media_resolution (detail parameter) handling - Issue #17084 class TestMediaResolution: """Tests for media_resolution handling in Gemini 2.x models""" diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py index 51cc2857252..b7265ed62e9 100644 --- a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py +++ b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py @@ -5246,11 +5246,14 @@ def test_mid_stream_429_error_raises_during_iteration(): # Iterate the stream: first chunks should succeed, then 429 error should be raised results = [] - with pytest.raises(VertexAIError) as exc_info: + def _drain(): for chunk in streaming_obj: if chunk is not None: results.append(chunk) + with pytest.raises(VertexAIError) as exc_info: + _drain() + # Verify: received normal chunks before the error assert ( len(results) >= 1 diff --git a/tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py b/tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py index b83d4742b64..c189cdd0ea7 100644 --- a/tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py +++ b/tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py @@ -33,7 +33,7 @@ def test_validate_vertex_location_accepts_valid(location): ["attacker.example/", "evil.com#", "us.attacker.example", "us/../..", "US", "us_central1", "-us", "", None], ) def test_validate_vertex_location_rejects_invalid(location): - with pytest.raises(ValueError): + with pytest.raises(ValueError, match="vertex_location is required|Invalid vertex_location format"): validate_vertex_location(location) diff --git a/tests/test_litellm/llms/volcengine/responses/test_volcengine_responses_transformation.py b/tests/test_litellm/llms/volcengine/responses/test_volcengine_responses_transformation.py index 891d1c15c61..7922331d19f 100644 --- a/tests/test_litellm/llms/volcengine/responses/test_volcengine_responses_transformation.py +++ b/tests/test_litellm/llms/volcengine/responses/test_volcengine_responses_transformation.py @@ -137,7 +137,7 @@ class TestVolcengineResponsesAPITransformation: monkeypatch.delenv("ARK_API_KEY", raising=False) monkeypatch.delenv("VOLCENGINE_API_KEY", raising=False) - with pytest.raises(ValueError): + with pytest.raises(ValueError, match='Volcengine API key is required\\. Set ARK_API_KEY /'): config.validate_environment(headers={}, model="volcengine/demo", litellm_params={}) def test_unsupported_params_are_dropped_with_extra_body(self): diff --git a/tests/test_litellm/llms/volcengine/test_volcengine_embedding.py b/tests/test_litellm/llms/volcengine/test_volcengine_embedding.py index 6a035bcd7f0..1670dac0e9d 100644 --- a/tests/test_litellm/llms/volcengine/test_volcengine_embedding.py +++ b/tests/test_litellm/llms/volcengine/test_volcengine_embedding.py @@ -198,10 +198,11 @@ def test_volcengine_embedding_error_scenarios(): mock_embedding.side_effect = ValueError("Unsupported encoding_format") # Test that errors are properly raised - with pytest.raises(Exception) as exc_info: - test_params = { - k: v for k, v in scenario.items() if k != "expected_error_pattern" - } + test_params = { + k: v for k, v in scenario.items() if k != "expected_error_pattern" + } + + with pytest.raises(Exception, match=f"(?i){scenario['expected_error_pattern']}") as exc_info: litellm.embedding(input=["test"], **test_params) # Verify error message contains expected pattern diff --git a/tests/test_litellm/llms/voyage/rerank/test_voyage_rerank_transformation.py b/tests/test_litellm/llms/voyage/rerank/test_voyage_rerank_transformation.py index 8f99609e3f5..f466b7e19b5 100644 --- a/tests/test_litellm/llms/voyage/rerank/test_voyage_rerank_transformation.py +++ b/tests/test_litellm/llms/voyage/rerank/test_voyage_rerank_transformation.py @@ -227,7 +227,7 @@ class TestVoyageRerankTransform: mock_logging = MagicMock() model_response = RerankResponse() - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Unauthorized') as exc_info: self.config.transform_rerank_response( model=self.model, raw_response=mock_response, @@ -248,7 +248,7 @@ class TestVoyageRerankTransform: mock_logging = MagicMock() model_response = RerankResponse() - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Failed to parse response: Invalid JSON response') as exc_info: self.config.transform_rerank_response( model=self.model, raw_response=mock_response, diff --git a/tests/test_litellm/llms/voyage/test_voyage_multimodal_embedding.py b/tests/test_litellm/llms/voyage/test_voyage_multimodal_embedding.py index f283e7fe0df..f3e6885cbe6 100644 --- a/tests/test_litellm/llms/voyage/test_voyage_multimodal_embedding.py +++ b/tests/test_litellm/llms/voyage/test_voyage_multimodal_embedding.py @@ -195,7 +195,7 @@ class TestVoyageMultimodalEmbeddings: monkeypatch.setattr(module, "get_secret_str", lambda name: None) config = VoyageMultimodalEmbeddingConfig() - with pytest.raises(ValueError) as exc_info: + with pytest.raises(ValueError, match='Voyage API key is required for multimodal embeddings\\. Set') as exc_info: config.validate_environment( {}, "voyage-multimodal-3.5", [], {}, {}, api_key=None ) @@ -207,7 +207,7 @@ class TestVoyageMultimodalEmbeddings: ) config = VoyageMultimodalEmbeddingConfig() - with pytest.raises(ValueError) as exc_info: + with pytest.raises(ValueError, match='Voyage multimodal embeddings require a non-empty') as exc_info: config._normalize_content_item({"type": "image_url", "image_url": {}}) assert "image_url" in str(exc_info.value) diff --git a/tests/test_litellm/llms/xai/test_xai_key_fallback.py b/tests/test_litellm/llms/xai/test_xai_key_fallback.py index 4c769c572ac..ec3eb83309c 100644 --- a/tests/test_litellm/llms/xai/test_xai_key_fallback.py +++ b/tests/test_litellm/llms/xai/test_xai_key_fallback.py @@ -168,7 +168,7 @@ def test_responses_config_raises_when_no_key_is_available(monkeypatch): monkeypatch.setattr(litellm, "api_key", None) monkeypatch.delenv("XAI_API_KEY", raising=False) - with pytest.raises(ValueError) as exc_info: + with pytest.raises(ValueError, match='XAI API key is required\\. Set api_key, litellm\\.xai_key') as exc_info: XAIResponsesAPIConfig().validate_environment({}, "xai/grok-3-mini", None) error_message = str(exc_info.value) diff --git a/tests/test_litellm/llms/xai/test_xai_oauth.py b/tests/test_litellm/llms/xai/test_xai_oauth.py index 45fa6a405f2..3fc4c35052e 100644 --- a/tests/test_litellm/llms/xai/test_xai_oauth.py +++ b/tests/test_litellm/llms/xai/test_xai_oauth.py @@ -556,7 +556,7 @@ def test_get_llm_provider_uses_single_xai_provider(monkeypatch): def test_xai_oauth_alias_is_not_a_provider(): - with pytest.raises(Exception): + with pytest.raises(litellm.BadRequestError): get_llm_provider("xai_oauth/grok-4") diff --git a/tests/test_litellm/models/test_models.py b/tests/test_litellm/models/test_models.py index 187c7aa7f5e..669dba8e466 100644 --- a/tests/test_litellm/models/test_models.py +++ b/tests/test_litellm/models/test_models.py @@ -39,6 +39,7 @@ from litellm.models.verification_token import ( LiteLLM_DeletedVerificationToken, LiteLLM_VerificationToken, ) +from pydantic import ValidationError class TestBudget: @@ -421,7 +422,7 @@ class TestBudgetTableFull: assert budget.max_budget == 10.0 def test_full_requires_created_at(self): - with pytest.raises(Exception): + with pytest.raises(ValidationError): LiteLLM_BudgetTableFull(budget_id="b1") @@ -480,7 +481,7 @@ class TestMCPServerTable: assert server.env == {} def test_mcp_server_requires_transport(self): - with pytest.raises(Exception): + with pytest.raises(ValidationError): LiteLLM_MCPServerTable(server_id="s1") @@ -538,7 +539,7 @@ class TestManagedTables: assert table.flat_model_file_ids == ["file-abc"] def test_managed_object_table_requires_purpose(self): - with pytest.raises(Exception): + with pytest.raises(ValidationError): LiteLLM_ManagedObjectTable( unified_object_id="o1", model_object_id="m1", file_object={} ) diff --git a/tests/test_litellm/passthrough/test_async_streaming_error_propagation.py b/tests/test_litellm/passthrough/test_async_streaming_error_propagation.py index d262063584b..faf4ea46c43 100644 --- a/tests/test_litellm/passthrough/test_async_streaming_error_propagation.py +++ b/tests/test_litellm/passthrough/test_async_streaming_error_propagation.py @@ -65,7 +65,7 @@ async def test_async_streaming_429_raises(): return mock_response chunks = [] - with pytest.raises(httpx.HTTPStatusError) as exc_info: + async def _drain(): async for chunk in _async_streaming( response=response_coro(), litellm_logging_obj=_make_mock_logging_obj(), @@ -73,6 +73,9 @@ async def test_async_streaming_429_raises(): ): chunks.append(chunk) + with pytest.raises(httpx.HTTPStatusError) as exc_info: + await _drain() + assert exc_info.value.response.status_code == 429 assert len(chunks) == 0 diff --git a/tests/test_litellm/passthrough/test_passthrough_main.py b/tests/test_litellm/passthrough/test_passthrough_main.py index 58a0185ea8c..965f9fd8f7d 100644 --- a/tests/test_litellm/passthrough/test_passthrough_main.py +++ b/tests/test_litellm/passthrough/test_passthrough_main.py @@ -721,10 +721,13 @@ async def test_allm_passthrough_route_429_streaming_raises(): # result is an async generator — consuming it must raise, not silently yield error bytes chunks = [] - with pytest.raises(httpx.HTTPStatusError) as exc_info: + async def _drain(): async for chunk in result: # type: ignore[union-attr] chunks.append(chunk) + with pytest.raises(httpx.HTTPStatusError) as exc_info: + await _drain() + assert exc_info.value.response.status_code == 429 assert len(chunks) == 0, "No chunks should be yielded before the 429 raises" diff --git a/tests/test_litellm/passthrough/test_streaming_interrupt_spend_tracking.py b/tests/test_litellm/passthrough/test_streaming_interrupt_spend_tracking.py index f3fe3ae5c38..3783e218e4e 100644 --- a/tests/test_litellm/passthrough/test_streaming_interrupt_spend_tracking.py +++ b/tests/test_litellm/passthrough/test_streaming_interrupt_spend_tracking.py @@ -164,7 +164,7 @@ async def test_async_streaming_flushes_on_upstream_exception_with_partial_data() provider_config = MagicMock() received = [] - with pytest.raises(httpx.ReadError): + async def _drain(): async for chunk in _async_streaming( response=response_coro(), litellm_logging_obj=mock_logging_obj, @@ -172,6 +172,9 @@ async def test_async_streaming_flushes_on_upstream_exception_with_partial_data() ): received.append(chunk) + with pytest.raises(httpx.ReadError): + await _drain() + assert received == partial_chunks await asyncio.sleep(0) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py index 0209abee510..d5936b2ae86 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py @@ -8185,7 +8185,7 @@ class TestGetUserObjectPermission: return_value=None, ), ): - with pytest.raises(ValueError): + with pytest.raises(ValueError, match="user 'human-dangling' names object_permission_id"): await MCPRequestHandler._get_user_object_permission(auth) async def test_no_user_id_places_no_ceiling(self): diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py index 34852850de6..b4d3782ba43 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py @@ -2695,13 +2695,6 @@ async def test_token_endpoint_respects_x_forwarded_host(): "443", "https://internal.local", ), - ( - "http://localhost:4000/", - "https", - "proxy.example.com", - "8443", - "https://proxy.example.com:8443", - ), ( "http://localhost:4000/", "https", diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py index 051df30dfcb..aa6e63b7ca0 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py @@ -4440,7 +4440,7 @@ async def test_call_mcp_tool_logs_failure_via_post_call_failure_hook(): proxy_logging_mock, ), ): - with pytest.raises(Exception): + with pytest.raises(Exception, match="boom"): await call_mcp_tool( name="test_server-any_tool", arguments={"x": 1}, diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_short_mcp_tool_prefix.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_short_mcp_tool_prefix.py index 6e3ac014840..941e5deee93 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_short_mcp_tool_prefix.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_short_mcp_tool_prefix.py @@ -79,7 +79,7 @@ class TestShortPrefixHelpers: assert compute_short_server_prefix("abc") != compute_short_server_prefix("abd") def test_short_prefix_requires_server_id(self): - with pytest.raises(ValueError): + with pytest.raises(ValueError, match='compute_short_server_prefix requires a non-empty server_id'): compute_short_server_prefix("") def test_flag_defaults_to_false(self): diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index ccbf00a67b9..762d2cbf3c7 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -3,8 +3,12 @@ import json import os import sys from types import SimpleNamespace +from typing import TYPE_CHECKING, Optional from unittest.mock import AsyncMock, MagicMock, patch +if TYPE_CHECKING: + from litellm.router import Router + sys.path.insert( 0, os.path.abspath("../../..") ) # Adds the parent directory to the system path @@ -264,7 +268,7 @@ def test_get_experimental_ui_login_jwt_auth_token_invalid( invalid_sso_user_defined_values, ): """Test generating JWT token with missing user role""" - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='User role is required for experimental UI login') as exc_info: ExperimentalUIJWTToken.get_experimental_ui_login_jwt_auth_token( invalid_sso_user_defined_values ) @@ -879,7 +883,7 @@ async def test_get_user_object_wraps_db_outage_as_valueerror_preserving_context( mock_cache.async_set_cache = AsyncMock() with patch("litellm.proxy.auth.auth_checks._should_check_db", return_value=True): - with pytest.raises(ValueError) as exc_info: + with pytest.raises(ValueError, match="User doesn't exist in db\\.") as exc_info: await get_user_object( user_id="outage-contract-probe-user", prisma_client=mock_prisma_client, @@ -6612,3 +6616,284 @@ def test_can_object_call_model_team_scoped_wildcard_accepts_bare_model_name(): ) is True ) + + +UNPRICED_UNDERLYING_MODEL = "openai/unpriced-model-lit4984-xyz" + + +def _router_with_priced_and_unpriced_models() -> "Router": + from litellm.router import Router + + return Router( + model_list=[ + { + "model_name": "priced-group", + "litellm_params": {"model": "gpt-3.5-turbo", "api_key": "sk-test"}, + }, + { + "model_name": "unpriced-group", + "litellm_params": {"model": UNPRICED_UNDERLYING_MODEL, "api_key": "sk-test"}, + }, + ] + ) + + +def test_model_has_no_cost_mapping_priced_model_is_false(): + from litellm.proxy.auth.auth_checks import model_has_no_cost_mapping + + router = _router_with_priced_and_unpriced_models() + + assert model_has_no_cost_mapping(model="priced-group", llm_router=router) is False + + +def test_model_has_no_cost_mapping_unpriced_model_is_true(): + from litellm.proxy.auth.auth_checks import model_has_no_cost_mapping + + router = _router_with_priced_and_unpriced_models() + + assert model_has_no_cost_mapping(model="unpriced-group", llm_router=router) is True + + +def test_model_has_no_cost_mapping_no_model_or_router_is_false(): + from litellm.proxy.auth.auth_checks import model_has_no_cost_mapping + + router = _router_with_priced_and_unpriced_models() + + assert model_has_no_cost_mapping(model=None, llm_router=router) is False + assert model_has_no_cost_mapping(model="unpriced-group", llm_router=None) is False + + +@pytest.mark.parametrize( + "underlying_model", + [ + "azure/speech/azure-tts", + "mistral/mistral-ocr-latest", + "vertex_ai/imagen-3.0-generate-001", + "dashscope/qwen-flash", + ], +) +def test_model_has_no_cost_mapping_non_token_priced_model_is_false(underlying_model): + from litellm.proxy.auth.auth_checks import model_has_no_cost_mapping + from litellm.router import Router + + router = Router( + model_list=[ + { + "model_name": "non-token-priced-group", + "litellm_params": {"model": underlying_model, "api_key": "sk-test"}, + } + ] + ) + + assert model_has_no_cost_mapping(model="non-token-priced-group", llm_router=router) is False + + +def test_model_has_no_cost_mapping_non_token_price_from_litellm_params_is_false(): + from litellm.proxy.auth.auth_checks import model_has_no_cost_mapping + from litellm.router import Router + + router = Router( + model_list=[ + { + "model_name": "custom-tts", + "litellm_params": { + "model": f"{UNPRICED_UNDERLYING_MODEL}-per-second", + "api_key": "sk-test", + "input_cost_per_second": 0.0001, + }, + } + ] + ) + + assert model_has_no_cost_mapping(model="custom-tts", llm_router=router) is False + + +@pytest.mark.parametrize("cost_field", ["input_cost_per_second", "input_cost_per_token"]) +def test_model_has_no_cost_mapping_explicit_zero_price_is_false(cost_field): + from litellm.proxy.auth.auth_checks import model_has_no_cost_mapping + from litellm.router import Router + + router = Router( + model_list=[ + { + "model_name": "free-group", + "litellm_params": { + "model": f"{UNPRICED_UNDERLYING_MODEL}-{cost_field}", + "api_key": "sk-test", + cost_field: 0, + }, + } + ] + ) + + assert model_has_no_cost_mapping(model="free-group", llm_router=router) is False + + +def test_model_has_no_cost_mapping_tiered_pricing_only_is_false(): + from litellm.proxy.auth.auth_checks import model_has_no_cost_mapping + from litellm.router import Router + + router = Router( + model_list=[ + { + "model_name": "tiered-group", + "litellm_params": { + "model": f"{UNPRICED_UNDERLYING_MODEL}-tiered", + "api_key": "sk-test", + "tiered_pricing": [ + {"range": [0, 128000], "input_cost_per_token": 2e-7, "output_cost_per_token": 6e-7}, + {"range": [128000, 256000], "input_cost_per_token": 4e-7, "output_cost_per_token": 12e-7}, + ], + }, + } + ] + ) + + assert model_has_no_cost_mapping(model="tiered-group", llm_router=router) is False + + +async def _run_common_checks( + model: Optional[str], llm_router: Optional["Router"], route: str = "/chat/completions" +) -> bool: + from fastapi import Request + + from litellm.proxy.auth.auth_checks import common_checks + + return await common_checks( + request_body={"model": model, "messages": [{"role": "user", "content": "hi"}]}, + team_object=None, + user_object=None, + end_user_object=None, + global_proxy_spend=None, + general_settings={}, + route=route, + llm_router=llm_router, + proxy_logging_obj=MagicMock(), + valid_token=UserAPIKeyAuth(token="test-token"), + request=MagicMock(spec=Request), + ) + + +@pytest.mark.asyncio +async def test_common_checks_blocks_unpriced_model_when_enabled(monkeypatch): + monkeypatch.setattr(litellm, "block_requests_for_models_without_pricing", True) + router = _router_with_priced_and_unpriced_models() + + with pytest.raises(ProxyException) as exc_info: + await _run_common_checks(model="unpriced-group", llm_router=router) + + assert exc_info.value.code == "403" + assert exc_info.value.type == ProxyErrorTypes.model_cost_map_missing + assert exc_info.value.param == "model" + assert "unpriced-group" in exc_info.value.message + assert "pricing" in exc_info.value.message.lower() + + +@pytest.mark.asyncio +async def test_common_checks_allows_unpriced_model_when_disabled(monkeypatch): + monkeypatch.setattr(litellm, "block_requests_for_models_without_pricing", False) + router = _router_with_priced_and_unpriced_models() + + result = await _run_common_checks(model="unpriced-group", llm_router=router) + + assert result is True + + +@pytest.mark.asyncio +async def test_common_checks_allows_priced_model_when_enabled(monkeypatch): + monkeypatch.setattr(litellm, "block_requests_for_models_without_pricing", True) + router = _router_with_priced_and_unpriced_models() + + result = await _run_common_checks(model="priced-group", llm_router=router) + + assert result is True + + +@pytest.mark.asyncio +async def test_common_checks_ignores_non_llm_route_when_enabled(monkeypatch): + monkeypatch.setattr(litellm, "block_requests_for_models_without_pricing", True) + router = _router_with_priced_and_unpriced_models() + + result = await _run_common_checks( + model="unpriced-group", llm_router=router, route="/model/new" + ) + + assert result is True + + +@pytest.mark.asyncio +async def test_common_checks_blocks_alias_resolving_to_unpriced_model(monkeypatch): + from litellm.router import Router + + monkeypatch.setattr(litellm, "block_requests_for_models_without_pricing", True) + router = Router( + model_list=[ + { + "model_name": "billed-underlying-group", + "litellm_params": {"model": UNPRICED_UNDERLYING_MODEL, "api_key": "sk-test"}, + } + ], + model_group_alias={"public-alias": "billed-underlying-group"}, + ) + + with pytest.raises(ProxyException) as exc_info: + await _run_common_checks(model="public-alias", llm_router=router) + + assert exc_info.value.code == "403" + assert exc_info.value.type == ProxyErrorTypes.model_cost_map_missing + assert "public-alias" in exc_info.value.message + + +@pytest.mark.asyncio +async def test_common_checks_blocks_comma_separated_request_carrying_an_unpriced_model(monkeypatch): + monkeypatch.setattr(litellm, "block_requests_for_models_without_pricing", True) + router = _router_with_priced_and_unpriced_models() + + with pytest.raises(ProxyException) as exc_info: + await _run_common_checks(model="priced-group,unpriced-group", llm_router=router) + + assert exc_info.value.code == "403" + assert exc_info.value.type == ProxyErrorTypes.model_cost_map_missing + assert "'unpriced-group'" in exc_info.value.message + assert "'priced-group'" not in exc_info.value.message + + +@pytest.mark.asyncio +async def test_common_checks_allows_comma_separated_request_when_every_model_is_priced(monkeypatch): + monkeypatch.setattr(litellm, "block_requests_for_models_without_pricing", True) + router = _router_with_priced_and_unpriced_models() + + result = await _run_common_checks(model="priced-group,priced-group", llm_router=router) + + assert result is True + + +def _router_with_a_group_priced_through_model_info() -> "Router": + from litellm.router import Router + + return Router( + model_list=[ + { + "model_name": "model-info-priced-group", + "litellm_params": {"model": f"{UNPRICED_UNDERLYING_MODEL}-model-info", "api_key": "sk-test"}, + "model_info": {"input_cost_per_token": 0, "output_cost_per_token": 0}, + } + ], + model_group_alias={"model-info-priced-alias": "model-info-priced-group"}, + ) + + +def test_model_has_no_cost_mapping_group_priced_through_model_info_is_false(): + from litellm.proxy.auth.auth_checks import model_has_no_cost_mapping + + router = _router_with_a_group_priced_through_model_info() + + assert model_has_no_cost_mapping(model="model-info-priced-group", llm_router=router) is False + + +def test_model_has_no_cost_mapping_alias_to_a_group_priced_through_model_info_is_false(): + from litellm.proxy.auth.auth_checks import model_has_no_cost_mapping + + router = _router_with_a_group_priced_through_model_info() + + assert model_has_no_cost_mapping(model="model-info-priced-alias", llm_router=router) is False diff --git a/tests/test_litellm/proxy/auth/test_auth_exception_handler.py b/tests/test_litellm/proxy/auth/test_auth_exception_handler.py index b4725a81823..721857e5411 100644 --- a/tests/test_litellm/proxy/auth/test_auth_exception_handler.py +++ b/tests/test_litellm/proxy/auth/test_auth_exception_handler.py @@ -338,12 +338,13 @@ async def test_handle_authentication_error_budget_exceeded(): mock_api_key = "test-key" # Test with budget exceeded error - with pytest.raises(ProxyException) as exc_info: - from litellm.exceptions import BudgetExceededError + from litellm.exceptions import BudgetExceededError - budget_error = BudgetExceededError( - message="Budget exceeded", current_cost=100, max_budget=100 - ) + budget_error = BudgetExceededError( + message="Budget exceeded", current_cost=100, max_budget=100 + ) + + with pytest.raises(ProxyException) as exc_info: await handler._handle_authentication_error( budget_error, mock_request, diff --git a/tests/test_litellm/proxy/auth/test_auth_hot_path_network_requests.py b/tests/test_litellm/proxy/auth/test_auth_hot_path_network_requests.py index 22752f767ce..6b2d2babedc 100644 --- a/tests/test_litellm/proxy/auth/test_auth_hot_path_network_requests.py +++ b/tests/test_litellm/proxy/auth/test_auth_hot_path_network_requests.py @@ -1,521 +1,521 @@ -""" -Test to count and track the number of network requests (DB queries, cache lookups) -made on the hot path for keys that have team_id and user_id attached. - -This test ensures we don't regress on the number of network requests made during -request authentication, which directly impacts proxy latency. - -The hot path covers auth functions called on every LLM API request: -- get_key_object: lookup the API key -- get_team_object: lookup the team (for keys with team_id) -- get_user_object: lookup the user (for keys with user_id) -- get_team_membership: lookup team member budget (when team_member_spend set) - -Each function does: cache read -> (on miss) DB query -> cache write. -We count these to catch regressions in the number of network requests. - -NOTE: This test does NOT require proxy extras (apscheduler, etc.) because -it tests at the auth_checks level, not the full proxy_server level. -""" - -import os -import sys -import time -from typing import Any, Dict, List, Optional -from unittest.mock import AsyncMock, MagicMock - -import pytest - -sys.path.insert(0, os.path.abspath("../../..")) - -from litellm.caching.dual_cache import DualCache -from litellm.caching.in_memory_cache import InMemoryCache -from litellm.proxy._types import ( - LiteLLM_TeamTableCachedObj, - LiteLLM_UserTable, - LitellmUserRoles, - UserAPIKeyAuth, - LiteLLM_TeamMembership, - hash_token, -) -from litellm.proxy.auth.auth_checks import ( - get_key_object, - get_team_membership, - get_team_object, - get_user_object, -) - - -class CacheCallTracker: - """ - Tracks cache read/write operations by wrapping DualCache methods. - This is used to count network-level operations on the hot path. - """ - - def __init__(self): - self.cache_reads: List[Dict[str, Any]] = [] - self.cache_writes: List[Dict[str, Any]] = [] - self.db_queries: List[Dict[str, Any]] = [] - - def get_summary(self) -> Dict[str, Any]: - return { - "total_cache_reads": len(self.cache_reads), - "total_cache_writes": len(self.cache_writes), - "total_db_queries": len(self.db_queries), - "total_network_requests": len(self.cache_reads) - + len(self.cache_writes) - + len(self.db_queries), - "cache_read_keys": [r["key"] for r in self.cache_reads], - "cache_write_keys": [w["key"] for w in self.cache_writes], - "db_query_details": self.db_queries, - } - - -def _wrap_cache_with_tracker(cache: DualCache, tracker: CacheCallTracker) -> DualCache: - """Wrap a DualCache to track all reads and writes.""" - original_async_get = cache.async_get_cache - original_async_set = cache.async_set_cache - - async def tracked_async_get(key, *args, **kwargs): - result = await original_async_get(key, *args, **kwargs) - tracker.cache_reads.append( - {"key": key, "hit": result is not None, "method": "async_get_cache"} - ) - return result - - async def tracked_async_set(key, value, *args, **kwargs): - tracker.cache_writes.append({"key": key, "method": "async_set_cache"}) - return await original_async_set(key, value, *args, **kwargs) - - cache.async_get_cache = tracked_async_get - cache.async_set_cache = tracked_async_set - return cache - - -def _create_valid_token( - api_key: str, - team_id: str, - user_id: str, - has_team_member_spend: bool = False, - org_id: Optional[str] = None, -) -> UserAPIKeyAuth: - """Create a UserAPIKeyAuth with team_id and user_id set.""" - hashed = hash_token(api_key) - return UserAPIKeyAuth( - token=hashed, - api_key=api_key, - team_id=team_id, - user_id=user_id, - org_id=org_id, - models=["gpt-4", "gpt-3.5-turbo"], - max_budget=100.0, - spend=10.0, - team_spend=50.0, - team_max_budget=1000.0, - team_models=["gpt-4", "gpt-3.5-turbo"], - team_member_spend=5.0 if has_team_member_spend else None, - last_refreshed_at=time.time(), - user_role=LitellmUserRoles.INTERNAL_USER, - ) - - -def _create_team_object(team_id: str) -> LiteLLM_TeamTableCachedObj: - """Create a team table object for caching.""" - return LiteLLM_TeamTableCachedObj( - team_id=team_id, - models=["gpt-4", "gpt-3.5-turbo"], - max_budget=1000.0, - spend=50.0, - tpm_limit=10000, - rpm_limit=100, - last_refreshed_at=time.time(), - ) - - -def _create_user_object(user_id: str) -> LiteLLM_UserTable: - """Create a user table object for caching.""" - return LiteLLM_UserTable( - user_id=user_id, - max_budget=500.0, - spend=25.0, - models=["gpt-4"], - tpm_limit=5000, - rpm_limit=50, - user_role=LitellmUserRoles.INTERNAL_USER, - user_email="test@example.com", - ) - - -# ============================================================================ -# TEST: get_key_object cache behavior -# ============================================================================ - - -@pytest.mark.asyncio -async def test_get_key_object_warm_cache(): - """ - Test get_key_object with a warm cache - should hit cache, no DB query. - """ - api_key = "sk-test-key-warm" - team_id = "team-123" - user_id = "user-456" - hashed_token = hash_token(api_key) - - valid_token = _create_valid_token(api_key, team_id, user_id) - - # Create cache with pre-populated data - cache = DualCache(in_memory_cache=InMemoryCache()) - await cache.async_set_cache(key=hashed_token, value=valid_token) - - # Track cache operations - tracker = CacheCallTracker() - tracked_cache = _wrap_cache_with_tracker(cache, tracker) - - # Mock prisma client (should NOT be called for warm cache) - mock_prisma = MagicMock() - mock_prisma.get_data = AsyncMock() - - result = await get_key_object( - hashed_token=hashed_token, - prisma_client=mock_prisma, - user_api_key_cache=tracked_cache, - parent_otel_span=None, - proxy_logging_obj=None, - ) - - summary = tracker.get_summary() - - # Should have exactly 1 cache read - assert summary["total_cache_reads"] == 1 - assert hashed_token in summary["cache_read_keys"] - - # Prisma should NOT have been called - mock_prisma.get_data.assert_not_called() - - # Result should be the cached token - assert result.token == hashed_token - - -@pytest.mark.asyncio -async def test_get_key_object_cold_cache(): - """ - Test get_key_object with a cold cache - should miss cache, query DB. - """ - api_key = "sk-test-key-cold" - team_id = "team-123" - user_id = "user-456" - hashed_token = hash_token(api_key) - - valid_token = _create_valid_token(api_key, team_id, user_id) - - # Create empty cache - cache = DualCache(in_memory_cache=InMemoryCache()) - - tracker = CacheCallTracker() - tracked_cache = _wrap_cache_with_tracker(cache, tracker) - - # Mock prisma client to return token on DB query - mock_prisma = MagicMock() - mock_prisma.get_data = AsyncMock(return_value=valid_token) - - await get_key_object( - hashed_token=hashed_token, - prisma_client=mock_prisma, - user_api_key_cache=tracked_cache, - parent_otel_span=None, - proxy_logging_obj=None, - ) - - summary = tracker.get_summary() - - # Should have 1 cache read (miss) and at least 1 cache write (populate cache) - assert summary["total_cache_reads"] >= 1 - - # Prisma SHOULD have been called - mock_prisma.get_data.assert_called_once() - - -# ============================================================================ -# TEST: get_team_object cache behavior -# ============================================================================ - - -@pytest.mark.asyncio -async def test_get_team_object_warm_cache(): - """ - Test get_team_object with a warm cache - should hit cache, no DB query. - """ - team_id = "team-warm-123" - team_obj = _create_team_object(team_id) - - cache = DualCache(in_memory_cache=InMemoryCache()) - cache_key = f"team_id:{team_id}" - await cache.async_set_cache(key=cache_key, value=team_obj) - - tracker = CacheCallTracker() - tracked_cache = _wrap_cache_with_tracker(cache, tracker) - - mock_prisma = MagicMock() - mock_prisma.db = MagicMock() - mock_prisma.db.litellm_teamtable = MagicMock() - mock_prisma.db.litellm_teamtable.find_unique = AsyncMock() - - await get_team_object( - team_id=team_id, - prisma_client=mock_prisma, - user_api_key_cache=tracked_cache, - parent_otel_span=None, - proxy_logging_obj=None, - ) - - summary = tracker.get_summary() - - assert summary["total_cache_reads"] >= 1 - assert cache_key in summary["cache_read_keys"] - - # DB should NOT have been called - mock_prisma.db.litellm_teamtable.find_unique.assert_not_called() - - -# ============================================================================ -# TEST: get_user_object cache behavior -# ============================================================================ - - -@pytest.mark.asyncio -async def test_get_user_object_warm_cache(): - """ - Test get_user_object with a warm cache - should hit cache, no DB query. - """ - user_id = "user-warm-456" - user_obj = _create_user_object(user_id) - - cache = DualCache(in_memory_cache=InMemoryCache()) - await cache.async_set_cache(key=user_id, value=user_obj) - - tracker = CacheCallTracker() - tracked_cache = _wrap_cache_with_tracker(cache, tracker) - - mock_prisma = MagicMock() - mock_prisma.db = MagicMock() - mock_prisma.db.litellm_usertable = MagicMock() - mock_prisma.db.litellm_usertable.find_unique = AsyncMock() - - await get_user_object( - user_id=user_id, - prisma_client=mock_prisma, - user_api_key_cache=tracked_cache, - parent_otel_span=None, - proxy_logging_obj=None, - user_id_upsert=False, - ) - - summary = tracker.get_summary() - - assert summary["total_cache_reads"] >= 1 - assert user_id in summary["cache_read_keys"] - - # DB should NOT have been called - mock_prisma.db.litellm_usertable.find_unique.assert_not_called() - - -# ============================================================================ -# TEST: get_team_membership cache behavior -# ============================================================================ - - -@pytest.mark.asyncio -async def test_get_team_membership_warm_cache(): - """ - Test get_team_membership with a warm cache - should hit cache, no DB query. - """ - user_id = "user-tm-456" - team_id = "team-tm-123" - - membership_dict = { - "user_id": user_id, - "team_id": team_id, - "spend": 3.0, - "budget_id": None, - "litellm_budget_table": None, - } - - cache = DualCache(in_memory_cache=InMemoryCache()) - # Cache key format used by get_team_membership - cache_key = f"team_membership:{user_id}:{team_id}" - await cache.async_set_cache(key=cache_key, value=membership_dict) - - tracker = CacheCallTracker() - tracked_cache = _wrap_cache_with_tracker(cache, tracker) - - mock_prisma = MagicMock() - mock_prisma.db = MagicMock() - mock_prisma.db.litellm_teammembership = MagicMock() - mock_prisma.db.litellm_teammembership.find_unique = AsyncMock() - - await get_team_membership( - user_id=user_id, - team_id=team_id, - prisma_client=mock_prisma, - user_api_key_cache=tracked_cache, - parent_otel_span=None, - proxy_logging_obj=None, - ) - - summary = tracker.get_summary() - - assert summary["total_cache_reads"] >= 1 - assert cache_key in summary["cache_read_keys"] - - # DB should NOT have been called - mock_prisma.db.litellm_teammembership.find_unique.assert_not_called() - - -# ============================================================================ -# TEST: Document duplicate team membership cache key issue -# ============================================================================ - - -@pytest.mark.asyncio -async def test_team_membership_cache_key_duplication(): - """ - Document the team membership duplicate cache key issue: - - Team membership is queried via TWO different cache keys: - 1. "{team_id}_{user_id}" - used in user_api_key_auth.py:1048 - 2. "team_membership:{user_id}:{team_id}" - used in auth_checks.py:960 (get_team_membership) - - This test documents that both keys refer to the same data but use different - cache key formats, potentially leading to duplicate lookups. - """ - user_id = "user-dup-456" - team_id = "team-dup-123" - - # The two different cache keys used for the same data - key_format_1 = f"{team_id}_{user_id}" # user_api_key_auth format - key_format_2 = f"team_membership:{user_id}:{team_id}" # auth_checks format - - _ = { - "user_id": user_id, - "team_id": team_id, - "spend": 3.0, - } - - # Document that these are different keys - assert ( - key_format_1 != key_format_2 - ), "Cache keys should be different (this is the bug)" - - # Document that these are different keys - assert ( - key_format_1 != key_format_2 - ), "Cache keys should be different (this is the bug)" - - -# ============================================================================ -# TEST: Full hot path network count summary -# ============================================================================ - - -@pytest.mark.asyncio -async def test_full_hot_path_network_count(): - """ - Summary test that counts all network operations when processing - a request with a key that has team_id and user_id attached. - - This test verifies the baseline number of cache operations expected - on a fully warm cache path. - """ - api_key = "sk-test-full-path" - team_id = "team-full-123" - user_id = "user-full-456" - hashed_token = hash_token(api_key) - - # Create all objects - valid_token = _create_valid_token( - api_key, team_id, user_id, has_team_member_spend=True - ) - team_obj = _create_team_object(team_id) - user_obj = _create_user_object(user_id) - membership_data = LiteLLM_TeamMembership( - user_id=user_id, - team_id=team_id, - spend=3.0, - budget_id=None, - litellm_budget_table=None, - ) - - # Pre-populate cache with all data - cache = DualCache(in_memory_cache=InMemoryCache()) - await cache.async_set_cache(key=hashed_token, value=valid_token) - await cache.async_set_cache(key=f"team_id:{team_id}", value=team_obj) - await cache.async_set_cache(key=user_id, value=user_obj) - await cache.async_set_cache( - key=f"team_membership:{user_id}:{team_id}", value=membership_data.model_dump() - ) - await cache.async_set_cache( - key=f"{team_id}_{user_id}", value=membership_data.model_dump() - ) - - # Create tracker AFTER populating cache - tracker = CacheCallTracker() - tracked_cache = _wrap_cache_with_tracker(cache, tracker) - - # Mock prisma (should not be called on warm cache) - mock_prisma = MagicMock() - - # Call each function to simulate the hot path - await get_key_object( - hashed_token=hashed_token, - prisma_client=mock_prisma, - user_api_key_cache=tracked_cache, - parent_otel_span=None, - proxy_logging_obj=None, - ) - - await get_team_object( - team_id=team_id, - prisma_client=mock_prisma, - user_api_key_cache=tracked_cache, - parent_otel_span=None, - proxy_logging_obj=None, - ) - - await get_user_object( - user_id=user_id, - prisma_client=mock_prisma, - user_api_key_cache=tracked_cache, - parent_otel_span=None, - proxy_logging_obj=None, - user_id_upsert=False, - ) - - await get_team_membership( - user_id=user_id, - team_id=team_id, - prisma_client=mock_prisma, - user_api_key_cache=tracked_cache, - parent_otel_span=None, - proxy_logging_obj=None, - ) - - summary = tracker.get_summary() - - # Assertions for expected baseline - # On warm cache: 4 reads (key, team, user, team_membership) - assert ( - summary["total_cache_reads"] == 4 - ), f"Expected 4 cache reads on warm path, got {summary['total_cache_reads']}" - - # No DB queries on warm cache - assert ( - summary["total_db_queries"] == 0 - ), f"Expected 0 DB queries on warm path, got {summary['total_db_queries']}" - - # Total network requests should be exactly 4 on warm cache - assert ( - summary["total_network_requests"] == 4 - ), f"Expected 4 total network requests on warm path, got {summary['total_network_requests']}" +""" +Test to count and track the number of network requests (DB queries, cache lookups) +made on the hot path for keys that have team_id and user_id attached. + +This test ensures we don't regress on the number of network requests made during +request authentication, which directly impacts proxy latency. + +The hot path covers auth functions called on every LLM API request: +- get_key_object: lookup the API key +- get_team_object: lookup the team (for keys with team_id) +- get_user_object: lookup the user (for keys with user_id) +- get_team_membership: lookup team member budget (when team_member_spend set) + +Each function does: cache read -> (on miss) DB query -> cache write. +We count these to catch regressions in the number of network requests. + +NOTE: This test does NOT require proxy extras (apscheduler, etc.) because +it tests at the auth_checks level, not the full proxy_server level. +""" + +import os +import sys +import time +from typing import Any, Dict, List, Optional +from unittest.mock import AsyncMock, MagicMock + +import pytest + +sys.path.insert(0, os.path.abspath("../../..")) + +from litellm.caching.dual_cache import DualCache +from litellm.caching.in_memory_cache import InMemoryCache +from litellm.proxy._types import ( + LiteLLM_TeamTableCachedObj, + LiteLLM_UserTable, + LitellmUserRoles, + UserAPIKeyAuth, + LiteLLM_TeamMembership, + hash_token, +) +from litellm.proxy.auth.auth_checks import ( + get_key_object, + get_team_membership, + get_team_object, + get_user_object, +) + + +class CacheCallTracker: + """ + Tracks cache read/write operations by wrapping DualCache methods. + This is used to count network-level operations on the hot path. + """ + + def __init__(self): + self.cache_reads: List[Dict[str, Any]] = [] + self.cache_writes: List[Dict[str, Any]] = [] + self.db_queries: List[Dict[str, Any]] = [] + + def get_summary(self) -> Dict[str, Any]: + return { + "total_cache_reads": len(self.cache_reads), + "total_cache_writes": len(self.cache_writes), + "total_db_queries": len(self.db_queries), + "total_network_requests": len(self.cache_reads) + + len(self.cache_writes) + + len(self.db_queries), + "cache_read_keys": [r["key"] for r in self.cache_reads], + "cache_write_keys": [w["key"] for w in self.cache_writes], + "db_query_details": self.db_queries, + } + + +def _wrap_cache_with_tracker(cache: DualCache, tracker: CacheCallTracker) -> DualCache: + """Wrap a DualCache to track all reads and writes.""" + original_async_get = cache.async_get_cache + original_async_set = cache.async_set_cache + + async def tracked_async_get(key, *args, **kwargs): + result = await original_async_get(key, *args, **kwargs) + tracker.cache_reads.append( + {"key": key, "hit": result is not None, "method": "async_get_cache"} + ) + return result + + async def tracked_async_set(key, value, *args, **kwargs): + tracker.cache_writes.append({"key": key, "method": "async_set_cache"}) + return await original_async_set(key, value, *args, **kwargs) + + cache.async_get_cache = tracked_async_get + cache.async_set_cache = tracked_async_set + return cache + + +def _create_valid_token( + api_key: str, + team_id: str, + user_id: str, + has_team_member_spend: bool = False, + org_id: Optional[str] = None, +) -> UserAPIKeyAuth: + """Create a UserAPIKeyAuth with team_id and user_id set.""" + hashed = hash_token(api_key) + return UserAPIKeyAuth( + token=hashed, + api_key=api_key, + team_id=team_id, + user_id=user_id, + org_id=org_id, + models=["gpt-4", "gpt-3.5-turbo"], + max_budget=100.0, + spend=10.0, + team_spend=50.0, + team_max_budget=1000.0, + team_models=["gpt-4", "gpt-3.5-turbo"], + team_member_spend=5.0 if has_team_member_spend else None, + last_refreshed_at=time.time(), + user_role=LitellmUserRoles.INTERNAL_USER, + ) + + +def _create_team_object(team_id: str) -> LiteLLM_TeamTableCachedObj: + """Create a team table object for caching.""" + return LiteLLM_TeamTableCachedObj( + team_id=team_id, + models=["gpt-4", "gpt-3.5-turbo"], + max_budget=1000.0, + spend=50.0, + tpm_limit=10000, + rpm_limit=100, + last_refreshed_at=time.time(), + ) + + +def _create_user_object(user_id: str) -> LiteLLM_UserTable: + """Create a user table object for caching.""" + return LiteLLM_UserTable( + user_id=user_id, + max_budget=500.0, + spend=25.0, + models=["gpt-4"], + tpm_limit=5000, + rpm_limit=50, + user_role=LitellmUserRoles.INTERNAL_USER, + user_email="test@example.com", + ) + + +# ============================================================================ +# TEST: get_key_object cache behavior +# ============================================================================ + + +@pytest.mark.asyncio +async def test_get_key_object_warm_cache(): + """ + Test get_key_object with a warm cache - should hit cache, no DB query. + """ + api_key = "sk-test-key-warm" + team_id = "team-123" + user_id = "user-456" + hashed_token = hash_token(api_key) + + valid_token = _create_valid_token(api_key, team_id, user_id) + + # Create cache with pre-populated data + cache = DualCache(in_memory_cache=InMemoryCache()) + await cache.async_set_cache(key=hashed_token, value=valid_token) + + # Track cache operations + tracker = CacheCallTracker() + tracked_cache = _wrap_cache_with_tracker(cache, tracker) + + # Mock prisma client (should NOT be called for warm cache) + mock_prisma = MagicMock() + mock_prisma.get_data = AsyncMock() + + result = await get_key_object( + hashed_token=hashed_token, + prisma_client=mock_prisma, + user_api_key_cache=tracked_cache, + parent_otel_span=None, + proxy_logging_obj=None, + ) + + summary = tracker.get_summary() + + # Should have exactly 1 cache read + assert summary["total_cache_reads"] == 1 + assert hashed_token in summary["cache_read_keys"] + + # Prisma should NOT have been called + mock_prisma.get_data.assert_not_called() + + # Result should be the cached token + assert result.token == hashed_token + + +@pytest.mark.asyncio +async def test_get_key_object_cold_cache(): + """ + Test get_key_object with a cold cache - should miss cache, query DB. + """ + api_key = "sk-test-key-cold" + team_id = "team-123" + user_id = "user-456" + hashed_token = hash_token(api_key) + + valid_token = _create_valid_token(api_key, team_id, user_id) + + # Create empty cache + cache = DualCache(in_memory_cache=InMemoryCache()) + + tracker = CacheCallTracker() + tracked_cache = _wrap_cache_with_tracker(cache, tracker) + + # Mock prisma client to return token on DB query + mock_prisma = MagicMock() + mock_prisma.get_data = AsyncMock(return_value=valid_token) + + await get_key_object( + hashed_token=hashed_token, + prisma_client=mock_prisma, + user_api_key_cache=tracked_cache, + parent_otel_span=None, + proxy_logging_obj=None, + ) + + summary = tracker.get_summary() + + # Should have 1 cache read (miss) and at least 1 cache write (populate cache) + assert summary["total_cache_reads"] >= 1 + + # Prisma SHOULD have been called + mock_prisma.get_data.assert_called_once() + + +# ============================================================================ +# TEST: get_team_object cache behavior +# ============================================================================ + + +@pytest.mark.asyncio +async def test_get_team_object_warm_cache(): + """ + Test get_team_object with a warm cache - should hit cache, no DB query. + """ + team_id = "team-warm-123" + team_obj = _create_team_object(team_id) + + cache = DualCache(in_memory_cache=InMemoryCache()) + cache_key = f"team_id:{team_id}" + await cache.async_set_cache(key=cache_key, value=team_obj) + + tracker = CacheCallTracker() + tracked_cache = _wrap_cache_with_tracker(cache, tracker) + + mock_prisma = MagicMock() + mock_prisma.db = MagicMock() + mock_prisma.db.litellm_teamtable = MagicMock() + mock_prisma.db.litellm_teamtable.find_unique = AsyncMock() + + await get_team_object( + team_id=team_id, + prisma_client=mock_prisma, + user_api_key_cache=tracked_cache, + parent_otel_span=None, + proxy_logging_obj=None, + ) + + summary = tracker.get_summary() + + assert summary["total_cache_reads"] >= 1 + assert cache_key in summary["cache_read_keys"] + + # DB should NOT have been called + mock_prisma.db.litellm_teamtable.find_unique.assert_not_called() + + +# ============================================================================ +# TEST: get_user_object cache behavior +# ============================================================================ + + +@pytest.mark.asyncio +async def test_get_user_object_warm_cache(): + """ + Test get_user_object with a warm cache - should hit cache, no DB query. + """ + user_id = "user-warm-456" + user_obj = _create_user_object(user_id) + + cache = DualCache(in_memory_cache=InMemoryCache()) + await cache.async_set_cache(key=user_id, value=user_obj) + + tracker = CacheCallTracker() + tracked_cache = _wrap_cache_with_tracker(cache, tracker) + + mock_prisma = MagicMock() + mock_prisma.db = MagicMock() + mock_prisma.db.litellm_usertable = MagicMock() + mock_prisma.db.litellm_usertable.find_unique = AsyncMock() + + await get_user_object( + user_id=user_id, + prisma_client=mock_prisma, + user_api_key_cache=tracked_cache, + parent_otel_span=None, + proxy_logging_obj=None, + user_id_upsert=False, + ) + + summary = tracker.get_summary() + + assert summary["total_cache_reads"] >= 1 + assert user_id in summary["cache_read_keys"] + + # DB should NOT have been called + mock_prisma.db.litellm_usertable.find_unique.assert_not_called() + + +# ============================================================================ +# TEST: get_team_membership cache behavior +# ============================================================================ + + +@pytest.mark.asyncio +async def test_get_team_membership_warm_cache(): + """ + Test get_team_membership with a warm cache - should hit cache, no DB query. + """ + user_id = "user-tm-456" + team_id = "team-tm-123" + + membership_dict = { + "user_id": user_id, + "team_id": team_id, + "spend": 3.0, + "budget_id": None, + "litellm_budget_table": None, + } + + cache = DualCache(in_memory_cache=InMemoryCache()) + # Cache key format used by get_team_membership + cache_key = f"team_membership:{user_id}:{team_id}" + await cache.async_set_cache(key=cache_key, value=membership_dict) + + tracker = CacheCallTracker() + tracked_cache = _wrap_cache_with_tracker(cache, tracker) + + mock_prisma = MagicMock() + mock_prisma.db = MagicMock() + mock_prisma.db.litellm_teammembership = MagicMock() + mock_prisma.db.litellm_teammembership.find_unique = AsyncMock() + + await get_team_membership( + user_id=user_id, + team_id=team_id, + prisma_client=mock_prisma, + user_api_key_cache=tracked_cache, + parent_otel_span=None, + proxy_logging_obj=None, + ) + + summary = tracker.get_summary() + + assert summary["total_cache_reads"] >= 1 + assert cache_key in summary["cache_read_keys"] + + # DB should NOT have been called + mock_prisma.db.litellm_teammembership.find_unique.assert_not_called() + + +# ============================================================================ +# TEST: Document duplicate team membership cache key issue +# ============================================================================ + + +@pytest.mark.asyncio +async def test_team_membership_cache_key_duplication(): + """ + Document the team membership duplicate cache key issue: + + Team membership is queried via TWO different cache keys: + 1. "{team_id}_{user_id}" - used in user_api_key_auth.py:1048 + 2. "team_membership:{user_id}:{team_id}" - used in auth_checks.py:960 (get_team_membership) + + This test documents that both keys refer to the same data but use different + cache key formats, potentially leading to duplicate lookups. + """ + user_id = "user-dup-456" + team_id = "team-dup-123" + + # The two different cache keys used for the same data + key_format_1 = f"{team_id}_{user_id}" # user_api_key_auth format + key_format_2 = f"team_membership:{user_id}:{team_id}" # auth_checks format + + _ = { + "user_id": user_id, + "team_id": team_id, + "spend": 3.0, + } + + # Document that these are different keys + assert ( + key_format_1 != key_format_2 + ), "Cache keys should be different (this is the bug)" + + # Document that these are different keys + assert ( + key_format_1 != key_format_2 + ), "Cache keys should be different (this is the bug)" + + +# ============================================================================ +# TEST: Full hot path network count summary +# ============================================================================ + + +@pytest.mark.asyncio +async def test_full_hot_path_network_count(): + """ + Summary test that counts all network operations when processing + a request with a key that has team_id and user_id attached. + + This test verifies the baseline number of cache operations expected + on a fully warm cache path. + """ + api_key = "sk-test-full-path" + team_id = "team-full-123" + user_id = "user-full-456" + hashed_token = hash_token(api_key) + + # Create all objects + valid_token = _create_valid_token( + api_key, team_id, user_id, has_team_member_spend=True + ) + team_obj = _create_team_object(team_id) + user_obj = _create_user_object(user_id) + membership_data = LiteLLM_TeamMembership( + user_id=user_id, + team_id=team_id, + spend=3.0, + budget_id=None, + litellm_budget_table=None, + ) + + # Pre-populate cache with all data + cache = DualCache(in_memory_cache=InMemoryCache()) + await cache.async_set_cache(key=hashed_token, value=valid_token) + await cache.async_set_cache(key=f"team_id:{team_id}", value=team_obj) + await cache.async_set_cache(key=user_id, value=user_obj) + await cache.async_set_cache( + key=f"team_membership:{user_id}:{team_id}", value=membership_data.model_dump() + ) + await cache.async_set_cache( + key=f"{team_id}_{user_id}", value=membership_data.model_dump() + ) + + # Create tracker AFTER populating cache + tracker = CacheCallTracker() + tracked_cache = _wrap_cache_with_tracker(cache, tracker) + + # Mock prisma (should not be called on warm cache) + mock_prisma = MagicMock() + + # Call each function to simulate the hot path + await get_key_object( + hashed_token=hashed_token, + prisma_client=mock_prisma, + user_api_key_cache=tracked_cache, + parent_otel_span=None, + proxy_logging_obj=None, + ) + + await get_team_object( + team_id=team_id, + prisma_client=mock_prisma, + user_api_key_cache=tracked_cache, + parent_otel_span=None, + proxy_logging_obj=None, + ) + + await get_user_object( + user_id=user_id, + prisma_client=mock_prisma, + user_api_key_cache=tracked_cache, + parent_otel_span=None, + proxy_logging_obj=None, + user_id_upsert=False, + ) + + await get_team_membership( + user_id=user_id, + team_id=team_id, + prisma_client=mock_prisma, + user_api_key_cache=tracked_cache, + parent_otel_span=None, + proxy_logging_obj=None, + ) + + summary = tracker.get_summary() + + # Assertions for expected baseline + # On warm cache: 4 reads (key, team, user, team_membership) + assert ( + summary["total_cache_reads"] == 4 + ), f"Expected 4 cache reads on warm path, got {summary['total_cache_reads']}" + + # No DB queries on warm cache + assert ( + summary["total_db_queries"] == 0 + ), f"Expected 0 DB queries on warm path, got {summary['total_db_queries']}" + + # Total network requests should be exactly 4 on warm cache + assert ( + summary["total_network_requests"] == 4 + ), f"Expected 4 total network requests on warm path, got {summary['total_network_requests']}" # ============================================================================ @@ -540,7 +540,7 @@ async def test_get_user_object_missing_user_negative_cache(): mock_prisma.db.litellm_usertable.find_unique = AsyncMock(return_value=None) for _ in range(3): - with pytest.raises(ValueError): + with pytest.raises(ValueError, match="User doesn't exist in db\\."): await get_user_object( user_id=user_id, prisma_client=mock_prisma, @@ -570,7 +570,7 @@ async def test_get_user_object_missing_user_rechecks_after_expiry(): mock_prisma.db.litellm_usertable = MagicMock() mock_prisma.db.litellm_usertable.find_unique = AsyncMock(return_value=None) - with pytest.raises(ValueError): + with pytest.raises(ValueError, match="User doesn't exist in db\\."): await get_user_object( user_id=user_id, prisma_client=mock_prisma, @@ -586,7 +586,7 @@ async def test_get_user_object_missing_user_rechecks_after_expiry(): time.time() - (db_cache_expiry + 1), ) - with pytest.raises(ValueError): + with pytest.raises(ValueError, match="User doesn't exist in db\\."): await get_user_object( user_id=user_id, prisma_client=mock_prisma, diff --git a/tests/test_litellm/proxy/auth/test_auth_utils.py b/tests/test_litellm/proxy/auth/test_auth_utils.py index ecf7f89d487..9301176f3ed 100644 --- a/tests/test_litellm/proxy/auth/test_auth_utils.py +++ b/tests/test_litellm/proxy/auth/test_auth_utils.py @@ -1588,7 +1588,7 @@ class TestCheckCompleteCredentialsBlocksSSRF: "litellm.proxy.auth.auth_utils.validate_url", side_effect=SSRFError(f"blocked: {blocked_url}"), ): - with pytest.raises(ValueError) as exc_info: + with pytest.raises(ValueError, match='is rejected by the SSRF guard') as exc_info: check_complete_credentials( { "model": "gpt-4", @@ -2144,7 +2144,7 @@ class TestIsRequestBodySafeBlocksEndpointTargetingFields: ], ) def test_endpoint_targeting_field_in_request_body_is_rejected(self, field): - with pytest.raises(ValueError) as exc: + with pytest.raises(ValueError, match='Rejected Request') as exc: is_request_body_safe( request_body={"model": "gpt-4", field: "https://attacker.example"}, general_settings={}, @@ -2165,7 +2165,7 @@ class TestIsRequestBodySafeBlocksEndpointTargetingFields: # on the blocklist into an SSRF / credential-exfil hole. Verify # that supplying an api_key (alongside the banned param) does NOT # bypass the gate — it can only be opened by an admin opt-in. - with pytest.raises(ValueError) as exc: + with pytest.raises(ValueError, match='Rejected Request') as exc: is_request_body_safe( request_body={ "model": "gpt-4", @@ -2722,7 +2722,7 @@ class TestObservabilityCallbackBans: ], ) def test_observability_field_in_request_body_root_is_rejected(self, field): - with pytest.raises(ValueError) as exc: + with pytest.raises(ValueError, match='Rejected Request') as exc: is_request_body_safe( request_body={"model": "gpt-4", field: "attacker-value"}, general_settings={}, @@ -2752,7 +2752,7 @@ class TestObservabilityCallbackBans: # Verifies the metadata walk: a value smuggled inside ``metadata`` # or ``litellm_metadata`` is just as dangerous as the same field # at the body root, and must hit the same gate. - with pytest.raises(ValueError) as exc: + with pytest.raises(ValueError, match='Rejected Request') as exc: is_request_body_safe( request_body={ "model": "gpt-4", @@ -2787,7 +2787,7 @@ class TestObservabilityCallbackBans: ) def test_observability_field_in_litellm_params_metadata_is_rejected(self): - with pytest.raises(ValueError) as exc: + with pytest.raises(ValueError, match='Rejected Request: turn_off_message_logging is not allowed') as exc: is_request_body_safe( request_body={ "model": "gpt-4", @@ -2814,7 +2814,7 @@ class TestObservabilityCallbackBans: # the ``isinstance(dict)`` guard. import json - with pytest.raises(ValueError) as exc: + with pytest.raises(ValueError, match='Rejected Request: langfuse_host is not allowed in request') as exc: is_request_body_safe( request_body={ "model": "gpt-4", @@ -2887,7 +2887,7 @@ def test_model_level_allow_does_not_skip_subsequent_banned_params(monkeypatch): lambda model, param, request_body_value, llm_router: param == "api_base", ) - with pytest.raises(ValueError) as exc: + with pytest.raises(ValueError, match='Rejected Request: langfuse_host is not allowed in request') as exc: is_request_body_safe( request_body={ "model": "gpt-4", @@ -2958,7 +2958,7 @@ class TestPricingInjectionBlocked: ], ) def test_pricing_field_rejected_by_default(self, field, value): - with pytest.raises(ValueError) as exc: + with pytest.raises(ValueError, match='Rejected Request') as exc: is_request_body_safe( request_body={"model": "gpt-4", field: value}, general_settings={}, diff --git a/tests/test_litellm/proxy/auth/test_handle_jwt.py b/tests/test_litellm/proxy/auth/test_handle_jwt.py index a9e12beb54b..99a0a4c0a8b 100644 --- a/tests/test_litellm/proxy/auth/test_handle_jwt.py +++ b/tests/test_litellm/proxy/auth/test_handle_jwt.py @@ -2589,7 +2589,7 @@ async def test_find_and_validate_raises_when_required_team_not_found(): # Token without team info jwt_token = {"sub": "user-1"} - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match="No team found in token\\. Checked team_id field 'None' and") as exc_info: await JWTAuthManager.find_and_validate_specific_team_id( jwt_handler=jwt_handler, jwt_valid_token=jwt_token, @@ -2916,7 +2916,7 @@ async def test_find_and_validate_specific_team_id_hints_bracket_notation(): # token has roles as a list — dot-notation won't find anything token = {"roles": ["team1"]} - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match="is not supported\\. Use 'roles' instead — LiteLLM") as exc_info: await JWTAuthManager.find_and_validate_specific_team_id( jwt_handler=handler, jwt_valid_token=token, @@ -2947,7 +2947,7 @@ async def test_find_and_validate_specific_team_id_hints_bracket_index_notation() handler = _make_jwt_handler("roles[0]") token = {"roles": ["team1"]} - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match="is not supported in team_id_jwt_field\\. Use 'roles' instead") as exc_info: await JWTAuthManager.find_and_validate_specific_team_id( jwt_handler=handler, jwt_valid_token=token, @@ -2977,7 +2977,7 @@ async def test_find_and_validate_specific_team_id_no_hint_for_valid_field(): handler = _make_jwt_handler("appid") token = {} # no appid — triggers the "no team found" path - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match="No team found in token\\. Checked team_id field 'appid' and") as exc_info: await JWTAuthManager.find_and_validate_specific_team_id( jwt_handler=handler, jwt_valid_token=token, @@ -4807,7 +4807,7 @@ async def test_multi_issuer_jwt_unknown_issuer_falls_back_to_global_jwks(monkeyp kid="issuer-key", ) - with pytest.raises(Exception) as exc: + with pytest.raises(Exception, match='Missing JWT Public Key URL from environment\\.') as exc: await jwt_handler.auth_jwt(token=token) assert "Missing JWT Public Key URL from environment." in str(exc.value) @@ -4838,7 +4838,7 @@ async def test_multi_issuer_jwt_rejects_wrong_audience(monkeypatch): kid="issuer-key", ) - with pytest.raises(Exception) as exc: + with pytest.raises(Exception, match="Validation fails: Audience doesn't match") as exc: await jwt_handler.auth_jwt(token=token) assert "Validation fails" in str(exc.value) @@ -4881,7 +4881,7 @@ async def test_multi_issuer_jwt_same_kid_does_not_cross_issuer_keys(monkeypatch) kid=shared_kid, ) - with pytest.raises(Exception) as exc: + with pytest.raises(Exception, match='Validation fails: Signature verification failed') as exc: await jwt_handler.auth_jwt(token=token) assert "Validation fails" in str(exc.value) @@ -4936,7 +4936,7 @@ def test_multi_issuer_jwt_requires_audience_unless_explicitly_disabled( issuer = "https://issuer.example.com" jwks_url = f"{issuer}/keys" - with pytest.raises(Exception) as exc: + with pytest.raises(Exception, match='must configure audience or set') as exc: LiteLLM_JWTAuth( issuers=[ { @@ -4953,7 +4953,7 @@ def test_multi_issuer_jwt_rejects_audience_with_disable_audience_validation(): issuer = "https://issuer.example.com" jwks_url = f"{issuer}/keys" - with pytest.raises(Exception) as exc: + with pytest.raises(Exception, match='cannot set audience and disable_audience_validation=True') as exc: LiteLLM_JWTAuth( issuers=[ { diff --git a/tests/test_litellm/proxy/auth/test_oauth2_proxy_hook.py b/tests/test_litellm/proxy/auth/test_oauth2_proxy_hook.py index dcbfd281e01..2d81d48de1e 100644 --- a/tests/test_litellm/proxy/auth/test_oauth2_proxy_hook.py +++ b/tests/test_litellm/proxy/auth/test_oauth2_proxy_hook.py @@ -141,7 +141,7 @@ async def test_refuses_to_map_non_identity_fields(configure_proxy, privileged_fi configure_proxy(mappings={privileged_field: f"x-{privileged_field}"}) request = _request_with_headers({f"x-{privileged_field}": "proxy_admin"}) - with pytest.raises(ValueError) as exc: + with pytest.raises(ValueError, match='proxy auth refuses to map non-identity UserAPIKeyAuth') as exc: await handle_oauth2_proxy_request(request) assert privileged_field in str(exc.value) diff --git a/tests/test_litellm/proxy/auth/test_route_checks.py b/tests/test_litellm/proxy/auth/test_route_checks.py index 2f4c2d3870f..636c5480d67 100644 --- a/tests/test_litellm/proxy/auth/test_route_checks.py +++ b/tests/test_litellm/proxy/auth/test_route_checks.py @@ -42,7 +42,7 @@ def test_non_admin_config_update_route_rejected(): request.query_params = {} # Test that calling /config/update route raises HTTPException with 403 status - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Only proxy admin can be used to generate, delete, update') as exc_info: RouteChecks.non_proxy_admin_allowed_routes_check( user_obj=user_obj, _user_role=LitellmUserRoles.INTERNAL_USER.value, @@ -134,7 +134,7 @@ def test_user_banner_update_rejected_for_non_admin(): request = MagicMock(spec=Request) request.query_params = {} - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Only proxy admin can be used to generate, delete, update') as exc_info: RouteChecks.non_proxy_admin_allowed_routes_check( user_obj=user_obj, _user_role=LitellmUserRoles.INTERNAL_USER.value, @@ -1814,7 +1814,7 @@ def test_internal_user_blocked_from_global_spend_routes(route): request = MagicMock(spec=Request) request.query_params = {} - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Only proxy admin can be used to generate, delete, update') as exc_info: RouteChecks.non_proxy_admin_allowed_routes_check( user_obj=user_obj, _user_role=LitellmUserRoles.INTERNAL_USER.value, @@ -1843,7 +1843,7 @@ def test_internal_user_view_only_blocked_from_global_spend_routes(route): request = MagicMock(spec=Request) request.query_params = {} - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Only proxy admin can be used to generate, delete, update') as exc_info: RouteChecks.non_proxy_admin_allowed_routes_check( user_obj=user_obj, _user_role=LitellmUserRoles.INTERNAL_USER_VIEW_ONLY.value, @@ -2046,7 +2046,7 @@ def test_internal_user_blocked_from_admin_viewer_logs_routes(route): if route not in INTERNAL_USER_BLOCKED_SUBSET: return - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Only proxy admin can be used to generate, delete, update') as exc_info: RouteChecks.non_proxy_admin_allowed_routes_check( user_obj=user_obj, _user_role=LitellmUserRoles.INTERNAL_USER.value, @@ -2530,7 +2530,7 @@ def test_non_admin_non_team_admin_cannot_access_config_update_but_can_attempt_re ) # /config/update is still blocked - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Only proxy admin can be used to generate, delete, update') as exc_info: RouteChecks.non_proxy_admin_allowed_routes_check( user_obj=user_obj, _user_role=LitellmUserRoles.INTERNAL_USER.value, @@ -2809,7 +2809,7 @@ def test_team_update_gate_rejects_without_org_context(): request.method = "POST" request.query_params = {} - with pytest.raises(Exception): + with pytest.raises(Exception, match="Only proxy admin can be used to generate"): RouteChecks.non_proxy_admin_allowed_routes_check( user_obj=user_obj, _user_role=LitellmUserRoles.INTERNAL_USER.value, @@ -2829,7 +2829,7 @@ def test_team_update_gate_rejects_cross_org_admin_with_resolved_org(): request.method = "POST" request.query_params = {} - with pytest.raises(Exception): + with pytest.raises(Exception, match="Only proxy admin can be used to generate"): RouteChecks.non_proxy_admin_allowed_routes_check( user_obj=user_obj, _user_role=LitellmUserRoles.INTERNAL_USER.value, @@ -2951,7 +2951,7 @@ def test_patch_team_gate_rejects_regular_internal_user(): ) valid_token = UserAPIKeyAuth(user_id="regular-user", user_role=LitellmUserRoles.INTERNAL_USER.value) - with pytest.raises(Exception): + with pytest.raises(Exception, match="Only proxy admin can be used to generate"): RouteChecks.non_proxy_admin_allowed_routes_check( user_obj=user_obj, _user_role=LitellmUserRoles.INTERNAL_USER.value, @@ -2967,7 +2967,7 @@ def test_patch_team_gate_rejects_cross_org_admin(): user_obj = _make_org_admin_user("org-1") valid_token = UserAPIKeyAuth(user_id="org-admin-user", user_role=LitellmUserRoles.INTERNAL_USER.value) - with pytest.raises(Exception): + with pytest.raises(Exception, match="Only proxy admin can be used to generate"): RouteChecks.non_proxy_admin_allowed_routes_check( user_obj=user_obj, _user_role=LitellmUserRoles.INTERNAL_USER.value, @@ -2987,7 +2987,7 @@ def test_patch_team_gate_rejects_view_only_admin(): ) valid_token = UserAPIKeyAuth(user_id="viewer", user_role=LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY.value) - with pytest.raises(Exception): + with pytest.raises(HTTPException): RouteChecks.non_proxy_admin_allowed_routes_check( user_obj=user_obj, _user_role=LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY.value, @@ -3188,7 +3188,7 @@ def test_internal_user_blocked_from_search_tool_writes(route): request = MagicMock(spec=Request) request.query_params = {} - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Only proxy admin can be used to generate, delete, update') as exc_info: RouteChecks.non_proxy_admin_allowed_routes_check( user_obj=user_obj, _user_role=LitellmUserRoles.INTERNAL_USER.value, diff --git a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py index ab7e3d9701c..043bbb5b76a 100644 --- a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py +++ b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py @@ -5337,7 +5337,7 @@ async def test_random_non_sk_token_is_rejected(monkeypatch): patch("litellm.proxy.proxy_server.master_key", "sk-master"), patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), ): - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='LiteLLM Virtual Key expected\\.') as exc_info: await user_api_key_auth( request=mock_request, api_key="Bearer not-a-real-token", @@ -5539,7 +5539,7 @@ async def test_real_jwt_still_requires_license_when_jwt_auth_enabled(monkeypatch patch("litellm.proxy.proxy_server.master_key", "sk-master"), patch("litellm.proxy.proxy_server.prisma_client", None), ): - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='JWT Auth is an enterprise only feature\\. You must be a') as exc_info: await user_api_key_auth( request=mock_request, api_key=f"Bearer {jwt_token}", diff --git a/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py b/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py index b216071828e..863d9204cff 100644 --- a/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py @@ -979,7 +979,7 @@ async def test_create__exception_calls_failure_hook(harness, openai_env_creds): ) harness.litellm_acreate.side_effect = ValueError("provider boom") - with pytest.raises(Exception): + with pytest.raises(ProxyException): await call_create(harness) harness.logging.post_call_failure_hook.assert_called_once() @@ -1437,7 +1437,7 @@ async def test_retrieve__uses_aretrieve_batch_route_type(retrieve_harness, opena async def test_retrieve__exception_calls_failure_hook(retrieve_harness, openai_env_creds): retrieve_harness.litellm_aretrieve.side_effect = ValueError("provider boom") - with pytest.raises(Exception): + with pytest.raises(ProxyException): await call_retrieve(retrieve_harness, "batch-raw-xyz") retrieve_harness.logging.post_call_failure_hook.assert_called_once() @@ -1844,7 +1844,7 @@ async def test_list__uses_alist_batches_route_type(list_harness): async def test_list__exception_calls_failure_hook(list_harness): list_harness.litellm_alist.side_effect = ValueError("provider boom") - with pytest.raises(Exception): + with pytest.raises(ProxyException): await call_list(list_harness) list_harness.logging.post_call_failure_hook.assert_called_once() @@ -2233,7 +2233,7 @@ async def test_cancel__uses_acancel_batch_route_type(cancel_harness, openai_env_ async def test_cancel__exception_calls_failure_hook(cancel_harness, openai_env_creds): cancel_harness.litellm_acancel.side_effect = ValueError("provider boom") - with pytest.raises(Exception): + with pytest.raises(ProxyException): await call_cancel(cancel_harness, "batch-raw-xyz") cancel_harness.logging.post_call_failure_hook.assert_called_once() diff --git a/tests/test_litellm/proxy/client/cli/test_auth_commands.py b/tests/test_litellm/proxy/client/cli/test_auth_commands.py index eb40f54a1f3..be29269fe25 100644 --- a/tests/test_litellm/proxy/client/cli/test_auth_commands.py +++ b/tests/test_litellm/proxy/client/cli/test_auth_commands.py @@ -84,7 +84,7 @@ class TestPollingErrorSurfacing: } with patch("requests.get", return_value=mock_response) as mock_get, patch("time.sleep"): - with pytest.raises(ValueError) as exc_info: + with pytest.raises(ValueError, match='Your litellm CLI is out of date and uses a login flow') as exc_info: _poll_for_ready_data("http://test/sso/cli/poll/sk-legacy") assert mock_get.call_count == 1 @@ -151,7 +151,7 @@ class TestStartCliSsoFlowErrors: mock_response.status_code = 404 with patch("requests.post", return_value=mock_response): - with pytest.raises(ValueError) as exc_info: + with pytest.raises(ValueError, match='Either --base-url is wrong, or the proxy is older than') as exc_info: _start_cli_sso_flow("https://old-proxy.example.com") message = str(exc_info.value) @@ -167,7 +167,7 @@ class TestStartCliSsoFlowErrors: mock_response.json.return_value = {"detail": "Too many CLI login attempts. Try again later."} with patch("requests.post", return_value=mock_response): - with pytest.raises(ValueError) as exc_info: + with pytest.raises(ValueError, match='Too many CLI login attempts\\. Try again later\\.') as exc_info: _start_cli_sso_flow("https://test.example.com") assert "HTTP 429" in str(exc_info.value) @@ -183,7 +183,7 @@ class TestStartCliSsoFlowErrors: mock_response.text = "Sign in to corporate VPN" with patch("requests.post", return_value=mock_response): - with pytest.raises(ValueError) as exc_info: + with pytest.raises(ValueError, match='A proxy, load balancer, or auth gateway in front of') as exc_info: _start_cli_sso_flow("https://test.example.com") message = str(exc_info.value) @@ -197,7 +197,7 @@ class TestStartCliSsoFlowErrors: from litellm.proxy.client.cli.commands.auth import _start_cli_sso_flow with patch("requests.post", side_effect=requests.ConnectionError("Connection refused")): - with pytest.raises(ValueError) as exc_info: + with pytest.raises(ValueError, match='Connection refused\\. Check that the proxy is running') as exc_info: _start_cli_sso_flow("https://unreachable.example.com") message = str(exc_info.value) diff --git a/tests/test_litellm/proxy/client/cli/test_pkce_login.py b/tests/test_litellm/proxy/client/cli/test_pkce_login.py index 70f481d5cfa..f58bd0ff412 100644 --- a/tests/test_litellm/proxy/client/cli/test_pkce_login.py +++ b/tests/test_litellm/proxy/client/cli/test_pkce_login.py @@ -622,7 +622,7 @@ def test_fresh_api_key_never_hands_out_a_rotated_key_it_could_not_save(): def save(_record): raise OSError("disk full") - with pytest.raises(OSError): + with pytest.raises(OSError, match="disk full"): _fresh(STORED, save, http, now=lambda: 999_950.0) diff --git a/tests/test_litellm/proxy/client/test_models.py b/tests/test_litellm/proxy/client/test_models.py index b2485032a37..33f963b74af 100644 --- a/tests/test_litellm/proxy/client/test_models.py +++ b/tests/test_litellm/proxy/client/test_models.py @@ -472,14 +472,14 @@ def test_get_invalid_params(): client = ModelsManagementClient(base_url="http://localhost:8000") # Test with no parameters - with pytest.raises(ValueError) as exc_info: + with pytest.raises(ValueError, match='Exactly one of model_id or model_name must be provided') as exc_info: client.get() assert "Exactly one of model_id or model_name must be provided" in str( exc_info.value ) # Test with both parameters - with pytest.raises(ValueError) as exc_info: + with pytest.raises(ValueError, match='Exactly one of model_id or model_name must be provided') as exc_info: client.get(model_id="123", model_name="gpt-4") assert "Exactly one of model_id or model_name must be provided" in str( exc_info.value diff --git a/tests/test_litellm/proxy/common_utils/test_callback_utils.py b/tests/test_litellm/proxy/common_utils/test_callback_utils.py index 515a7b27c7b..77ada4c11a9 100644 --- a/tests/test_litellm/proxy/common_utils/test_callback_utils.py +++ b/tests/test_litellm/proxy/common_utils/test_callback_utils.py @@ -586,7 +586,7 @@ def test_initialize_callbacks_on_proxy_rejects_class_valued_entry(probe_config_p silently never run the hook. Config load must fail instead.""" entry = f"{_PROBE_MODULE_NAME}.FloorMaxTokens" - with pytest.raises(ValueError) as exc_info: + with pytest.raises(ValueError, match='litellm_settings\\.callbacks entry') as exc_info: _load_callbacks([entry], probe_config_path) message = str(exc_info.value) @@ -609,7 +609,7 @@ def test_initialize_callbacks_on_proxy_rejects_non_dispatchable_values( ): entry = f"{_PROBE_MODULE_NAME}.{attribute}" - with pytest.raises(ValueError) as exc_info: + with pytest.raises(ValueError, match='litellm_settings\\.callbacks entry') as exc_info: _load_callbacks([entry], probe_config_path) message = str(exc_info.value) @@ -621,7 +621,7 @@ def test_initialize_callbacks_on_proxy_rejects_non_dispatchable_values( def test_initialize_callbacks_on_proxy_rejects_class_valued_non_list_value(probe_config_path): entry = f"{_PROBE_MODULE_NAME}.FloorMaxTokens" - with pytest.raises(ValueError) as exc_info: + with pytest.raises(ValueError, match='litellm_settings\\.callbacks entry') as exc_info: _load_callbacks(entry, probe_config_path) assert entry in str(exc_info.value) diff --git a/tests/test_litellm/proxy/common_utils/test_path_utils.py b/tests/test_litellm/proxy/common_utils/test_path_utils.py index c8d58fa8259..8936d910777 100644 --- a/tests/test_litellm/proxy/common_utils/test_path_utils.py +++ b/tests/test_litellm/proxy/common_utils/test_path_utils.py @@ -42,5 +42,5 @@ class TestSafeFilename: safe_filename("..") def test_empty_rejected(self): - with pytest.raises(ValueError): + with pytest.raises(ValueError, match='Empty or unsafe filename'): safe_filename("") diff --git a/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py b/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py index c5bc4e29f81..8233b0d3864 100644 --- a/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py +++ b/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py @@ -16,6 +16,11 @@ sys.path.insert(0, os.path.abspath("../../..")) # Adds the parent directory to from litellm.proxy._types import LiteLLM_VerificationToken from litellm.proxy.common_utils import reset_budget_job as reset_budget_job_module +from litellm.constants import ( + PROXY_BUDGET_RESCHEDULER_MIN_TIME, + RESET_BUDGET_JOB_LOCK_TTL_SECONDS, + RESET_BUDGET_JOB_NAME, +) from litellm.proxy.common_utils.reset_budget_job import ResetBudgetJob from litellm.proxy.common_utils.timezone_utils import BudgetResetSettings @@ -1906,10 +1911,7 @@ def test_key_reset_keeps_paging_when_some_rows_in_a_chunk_fail(monkeypatch): assert client.fetches_by_table["key"] == 2 assert [w["where"]["token"] for w in _batch_writes(client, "key", op="update")] == ["k1", "k2"] assert [call["call_type"] for call in logging_obj.service_logging_obj.failure_calls] == ["reset_budget_keys"] - assert set(logging_obj.service_logging_obj.failure_calls[0]["event_metadata"]) == { - "num_keys_found", - "keys_found", - } + assert set(logging_obj.service_logging_obj.failure_calls[0]["event_metadata"]) == {"num_keys_found"} assert [call["call_type"] for call in logging_obj.service_logging_obj.success_calls] == ["reset_budget_keys"] @@ -1936,6 +1938,352 @@ def test_user_and_team_chunks_report_progress_despite_a_failed_row( assert [call["call_type"] for call in logging_obj.service_logging_obj.failure_calls] == [call_type] +class FakePodLockManager: + """Stands in for the redis-backed PodLockManager. + + Lets a test pick which of the three states a pod lands in: it wins the + lease, another pod already holds it, or redis cannot answer at all. + """ + + def __init__(self, *, acquired: bool, held_by_other: bool = False, has_redis: bool = True): + self.redis_cache = MagicMock() if has_redis else None + if self.redis_cache is not None: + self.redis_cache.async_get_cache = AsyncMock(return_value="another-pod" if held_by_other else None) + self._acquired = acquired + self.acquire_calls: List[Dict[str, Any]] = [] + self.release_calls: List[str] = [] + + @staticmethod + def get_redis_lock_key(cronjob_id: str) -> str: + return f"cronjob_lock:{cronjob_id}" + + async def acquire_lock(self, cronjob_id: str, ttl: Any = None) -> bool: + self.acquire_calls.append({"cronjob_id": cronjob_id, "ttl": ttl}) + return self._acquired + + async def release_lock(self, cronjob_id: str) -> None: + self.release_calls.append(cronjob_id) + + +def _make_leader_election_job(monkeypatch, pod_lock_manager): + """A ResetBudgetJob wired to one lock manager, with every read observable. + + `prisma_client.get_data_calls` plus `prisma_client.db.query_raw` together + cover every read the sweep makes, so a pod that skipped the tick leaves + both untouched. + """ + prisma_client = MockPrismaClient() + prisma_client.db.query_raw = AsyncMock(return_value=[]) + + spend_counter_cache = MagicMock() + spend_counter_cache.redis_cache = None + fake_module = types.ModuleType("litellm.proxy.proxy_server") + fake_module.spend_counter_cache = spend_counter_cache + monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", fake_module) + + job = ResetBudgetJob( + proxy_logging_obj=MockProxyLogging(), + prisma_client=prisma_client, + pod_lock_manager=pod_lock_manager, + ) + return job, prisma_client + + +def _swept(prisma_client) -> bool: + return bool(prisma_client.get_data_calls) or prisma_client.db.query_raw.await_count > 0 + + +def test_reset_budget_sweeps_and_releases_when_it_wins_the_lease(monkeypatch): + """The elected pod does the work and hands the lease back, so the next tick + can elect any pod rather than waiting out the TTL.""" + lock = FakePodLockManager(acquired=True) + job, prisma_client = _make_leader_election_job(monkeypatch, lock) + + asyncio.run(job.reset_budget()) + + assert _swept(prisma_client) + assert [call["cronjob_id"] for call in lock.acquire_calls] == [RESET_BUDGET_JOB_NAME] + assert lock.release_calls == [RESET_BUDGET_JOB_NAME] + + +def test_reset_budget_does_nothing_when_another_pod_holds_the_lease(monkeypatch): + """The whole point of the lease: a fleet must not multiply one sweep by its + replica count. A pod that loses the election issues no query at all, and + must not release a lease it never took.""" + lock = FakePodLockManager(acquired=False, held_by_other=True) + job, prisma_client = _make_leader_election_job(monkeypatch, lock) + + asyncio.run(job.reset_budget()) + + assert not _swept(prisma_client) + assert lock.release_calls == [] + + +def test_reset_budget_sweeps_unguarded_when_redis_cannot_answer(monkeypatch): + """acquire_lock reports contention and an unreachable redis identically, so + reading a failed acquire as contention would strand every expired budget at + its cap on every pod for as long as redis is down. No holder means sweep.""" + lock = FakePodLockManager(acquired=False, held_by_other=False) + job, prisma_client = _make_leader_election_job(monkeypatch, lock) + + asyncio.run(job.reset_budget()) + + assert _swept(prisma_client) + assert lock.release_calls == [] + + +def test_reset_budget_sweeps_when_the_deployment_has_no_redis(monkeypatch): + """A single-pod or redis-less deployment keeps its pre-election behavior.""" + lock = FakePodLockManager(acquired=False, has_redis=False) + job, prisma_client = _make_leader_election_job(monkeypatch, lock) + + asyncio.run(job.reset_budget()) + + assert _swept(prisma_client) + assert lock.acquire_calls == [] + assert lock.release_calls == [] + + +def test_reset_budget_sweeps_when_no_lock_manager_is_injected(monkeypatch): + """Callers that construct the job without a lock manager still sweep.""" + job, prisma_client = _make_leader_election_job(monkeypatch, None) + + asyncio.run(job.reset_budget()) + + assert _swept(prisma_client) + + +def test_reset_budget_releases_the_lease_when_a_phase_raises(monkeypatch): + """A crash mid-sweep must not hold the lease for its whole TTL, which would + stop every pod resetting budgets until it expired.""" + lock = FakePodLockManager(acquired=True) + job, _ = _make_leader_election_job(monkeypatch, lock) + + async def boom() -> None: + raise RuntimeError("phase exploded") + + monkeypatch.setattr(job, "reset_budget_for_litellm_keys", boom) + + with pytest.raises(RuntimeError): + asyncio.run(job.reset_budget()) + + assert lock.release_calls == [RESET_BUDGET_JOB_NAME] + + +def test_reset_budget_lease_outlives_one_scheduler_tick(monkeypatch): + """A lease shorter than the gap between ticks expires mid-sweep and lets a + second pod start sweeping, which is the amplification the lease removes.""" + lock = FakePodLockManager(acquired=True) + job, _ = _make_leader_election_job(monkeypatch, lock) + + asyncio.run(job.reset_budget()) + + assert lock.acquire_calls[0]["ttl"] == RESET_BUDGET_JOB_LOCK_TTL_SECONDS + assert RESET_BUDGET_JOB_LOCK_TTL_SECONDS > PROXY_BUDGET_RESCHEDULER_MIN_TIME + + +def _window_row(source_id_column: str, row_id: str, reset_at: datetime) -> Dict[str, Any]: + return { + source_id_column: row_id, + "budget_limits": [{"budget_duration": "1h", "reset_at": reset_at.isoformat(), "max_budget": 10}], + } + + +def _paginating_window_job(monkeypatch, pages_by_table: Dict[str, List[List[Dict[str, Any]]]]): + """Serve each table a canned sequence of pages and record every query. + + Returns (job, calls) where calls is a list of (sql, cursor, limit). + """ + prisma_client = MagicMock() + remaining = {table: list(pages) for table, pages in pages_by_table.items()} + calls: List[Dict[str, Any]] = [] + + async def fake_query_raw(query: str, *args, **kwargs): + table = "key" if '"LiteLLM_VerificationToken"' in query else "team" + calls.append({"table": table, "sql": query, "cursor": args[0], "limit": args[1]}) + pages = remaining[table] + return pages.pop(0) if pages else [] + + prisma_client.db.query_raw = AsyncMock(side_effect=fake_query_raw) + prisma_client.db.litellm_verificationtoken.update = AsyncMock(return_value=None) + prisma_client.db.litellm_teamtable.update = AsyncMock(return_value=None) + + spend_counter_cache = MagicMock() + spend_counter_cache.redis_cache = None + fake_module = types.ModuleType("litellm.proxy.proxy_server") + fake_module.spend_counter_cache = spend_counter_cache + monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", fake_module) + + job = ResetBudgetJob(proxy_logging_obj=MagicMock(), prisma_client=prisma_client) + return job, calls + + +def test_reset_budget_windows_pages_by_cursor_instead_of_reading_the_table(monkeypatch): + """The window scan used to read every row carrying budget_limits in one + statement, so its memory and its statement cost grew with the deployment's + key count. It now walks pages, and each page resumes past the last row. + """ + monkeypatch.setattr(reset_budget_job_module, "RESET_BUDGET_JOB_BATCH_SIZE", 2) + past = datetime.utcnow() - timedelta(hours=2) + job, calls = _paginating_window_job( + monkeypatch, + { + "key": [ + [_window_row("token", "k1", past), _window_row("token", "k2", past)], + [_window_row("token", "k3", past)], + ], + "team": [[]], + }, + ) + + asyncio.run(job.reset_budget_windows()) + + key_calls = [call for call in calls if call["table"] == "key"] + assert [call["cursor"] for call in key_calls] == ["", "k2"], "second page must resume past the last row read" + assert {call["limit"] for call in key_calls} == {2} + assert all("LIMIT $2" in call["sql"] for call in key_calls) + # the short second page ends the scan; a third query would re-read forever + assert len(key_calls) == 2 + + +def test_reset_budget_windows_pages_to_the_end_of_a_large_table(monkeypatch): + """The scan must reach the last row within one tick. + + Capping the pages per run would need a resume position, and that position + cannot live in the process: the lease is released after every sweep, so a + later tick can elect a pod whose position is unset, restart at the first + row, and leave the tail pinned at its cap forever. Paging alone bounds the + memory, so the walk runs to completion instead. + """ + monkeypatch.setattr(reset_budget_job_module, "RESET_BUDGET_JOB_BATCH_SIZE", 1) + # the table needs far more pages than any per-run cap would allow, so a + # capped walk stops short and only an uncapped one reaches the last row + monkeypatch.setattr(reset_budget_job_module, "RESET_BUDGET_JOB_MAX_CHUNKS_PER_RUN", 3) + past = datetime.utcnow() - timedelta(hours=2) + rows = [_window_row("token", f"k{i:03d}", past) for i in range(1, 26)] + job, visited = _cursor_paginating_window_job(monkeypatch, rows) + + asyncio.run(job.reset_budget_windows()) + + assert visited == [f"k{i:03d}" for i in range(1, 26)], visited + + +def test_reset_budget_windows_survives_one_table_failing(monkeypatch): + """A broken key scan must not cost the team scan its sweep.""" + prisma_client = MagicMock() + + async def fake_query_raw(query: str, *args, **kwargs): + if '"LiteLLM_VerificationToken"' in query: + raise RuntimeError("key scan exploded") + return [] + + prisma_client.db.query_raw = AsyncMock(side_effect=fake_query_raw) + spend_counter_cache = MagicMock() + spend_counter_cache.redis_cache = None + fake_module = types.ModuleType("litellm.proxy.proxy_server") + fake_module.spend_counter_cache = spend_counter_cache + monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", fake_module) + job = ResetBudgetJob(proxy_logging_obj=MagicMock(), prisma_client=prisma_client) + + asyncio.run(job.reset_budget_windows()) + + queried = [call.args[0] for call in prisma_client.db.query_raw.await_args_list] + assert any('"LiteLLM_TeamTable"' in sql for sql in queried) + + +def test_row_payloads_stay_out_of_reset_job_event_metadata(monkeypatch): + """Every found and updated row used to be JSON-serialized into the service + hook's metadata on every chunk, on the event loop, whether or not any + consumer read it. Only the counts are reported now.""" + client = ChunkedPrismaClient({"key": [[_key_row("k1"), _key_row("k2")]]}) + logging_obj = RecordingProxyLogging() + job = ResetBudgetJob(proxy_logging_obj=logging_obj, prisma_client=client) + + _run_and_drain_hooks(job.reset_budget_for_litellm_keys) + + metadata = logging_obj.service_logging_obj.success_calls[0]["event_metadata"] + assert metadata["num_keys_found"] == 2 + assert metadata["num_keys_updated"] == 2 + assert {"keys_found", "keys_updated", "keys_failed"}.isdisjoint(metadata) + assert all(isinstance(value, int) for value in metadata.values()), metadata + + +def test_debug_row_dump_is_deferred_until_a_record_is_emitted(): + """`logger.debug("%s", json.dumps(rows))` serializes before the logger drops + the record, so the sweep paid for a full dump of every chunk at any log + level. The wrapper defers the work to the formatter.""" + serialized = [] + + class Tracked: + def __repr__(self) -> str: + serialized.append("serialized") + return "tracked" + + lazy = reset_budget_job_module._LazyJson([Tracked()]) + assert serialized == [], "constructing the wrapper must not serialize" + + assert "tracked" in str(lazy) + assert serialized == ["serialized"] + + +def _cursor_paginating_window_job(monkeypatch, key_rows: List[Dict[str, Any]]): + """Serve real keyset pages out of one ordered table, honouring the cursor. + + Unlike the canned-page helper above, this models the database: a page is + whatever rows sort after the cursor, so a scan that forgets its cursor + genuinely re-reads the same prefix. + """ + prisma_client = MagicMock() + ordered = sorted(key_rows, key=lambda r: r["token"]) + visited: List[str] = [] + + async def fake_query_raw(query: str, *args, **kwargs): + if '"LiteLLM_TeamTable"' in query: + return [] + cursor, limit = args[0], args[1] + page = [row for row in ordered if row["token"] > cursor][:limit] + visited.extend(row["token"] for row in page) + return page + + prisma_client.db.query_raw = AsyncMock(side_effect=fake_query_raw) + prisma_client.db.litellm_verificationtoken.update = AsyncMock(return_value=None) + prisma_client.db.litellm_teamtable.update = AsyncMock(return_value=None) + + spend_counter_cache = MagicMock() + spend_counter_cache.redis_cache = None + fake_module = types.ModuleType("litellm.proxy.proxy_server") + fake_module.spend_counter_cache = spend_counter_cache + monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", fake_module) + + job = ResetBudgetJob(proxy_logging_obj=MagicMock(), prisma_client=prisma_client) + return job, visited + + +def test_every_tick_sweeps_the_whole_window_table_whichever_pod_won(monkeypatch): + """Coverage must not depend on which pod was elected. + + The lease is released after each sweep, so consecutive ticks routinely run + on different pods. A scan carrying a resume position in process memory would + have a fresh pod start over at the first row, so rows past one run's reach + would never be swept by anyone. Two independent job instances, standing in + for two pods, must each cover the table end to end. + """ + monkeypatch.setattr(reset_budget_job_module, "RESET_BUDGET_JOB_BATCH_SIZE", 2) + monkeypatch.setattr(reset_budget_job_module, "RESET_BUDGET_JOB_MAX_CHUNKS_PER_RUN", 2) + past = datetime.utcnow() - timedelta(hours=2) + rows = [_window_row("token", f"k{i:03d}", past) for i in range(1, 12)] + expected = [f"k{i:03d}" for i in range(1, 12)] + + pod_a, visited_a = _cursor_paginating_window_job(monkeypatch, rows) + pod_b, visited_b = _cursor_paginating_window_job(monkeypatch, rows) + + asyncio.run(pod_a.reset_budget_windows()) + asyncio.run(pod_b.reset_budget_windows()) + + assert visited_a == expected, visited_a + assert visited_b == expected, visited_b + + class FlakyPrismaClient(MockPrismaClient): """A client whose first N reads (or first N batch commits) fail with a transport error, and which records every reconnect attempt. diff --git a/tests/test_litellm/proxy/common_utils/test_timezone_utils.py b/tests/test_litellm/proxy/common_utils/test_timezone_utils.py index 7f686c53c95..dc3917cb48e 100644 --- a/tests/test_litellm/proxy/common_utils/test_timezone_utils.py +++ b/tests/test_litellm/proxy/common_utils/test_timezone_utils.py @@ -130,16 +130,16 @@ def test_parse_budget_reset_time_unset_defaults_to_midnight(): def test_parse_budget_reset_time_invalid_string_raises(): - with pytest.raises(ValueError): + with pytest.raises(ValueError, match="hour 'HH:MM' or 'HH:MM:SS' string, e\\.g\\."): parse_budget_reset_time("25:00") - with pytest.raises(ValueError): + with pytest.raises(ValueError, match="Invalid budget_reset_time 'noon'; expected a"): parse_budget_reset_time("noon") def test_parse_budget_reset_time_non_string_raises(): # Unquoted "12:00" in YAML parses to the int 720; it must fail loudly, # not silently fall back to midnight. - with pytest.raises(ValueError): + with pytest.raises(ValueError, match="hour 'HH:MM' string, e\\.g\\."): parse_budget_reset_time(720) diff --git a/tests/test_litellm/proxy/conftest.py b/tests/test_litellm/proxy/conftest.py index 61752997f0f..65e12b7d777 100644 --- a/tests/test_litellm/proxy/conftest.py +++ b/tests/test_litellm/proxy/conftest.py @@ -18,6 +18,7 @@ from prisma.errors import ClientNotConnectedError _PROXY_MODULE_GLOBALS_TO_ISOLATE = ( "master_key", "prisma_client", + "llm_router", ) @@ -56,7 +57,10 @@ def pytest_runtest_setup(item): Without this, a leaked value (e.g. master_key set by a sibling test) flips the auth short-circuit in user_api_key_auth and causes unrelated - tests in the same xdist worker to return 401 instead of 200. + tests in the same xdist worker to return 401 instead of 200. A leaked + llm_router does the same to anything that reads the running router out + of sys.modules, such as the PTU rollup's deployment scan, which then + counts a sibling test's deployments as if the proxy owned them. This must be a hook pair, not an autouse fixture: an autouse fixture in the root conftest requests monkeypatch, so monkeypatch's undo stack diff --git a/tests/test_litellm/proxy/db/db_transaction_queue/test_spend_logs_partition_manager.py b/tests/test_litellm/proxy/db/db_transaction_queue/test_spend_logs_partition_manager.py index e949afce57b..4ea655b8871 100644 --- a/tests/test_litellm/proxy/db/db_transaction_queue/test_spend_logs_partition_manager.py +++ b/tests/test_litellm/proxy/db/db_transaction_queue/test_spend_logs_partition_manager.py @@ -308,9 +308,9 @@ async def test_partition_maintenance_issues_nothing_when_the_budget_is_already_s def test_unsupported_interval_raises(): - with pytest.raises(ValueError): + with pytest.raises(ValueError, match='Unsupported partition interval: year'): period_start(date(2026, 6, 1), "year") - with pytest.raises(ValueError): + with pytest.raises(ValueError, match='Unsupported partition interval: year'): next_period_start(date(2026, 6, 1), "year") diff --git a/tests/test_litellm/proxy/db/test_prisma_planned_engine_restart.py b/tests/test_litellm/proxy/db/test_prisma_planned_engine_restart.py index c8e0338eeaa..95e794012ec 100644 --- a/tests/test_litellm/proxy/db/test_prisma_planned_engine_restart.py +++ b/tests/test_litellm/proxy/db/test_prisma_planned_engine_restart.py @@ -913,7 +913,7 @@ async def test_health_check_alerts_for_non_connection_errors_during_a_replacemen await _yield_to_loop() assert wrapper._reconnection_lock.locked() is True - with pytest.raises(ValueError): + with pytest.raises(ValueError, match='malformed SELECT'): await client.health_check() gate.set() diff --git a/tests/test_litellm/proxy/db/test_prisma_self_heal.py b/tests/test_litellm/proxy/db/test_prisma_self_heal.py index e8e7568d6d8..cc47cf4a7e4 100644 --- a/tests/test_litellm/proxy/db/test_prisma_self_heal.py +++ b/tests/test_litellm/proxy/db/test_prisma_self_heal.py @@ -507,7 +507,7 @@ async def test_engine_confirmed_dead_persists_across_failed_heavy_reconnect( client._reap_all_zombies = MagicMock() with patch.dict(os.environ, {"DATABASE_URL": "postgresql://test"}): - with pytest.raises(Exception): + with pytest.raises(RuntimeError): await client._run_reconnect_cycle(timeout_seconds=5.0) # The flag must STILL be True so the next attempt re-enters the heavy diff --git a/tests/test_litellm/proxy/db/test_spend_log_tool_index.py b/tests/test_litellm/proxy/db/test_spend_log_tool_index.py index 71073fd216e..9c4fbbf41aa 100644 --- a/tests/test_litellm/proxy/db/test_spend_log_tool_index.py +++ b/tests/test_litellm/proxy/db/test_spend_log_tool_index.py @@ -327,7 +327,7 @@ class TestFlushToolUsageTransactions: async def test_non_connection_errors_do_not_retry(self): prisma = MagicMock() prisma.db.batch_ = MagicMock(side_effect=ValueError("bad data")) - with pytest.raises(ValueError): + with pytest.raises(ValueError, match="bad data"): await flush_tool_usage_transactions(prisma_client=prisma, transactions=[_transaction("r1")]) prisma.db.batch_.assert_called_once() diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_moderations.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_moderations.py index 9002d1f81a3..729dcb54309 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_moderations.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_moderations.py @@ -487,17 +487,18 @@ async def test_openai_moderation_guardrail_streaming_harmful_content(): # Should raise HTTPException when processing streaming harmful content from fastapi import HTTPException - with pytest.raises(HTTPException) as exc_info: + async def _drain(): result_chunks = [] - async for ( - chunk - ) in unified_guardrail.async_post_call_streaming_iterator_hook( + async for chunk in unified_guardrail.async_post_call_streaming_iterator_hook( user_api_key_dict=user_api_key_dict, response=mock_stream(), request_data=request_data, ): result_chunks.append(chunk) + with pytest.raises(HTTPException) as exc_info: + await _drain() + assert exc_info.value.status_code == 400 assert "Violated OpenAI moderation policy" in str(exc_info.value.detail) diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py index e3516b6eda7..5d971bf1212 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py @@ -2113,7 +2113,7 @@ async def test_make_bedrock_api_request_forwards_guardrail_action(): ): mock_post.return_value = mock_bedrock_response - with pytest.raises(Exception): + with pytest.raises(Exception, match="blocked"): await guardrail.make_bedrock_api_request( source="INPUT", messages=request_data["messages"], diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_invoke_guardrail_checks.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_invoke_guardrail_checks.py index 3adf8b8407d..ceb59571389 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_invoke_guardrail_checks.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_invoke_guardrail_checks.py @@ -56,7 +56,7 @@ def _patched(guardrail: BedrockGuardrail, http_response): def test_init_rejects_both_identifier_and_checks(): - with pytest.raises(ValueError): + with pytest.raises(ValueError, match='Bedrock guardrail accepts either'): BedrockGuardrail(guardrailIdentifier="gid", checks=CONTENT_FILTER_CHECKS) @@ -304,7 +304,7 @@ async def test_truncated_pii_ignored_when_pii_check_not_configured(): @pytest.mark.asyncio async def test_checks_with_guardrail_version_rejected(): - with pytest.raises(ValueError): + with pytest.raises(ValueError, match='Bedrock guardrail accepts either'): BedrockGuardrail(checks=CONTENT_FILTER_CHECKS, guardrailVersion="DRAFT") diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_cato_networks.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_cato_networks.py index 428f2faf041..c23fbc0234e 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_cato_networks.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_cato_networks.py @@ -93,26 +93,26 @@ async def test_block_callback(mode: str): ], } - with pytest.raises(HTTPException, match="Jailbreak detected"): - with patch( - "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", - return_value=Response( - json={ - "analysis_result": { - "analysis_time_ms": 212, - "policy_drill_down": {}, - "session_entities": [], - }, - "required_action": { - "action_type": "block_action", - "detection_message": "Jailbreak detected", - "policy_name": "blocking policy", - }, + with patch( + "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", + return_value=Response( + json={ + "analysis_result": { + "analysis_time_ms": 212, + "policy_drill_down": {}, + "session_entities": [], }, - status_code=200, - request=Request(method="POST", url="http://cato"), - ), - ): + "required_action": { + "action_type": "block_action", + "detection_message": "Jailbreak detected", + "policy_name": "blocking policy", + }, + }, + status_code=200, + request=Request(method="POST", url="http://cato"), + ), + ): + async def _call_guardrail(): if mode == "pre_call": await cato_guardrail.async_pre_call_hook( data=data, @@ -127,6 +127,9 @@ async def test_block_callback(mode: str): call_type="completion", ) + with pytest.raises(HTTPException, match="Jailbreak detected"): + await _call_guardrail() + @pytest.mark.asyncio @pytest.mark.parametrize("mode", ["pre_call", "during_call"]) diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_enkryptai.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_enkryptai.py index e6c94a4c3cd..d31f462a185 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_enkryptai.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_enkryptai.py @@ -178,7 +178,7 @@ class TestEnkryptAIGuardrailHooks: with patch.object( enkryptai_guardrail.async_handler, "post", return_value=mock_response ): - with pytest.raises(ValueError) as exc_info: + with pytest.raises(ValueError, match='violation\\(s\\) detected') as exc_info: await enkryptai_guardrail.async_pre_call_hook( user_api_key_dict=mock_user_api_key_dict, cache=MagicMock(), diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py index 5be0d43c250..523ec1a37b4 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py @@ -767,7 +767,7 @@ class TestErrorHandling: "API Error", request=MagicMock(), response=MagicMock(status_code=500) ), ): - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Generic Guardrail API failed: API Error') as exc_info: await generic_guardrail.apply_guardrail( inputs={"texts": ["test"]}, request_data=mock_request_data_input, @@ -786,7 +786,7 @@ class TestErrorHandling: "post", side_effect=httpx.RequestError("Connection failed", request=MagicMock()), ): - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Generic Guardrail API failed: Connection failed') as exc_info: await generic_guardrail.apply_guardrail( inputs={"texts": ["test"]}, request_data=mock_request_data_input, @@ -810,7 +810,7 @@ class TestErrorHandling: "post", side_effect=httpx.RequestError("Connection failed", request=MagicMock()), ): - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Generic Guardrail API failed: Connection failed') as exc_info: await guardrail.apply_guardrail( inputs={"texts": ["test"]}, request_data=mock_request_data_input, diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_microsoft_purview.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_microsoft_purview.py index cc89cea58d2..4a7a14fceaa 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_microsoft_purview.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_microsoft_purview.py @@ -2441,7 +2441,7 @@ class TestStreamingIteratorHook: ), ): chunks = [] - with pytest.raises(HTTPException) as exc_info: + async def _drain(): async for chunk in guardrail.async_post_call_streaming_iterator_hook( user_api_key_dict=UserAPIKeyAuth( api_key="test", user_id="user-123" @@ -2451,6 +2451,9 @@ class TestStreamingIteratorHook: ): chunks.append(chunk) + with pytest.raises(HTTPException) as exc_info: + await _drain() + assert exc_info.value.status_code == 400 assert len(chunks) == 0 # No chunks yielded before the block @@ -2477,7 +2480,7 @@ class TestStreamingIteratorHook: "litellm.main.stream_chunk_builder", return_value=assembled_response ): chunks = [] - with pytest.raises(HTTPException) as exc_info: + async def _drain(): async for chunk in guardrail.async_post_call_streaming_iterator_hook( user_api_key_dict=UserAPIKeyAuth(api_key="test"), # no user_id response=fake_response_stream(), @@ -2485,6 +2488,9 @@ class TestStreamingIteratorHook: ): chunks.append(chunk) + with pytest.raises(HTTPException) as exc_info: + await _drain() + assert exc_info.value.status_code == 400 assert len(chunks) == 0 @@ -2625,7 +2631,7 @@ class TestStreamingIteratorHook: ), ): chunks = [] - with pytest.raises(HTTPException) as exc_info: + async def _drain(): async for chunk in guardrail.async_post_call_streaming_iterator_hook( user_api_key_dict=UserAPIKeyAuth( api_key="test", user_id="user-123" @@ -2635,6 +2641,9 @@ class TestStreamingIteratorHook: ): chunks.append(chunk) + with pytest.raises(HTTPException) as exc_info: + await _drain() + assert exc_info.value.status_code == 400 assert len(chunks) == 0 diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_model_armor.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_model_armor.py index 89b6af27719..14c0d2f9435 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_model_armor.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_model_armor.py @@ -2334,7 +2334,7 @@ async def test_async_moderation_hook_api_error_fail_on_error_true(): } # Should raise the exception since fail_on_error is True - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match="API Error") as exc_info: await guardrail.async_moderation_hook( data=request_data, user_api_key_dict=mock_user_api_key_dict, @@ -2374,7 +2374,7 @@ async def test_async_moderation_hook_api_error_fail_on_error_false(): # Even with fail_on_error=False, the decorator may still raise the exception # This test verifies that the exception is properly logged and handled - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match="API Error") as exc_info: await guardrail.async_moderation_hook( data=request_data, user_api_key_dict=mock_user_api_key_dict, @@ -2865,7 +2865,7 @@ async def test_skip_unscannable_still_fails_closed_on_api_error(): "post", AsyncMock(side_effect=Exception("model armor upstream 500")), ): - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='model armor upstream') as exc_info: await guardrail.async_pre_call_hook( user_api_key_dict=UserAPIKeyAuth(), cache=MagicMock(spec=DualCache), diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_panw_prisma_airs.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_panw_prisma_airs.py index 86a7ac1dabe..8f29ba66814 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_panw_prisma_airs.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_panw_prisma_airs.py @@ -5652,7 +5652,7 @@ class TestPanwAirsTimeoutCoercion: assert isinstance(params.timeout, float) def test_litellm_params_rejects_garbage_timeout(self): - with pytest.raises(ValueError): + with pytest.raises(ValueError, match='validation error for LitellmParams'): LitellmParams( guardrail="panw_prisma_airs", mode="pre_call", @@ -5902,7 +5902,7 @@ class TestPanwAirsBlockedErrorDetailPassthrough: with patch.object( base_handler, "_call_panw_api", return_value=copy.deepcopy(self._FULL_BLOCK_RESPONSE) ): - with pytest.raises(HTTPException) as exc_info: + async def _call_hook(): if is_response: await base_handler.async_post_call_success_hook( data=safe_prompt_data, @@ -5917,6 +5917,9 @@ class TestPanwAirsBlockedErrorDetailPassthrough: call_type="completion", ) + with pytest.raises(HTTPException) as exc_info: + await _call_hook() + error = exc_info.value.detail["error"] for field, value in self._FULL_BLOCK_RESPONSE.items(): if field == "category": diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py index 253d989f203..c779150ad3e 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py @@ -22,6 +22,7 @@ from litellm.proxy.guardrails.guardrail_hooks.presidio import ( from litellm.exceptions import GuardrailRaisedException from litellm.types.guardrails import LitellmParams, PiiAction, PiiEntityType from litellm.types.utils import Choices, Message, ModelResponse +from litellm.exceptions import BlockedPiiEntityError def _make_mock_session_iterator( @@ -1345,7 +1346,7 @@ def test_blocking_respects_threshold_filter(): {"entity_type": PiiEntityType.CREDIT_CARD, "score": 0.95, "start": 0, "end": 4} ] filtered_high = guardrail.filter_analyze_results_by_score(high_score_results) - with pytest.raises(Exception): + with pytest.raises(BlockedPiiEntityError): guardrail.raise_exception_if_blocked_entities_detected(filtered_high) diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_straiker.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_straiker.py index a3d86034f70..81604e22c87 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_straiker.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_straiker.py @@ -90,12 +90,12 @@ def test_config_model_wiring(): def test_init_rejects_empty_api_key(): - with pytest.raises(ValueError): + with pytest.raises(ValueError, match='api_key must be non-empty'): StraikerGuardrail(api_key="") def test_init_rejects_invalid_fallback(): - with pytest.raises(ValueError): + with pytest.raises(ValueError, match="unreachable_fallback must be 'fail_open' or 'fail_closed';"): StraikerGuardrail(api_key="k", unreachable_fallback="nope") @@ -109,7 +109,7 @@ def test_supported_hooks_limited_to_pre_and_post(): def test_during_call_mode_rejected_at_init(): - with pytest.raises(ValueError): + with pytest.raises(ValueError, match='Event hook GuardrailEventHooks\\.during_call is not in the'): StraikerGuardrail(api_key="k", event_hook="during_call") diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_tool_permission.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_tool_permission.py index 4b381b67f0e..0c5addbc143 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_tool_permission.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_tool_permission.py @@ -124,7 +124,7 @@ class TestToolPermissionGuardrail: assert rule_id is None def test_rule_requires_name_or_type(self): - with pytest.raises(ValueError): + with pytest.raises(ValueError, match='validation error for ToolPermissionRule'): ToolPermissionGuardrail( guardrail_name="invalid-rule", rules=[{"id": "no_target", "decision": "allow"}], @@ -1042,7 +1042,7 @@ class TestToolPermissionGuardrailInMemoryUpdate: assert guardrail._check_tool_permission("Secret")[0] is False assert guardrail._check_tool_permission("Other")[0] is True - with pytest.raises(ValueError): + with pytest.raises(ValueError, match="Invalid regex for tool_name in rule 'bad': unterminated"): guardrail.update_in_memory_litellm_params( LitellmParams( guardrail="tool_permission", diff --git a/tests/test_litellm/proxy/guardrails/test_llm_as_a_judge.py b/tests/test_litellm/proxy/guardrails/test_llm_as_a_judge.py index 45dec4ddb2d..bd2553b3280 100644 --- a/tests/test_litellm/proxy/guardrails/test_llm_as_a_judge.py +++ b/tests/test_litellm/proxy/guardrails/test_llm_as_a_judge.py @@ -230,7 +230,7 @@ def test_parse_judge_verdict_reraises_when_no_json(): def test_parse_judge_verdict_rejects_json_non_object(): """Valid JSON that is not an object (e.g. a bare list) raises ValueError.""" - with pytest.raises(ValueError): + with pytest.raises(ValueError, match='judge response is not a JSON object'): _parse_judge_verdict("[1, 2, 3]") diff --git a/tests/test_litellm/proxy/guardrails/test_prompt_security_guardrails.py b/tests/test_litellm/proxy/guardrails/test_prompt_security_guardrails.py index f35d64b89e3..c8f22e6c15e 100644 --- a/tests/test_litellm/proxy/guardrails/test_prompt_security_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/test_prompt_security_guardrails.py @@ -472,9 +472,9 @@ async def test_file_sanitization_block(): async def mock_get(*args, **kwargs): return mock_poll_response - with pytest.raises(HTTPException) as excinfo: - with patch.object(guardrail.async_handler, "post", side_effect=mock_post): - with patch.object(guardrail.async_handler, "get", side_effect=mock_get): + with patch.object(guardrail.async_handler, "post", side_effect=mock_post): + with patch.object(guardrail.async_handler, "get", side_effect=mock_get): + with pytest.raises(HTTPException) as excinfo: await guardrail.apply_guardrail( inputs=inputs, request_data=request_data, diff --git a/tests/test_litellm/proxy/hooks/test_sensitive_data_routing.py b/tests/test_litellm/proxy/hooks/test_sensitive_data_routing.py index 78d2c3af0f3..35c0f8deaf1 100644 --- a/tests/test_litellm/proxy/hooks/test_sensitive_data_routing.py +++ b/tests/test_litellm/proxy/hooks/test_sensitive_data_routing.py @@ -227,7 +227,7 @@ class TestCustomGuardrailSensitiveDataRouting: request_data = {"model": "gpt-4"} - with pytest.raises(ValueError) as exc_info: + with pytest.raises(ValueError, match='Cannot route sensitive data without a session_id\\. Ensure') as exc_info: guardrail.raise_sensitive_data_route_exception( route_to_model="on-premise-model", request_data=request_data, diff --git a/tests/test_litellm/proxy/hooks/test_tpm_concurrent.py b/tests/test_litellm/proxy/hooks/test_tpm_concurrent.py index 8d03857c917..2839acab6b0 100644 --- a/tests/test_litellm/proxy/hooks/test_tpm_concurrent.py +++ b/tests/test_litellm/proxy/hooks/test_tpm_concurrent.py @@ -150,7 +150,7 @@ async def test_no_leak_on_over_limit_rejection(rate_limiter): f"estimated={estimated}, limit={user_api_key_dict.tpm_limit}" ) - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Limit type: tokens\\. Current limit') as exc_info: await handler.async_pre_call_hook( user_api_key_dict=user_api_key_dict, cache=cache, @@ -685,7 +685,7 @@ async def test_contentless_request_reserves_minimum(rate_limiter): f"counter should be 2, got {counter_after_two}" ) - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Limit type: tokens\\. Current limit') as exc_info: await handler.async_pre_call_hook( user_api_key_dict=user_api_key_dict, cache=cache, @@ -1319,7 +1319,7 @@ async def test_project_otpm_rejects_multiple_completion_candidates(rate_limiter) "n": 10, } - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Rate limit exceeded for model_per_project_otpm') as exc_info: await handler.async_pre_call_hook( user_api_key_dict=user_api_key_dict, cache=cache, @@ -1347,7 +1347,7 @@ async def test_project_otpm_reserves_largest_conflicting_output_cap(rate_limiter "max_completion_tokens": 100, } - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Rate limit exceeded for model_per_project_otpm') as exc_info: await handler.async_pre_call_hook( user_api_key_dict=user_api_key_dict, cache=cache, @@ -1377,7 +1377,7 @@ async def test_project_otpm_rejects_google_genai_native_output_cap( project_metadata={"model_otpm_limit": {model: 50}}, ) - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Rate limit exceeded for model_per_project_otpm') as exc_info: await handler.async_pre_call_hook( user_api_key_dict=user_api_key_dict, cache=cache, @@ -1411,7 +1411,7 @@ async def test_project_otpm_rejects_google_genai_native_candidate_count( project_metadata={"model_otpm_limit": {model: 150}}, ) - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Rate limit exceeded for model_per_project_otpm') as exc_info: await handler.async_pre_call_hook( user_api_key_dict=user_api_key_dict, cache=cache, @@ -1500,7 +1500,7 @@ async def test_project_otpm_over_limit_rolls_back_itpm_reservation(rate_limiter) "max_tokens": 500, # blows past the 10-token OTPM limit } - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Rate limit exceeded for model_per_project_otpm') as exc_info: await handler.async_pre_call_hook( user_api_key_dict=user_api_key_dict, cache=cache, @@ -2003,7 +2003,7 @@ async def test_otpm_rejection_does_not_double_refund_combined_tpm(rate_limiter): rate_limit_type="tokens", ) - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Rate limit exceeded for model_per_project_otpm') as exc_info: await handler.async_pre_call_hook( user_api_key_dict=user_api_key_dict, cache=cache, @@ -2060,7 +2060,7 @@ async def test_project_itpm_rejects_pretokenized_embedding_input( "input": embedding_input, } - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Rate limit exceeded for model_per_project_itpm') as exc_info: await handler.async_pre_call_hook( user_api_key_dict=user_api_key_dict, cache=cache, @@ -2250,7 +2250,7 @@ async def test_itpm_reservation_accounts_for_audio_content_not_just_text(rate_li ], } - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Rate limit exceeded for model_per_project_itpm') as exc_info: await handler.async_pre_call_hook( user_api_key_dict=user_api_key_dict, cache=cache, @@ -2361,7 +2361,7 @@ async def test_itpm_rejects_large_audio_payload_that_would_pass_flat_estimate( ], } - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Rate limit exceeded for model_per_project_itpm') as exc_info: await handler.async_pre_call_hook( user_api_key_dict=user_api_key_dict, cache=cache, @@ -2622,7 +2622,7 @@ async def test_explicit_zero_output_responses_call_reserves_effective_provider_m }, ) - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Rate limit exceeded for model_per_project_otpm') as exc_info: await handler.async_pre_call_hook( user_api_key_dict=user_api_key_dict, cache=cache, @@ -2850,7 +2850,7 @@ async def test_otpm_rejection_releases_stashed_parallel_slot(rate_limiter): "rate_limit": {"tokens_per_unit": 5, "window_size": 60}, } - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Rate limit exceeded for model_per_project_otpm') as exc_info: await handler._reserve_project_io_tokens_or_raise( descriptors=[otpm_descriptor], data=data, @@ -3296,7 +3296,7 @@ async def test_rerank_query_and_documents_enforce_project_itpm( project_metadata={"model_itpm_limit": {"rerank-model": 100}}, ) - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Rate limit exceeded for model_per_project_itpm') as exc_info: await handler.async_pre_call_hook( user_api_key_dict=user_api_key_dict, cache=cache, diff --git a/tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py b/tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py index 5caca95095f..0a9efd40b48 100644 --- a/tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py @@ -1,8 +1,13 @@ +import logging import time -from unittest.mock import AsyncMock +from collections.abc import Mapping +from itertools import chain +from typing import Final +from unittest.mock import AsyncMock, MagicMock, call import pytest from fastapi import HTTPException +from pytest_mock import MockerFixture from litellm.proxy._types import ( LiteLLM_TeamTable, @@ -53,6 +58,7 @@ from litellm.types.proxy.management_endpoints.scim_v2 import ( SCIMUserGroup, SCIMUserName, ) +from litellm.proxy._types import ProxyException @pytest.mark.asyncio @@ -71,6 +77,7 @@ async def test_create_user_existing_user_conflict(mocker): mock_prisma_client.db = mocker.MagicMock() mock_prisma_client.db.litellm_usertable = mocker.MagicMock() mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value={"user_id": "existing-user"}) + mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=()) # Mock the _get_prisma_client_or_raise_exception to return our mock mocker.patch( @@ -107,6 +114,7 @@ async def test_create_user_defaults_to_viewer(mocker, monkeypatch): mock_prisma_client.db = mocker.MagicMock() mock_prisma_client.db.litellm_usertable = mocker.MagicMock() mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) + mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=()) mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=None) monkeypatch.setattr("litellm.default_internal_user_params", None, raising=False) @@ -157,6 +165,7 @@ async def test_create_user_ingests_enterprise_extension(mocker, monkeypatch): mock_prisma_client.db = mocker.MagicMock() mock_prisma_client.db.litellm_usertable = mocker.MagicMock() mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) + mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=()) mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=None) monkeypatch.setattr("litellm.default_internal_user_params", None, raising=False) @@ -213,6 +222,7 @@ async def test_create_user_ingests_entitlements_and_roles(mocker, monkeypatch): mock_prisma_client.db = mocker.MagicMock() mock_prisma_client.db.litellm_usertable = mocker.MagicMock() mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) + mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=()) mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=None) monkeypatch.setattr("litellm.default_internal_user_params", None, raising=False) @@ -262,6 +272,7 @@ async def test_create_user_uses_default_internal_user_params_role(mocker, monkey mock_prisma_client.db = mocker.MagicMock() mock_prisma_client.db.litellm_usertable = mocker.MagicMock() mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) + mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=()) mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=None) # Set default_internal_user_params with a specific role @@ -361,6 +372,7 @@ async def test_scim_create_user_respects_default_role_set_via_ui(mocker, monkeyp mock_prisma_client.db = mocker.MagicMock() mock_prisma_client.db.litellm_usertable = mocker.MagicMock() mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) + mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=()) mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=None) mocker.patch( @@ -1281,6 +1293,7 @@ async def test_update_group_metadata_serialization_issue(mocker): mock_user.user_email = "user1@example.com" # Add proper string value for user_email mock_user.teams = [group_id] mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=mock_user) + mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=()) mock_prisma_client.db.litellm_usertable.update = AsyncMock(return_value=mock_user) # Mock the _get_prisma_client_or_raise_exception to return our mock @@ -1483,6 +1496,7 @@ async def test_update_group_e2e(mocker): mock_user = mocker.MagicMock() mock_user.user_id = "test-user" mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=mock_user) + mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=()) # Mock dependencies mocker.patch( @@ -1617,6 +1631,8 @@ async def test_create_group_with_nonexistent_users_rejects(mocker, monkeypatch): return None # new-user-1 and new-user-2 don't exist mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(side_effect=mock_user_lookup) + mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=None) + mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=[]) # Mock dependencies mocker.patch( @@ -1701,6 +1717,8 @@ async def test_update_group_with_nonexistent_users_rejects(mocker, monkeypatch): return None # new-user-3 and new-user-4 don't exist mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(side_effect=mock_user_lookup) + mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=None) + mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=[]) # Mock dependencies mocker.patch( @@ -1770,6 +1788,8 @@ async def test_create_group_with_nonexistent_users_creates_when_flag_true(mocker return None # new-user-1 and new-user-2 don't exist mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(side_effect=mock_user_lookup) + mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=None) + mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=[]) # Mock user creation created_user_1 = NewUserResponse(user_id="new-user-1", key="test-key-1") @@ -1858,6 +1878,8 @@ async def test_extract_group_member_ids_with_flag_true_creates_users(mocker, mon return None # new-user-1 doesn't exist mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(side_effect=mock_user_lookup) + mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=None) + mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=[]) # Mock user creation created_user = NewUserResponse(user_id="new-user-1", key="test-key-1") @@ -1926,6 +1948,8 @@ async def test_extract_group_member_ids_with_flag_false_rejects(mocker, monkeypa return None # new-user-1 doesn't exist mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(side_effect=mock_user_lookup) + mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=None) + mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=[]) # Mock dependencies mocker.patch( @@ -1975,6 +1999,8 @@ async def test_process_group_patch_operations_with_flag_true_creates_users(mocke # Mock user lookup - new-user-1 doesn't exist mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) + mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=None) + mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=[]) mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=None) # Mock user creation @@ -2030,6 +2056,8 @@ async def test_process_group_patch_operations_with_flag_false_rejects(mocker, mo # Mock user lookup - new-user-1 doesn't exist mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) + mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=None) + mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=[]) mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=None) # Execute the function - should raise HTTPException @@ -2069,6 +2097,7 @@ async def test_create_user_grants_admin_when_in_scim_admin_group(mocker, monkeyp mock_prisma_client.db = mocker.MagicMock() mock_prisma_client.db.litellm_usertable = mocker.MagicMock() mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) + mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=()) mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=None) mocker.patch( @@ -2113,6 +2142,7 @@ async def test_create_user_keeps_default_when_not_in_scim_admin_group(mocker, mo mock_prisma_client.db = mocker.MagicMock() mock_prisma_client.db.litellm_usertable = mocker.MagicMock() mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) + mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=()) mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=None) mocker.patch( @@ -2462,6 +2492,7 @@ def _scim_admin_prisma(mocker, *, user_teams): prisma.db = mocker.MagicMock() prisma.db.litellm_usertable = mocker.MagicMock() prisma.db.litellm_usertable.find_unique = AsyncMock(return_value=user) + prisma.db.litellm_usertable.find_many = AsyncMock(return_value=()) prisma.db.litellm_usertable.update = AsyncMock(return_value=user) prisma.db.litellm_teamtable = mocker.MagicMock() prisma.db.litellm_teamtable.find_unique = AsyncMock(side_effect=_team_find_unique) @@ -2560,6 +2591,7 @@ async def test_update_group_recomputes_roles_for_changed_members(mocker): mock_prisma_client.db.litellm_teamtable.update = AsyncMock(return_value=existing_team) mock_prisma_client.db.litellm_usertable = mocker.MagicMock() mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=mocker.MagicMock()) + mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=()) mocker.patch( "litellm.proxy.management_endpoints.scim.scim_v2._get_prisma_client_or_raise_exception", @@ -2617,6 +2649,7 @@ async def test_patch_group_recomputes_roles_for_changed_members(mocker): mock_prisma_client.db.litellm_teamtable.update = AsyncMock(return_value=existing_team) mock_prisma_client.db.litellm_usertable = mocker.MagicMock() mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=mocker.MagicMock()) + mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=()) mocker.patch( "litellm.proxy.management_endpoints.scim.scim_v2._get_prisma_client_or_raise_exception", @@ -2670,6 +2703,7 @@ async def test_delete_group_recomputes_roles_for_members(mocker): mock_prisma_client.db.litellm_teamtable.delete = AsyncMock() mock_prisma_client.db.litellm_usertable = mocker.MagicMock() mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=member) + mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=()) mock_prisma_client.db.litellm_usertable.update = AsyncMock() mocker.patch( @@ -2793,6 +2827,7 @@ async def test_create_user_existing_email_upsert_demotes_when_admin_group_set(mo mock_prisma_client.db = mocker.MagicMock() mock_prisma_client.db.litellm_usertable = mocker.MagicMock() mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) + mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=()) mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=existing_user) mock_prisma_client.db.litellm_usertable.update = AsyncMock(return_value={"user_id": "returning-user"}) @@ -2842,6 +2877,7 @@ async def test_create_group_recomputes_roles_for_members(mocker): mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=None) mock_prisma_client.db.litellm_usertable = mocker.MagicMock() mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=mocker.MagicMock()) + mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=()) mocker.patch( "litellm.proxy.management_endpoints.scim.scim_v2._get_prisma_client_or_raise_exception", @@ -2899,6 +2935,7 @@ async def test_update_group_rename_recomputes_retained_members(mocker): mock_prisma_client.db.litellm_teamtable.update = AsyncMock(return_value=existing_team) mock_prisma_client.db.litellm_usertable = mocker.MagicMock() mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=mocker.MagicMock()) + mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=()) mocker.patch( "litellm.proxy.management_endpoints.scim.scim_v2._get_prisma_client_or_raise_exception", @@ -2953,6 +2990,7 @@ async def test_patch_group_rename_recomputes_retained_members(mocker): mock_prisma_client.db.litellm_teamtable.update = AsyncMock(return_value=existing_team) mock_prisma_client.db.litellm_usertable = mocker.MagicMock() mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=mocker.MagicMock()) + mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=()) mocker.patch( "litellm.proxy.management_endpoints.scim.scim_v2._get_prisma_client_or_raise_exception", @@ -3018,6 +3056,7 @@ async def test_process_group_patch_operations_add_retains_existing_members(mocke mock_prisma_client.db.litellm_usertable = mocker.MagicMock() # new-user already exists in the DB mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=mocker.MagicMock(user_id="new-user")) + mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=()) _, final_members, _ = await _process_group_patch_operations( patch_ops=patch_ops, @@ -3058,6 +3097,7 @@ async def test_process_group_patch_operations_remove_uses_members_with_roles(moc mock_prisma_client.db = mocker.MagicMock() mock_prisma_client.db.litellm_usertable = mocker.MagicMock() mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=mocker.MagicMock(user_id="drop-user")) + mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=()) _, final_members, _ = await _process_group_patch_operations( patch_ops=patch_ops, @@ -3089,6 +3129,7 @@ async def test_get_groups_reports_members_from_members_with_roles(mocker): mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( return_value=mocker.MagicMock(user_id="member-1", user_email="member-1@example.com") ) + mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=()) mocker.patch( "litellm.proxy.management_endpoints.scim.scim_v2._get_prisma_client_or_raise_exception", @@ -3216,7 +3257,7 @@ async def test_delete_user_surfaces_prune_failure_and_keeps_user(mocker): AsyncMock(side_effect=Exception("database connection lost")), ) - with pytest.raises(Exception): + with pytest.raises(ProxyException): await delete_user(user_id=user_id) mock_prisma_client.db.litellm_usertable.delete.assert_not_awaited() @@ -3309,6 +3350,7 @@ async def test_patch_group_add_applies_delta_and_keeps_concurrent_add(mocker): mock_prisma_client.db.litellm_teamtable.update = AsyncMock(return_value=final_team) mock_prisma_client.db.litellm_usertable = mocker.MagicMock() mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=mocker.MagicMock()) + mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=()) mocker.patch( "litellm.proxy.management_endpoints.scim.scim_v2._get_prisma_client_or_raise_exception", @@ -3402,6 +3444,7 @@ async def test_patch_group_replace_stays_absolute_against_concurrent_roster(mock mock_prisma_client.db.litellm_teamtable.update = AsyncMock(return_value=final_team) mock_prisma_client.db.litellm_usertable = mocker.MagicMock() mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=mocker.MagicMock()) + mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=()) mocker.patch( "litellm.proxy.management_endpoints.scim.scim_v2._get_prisma_client_or_raise_exception", @@ -3502,6 +3545,7 @@ async def test_process_group_patch_remove_filtered_path_without_value(mocker): prisma_client.db = mocker.MagicMock() prisma_client.db.litellm_usertable = mocker.MagicMock() prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=LiteLLM_UserTable(user_id="user-1")) + prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=()) _, final_members, _ = await _process_group_patch_operations( patch_ops=patch_ops, @@ -3531,6 +3575,7 @@ async def test_process_group_patch_add_filtered_path_without_value(mocker): prisma_client.db = mocker.MagicMock() prisma_client.db.litellm_usertable = mocker.MagicMock() prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=LiteLLM_UserTable(user_id="user-3")) + prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=()) _, final_members, _ = await _process_group_patch_operations( patch_ops=patch_ops, @@ -3564,6 +3609,7 @@ async def test_process_group_patch_replace_empty_value_does_not_use_path_filter( prisma_client.db = mocker.MagicMock() prisma_client.db.litellm_usertable = mocker.MagicMock() prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=LiteLLM_UserTable(user_id="user-1")) + prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=()) _, final_members, _ = await _process_group_patch_operations( patch_ops=patch_ops, @@ -3574,7 +3620,16 @@ async def test_process_group_patch_replace_empty_value_does_not_use_path_filter( assert final_members == set() -def _member_resolution_prisma(mocker, *, users: set, teams: set, unmanaged_teams: frozenset = frozenset()): +def _member_resolution_prisma( + mocker: MockerFixture, + *, + users: set[str], + teams: set[str], + unmanaged_teams: frozenset[str] = frozenset(), + email_to_user_id: Mapping[str, str] | None = None, + email_to_user_ids: Mapping[str, tuple[str, ...]] | None = None, + sso_user_id_to_user_id: Mapping[str, str] | None = None, +) -> MagicMock: """Prisma mock where only the given ids resolve to a user row / team row. ``teams`` are teams a SCIM group write created, so they carry provenance; @@ -3588,14 +3643,78 @@ def _member_resolution_prisma(mocker, *, users: set, teams: set, unmanaged_teams return LiteLLM_TeamTable(team_id=team_id, metadata={}) return None + def user_row(where: Mapping[str, str]) -> LiteLLM_UserTable | None: + user_id: Final = where["user_id"] + if user_id in users: + return LiteLLM_UserTable(user_id=user_id) + return None + prisma_client = mocker.MagicMock() prisma_client.db = mocker.MagicMock() prisma_client.db.litellm_usertable = mocker.MagicMock() - prisma_client.db.litellm_usertable.find_unique = AsyncMock( - side_effect=lambda where: LiteLLM_UserTable(user_id=where["user_id"]) if where["user_id"] in users else None + prisma_client.db.litellm_usertable.find_unique = AsyncMock(side_effect=user_row) + + emails_to_ids: Final[Mapping[str, tuple[str, ...]]] = ( + dict(email_to_user_ids) + if email_to_user_ids is not None + else ({email: (user_id,) for email, user_id in email_to_user_id.items()} if email_to_user_id else {}) ) + ssos_to_ids: Final[Mapping[str, str]] = dict(sso_user_id_to_user_id) if sso_user_id_to_user_id else {} + + def identity_rows(where: Mapping[str, object], take: int | None = None) -> tuple[LiteLLM_UserTable, ...]: + """Stand-in for the cross-field lookup, honouring the comparison mode + production actually asks for per field, so a field that stops folding case, or + starts folding it, fails here instead of passing. + + A caller that must know which accounts match rather than merely how many + passes take=None, so an unbounded read returns every match. + """ + clauses: Final = where["OR"] + assert isinstance(clauses, list) + fields: Final = tuple(next(iter(clause)) for clause in clauses) + assert fields == ("sso_user_id", "user_email"), fields + + def comparison(clause: Mapping[str, object]) -> tuple[str, bool]: + """The needle and whether production asked for a case-insensitive compare, + read per field so a field that stops folding case fails here.""" + criterion = next(iter(clause.values())) + if isinstance(criterion, str): + return criterion, False + assert isinstance(criterion, dict), criterion + return criterion["equals"], criterion.get("mode") == "insensitive" + + sso_needle, sso_insensitive = comparison(clauses[0]) + email_needle, email_insensitive = comparison(clauses[1]) + + def same(stored: str, needle: str, insensitive: bool) -> bool: + return stored.casefold() == needle.casefold() if insensitive else stored == needle + + matched: Final = tuple( + chain( + ( + user_id + for sso_user_id, user_id in ssos_to_ids.items() + if same(sso_user_id, sso_needle, sso_insensitive) + ), + ( + user_id + for email, user_ids in emails_to_ids.items() + if same(email, email_needle, email_insensitive) + for user_id in user_ids + ), + ) + ) + found: Final = tuple(dict.fromkeys(matched)) + return tuple(LiteLLM_UserTable(user_id=user_id) for user_id in (found[:take] if take else found)) + + def team_lookup(where: Mapping[str, str]) -> LiteLLM_TeamTable | None: + team_id: Final = where["team_id"] + return team_row(team_id) + + prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=None) + prisma_client.db.litellm_usertable.find_many = AsyncMock(side_effect=identity_rows) prisma_client.db.litellm_teamtable = mocker.MagicMock() - prisma_client.db.litellm_teamtable.find_unique = AsyncMock(side_effect=lambda where: team_row(where["team_id"])) + prisma_client.db.litellm_teamtable.find_unique = AsyncMock(side_effect=team_lookup) return prisma_client @@ -4362,6 +4481,581 @@ async def test_resolve_group_member_ids_dedupes_repeated_member(mocker, scim_ups assert result.all_member_ids == ["dup-user"] +def _identity_lookup(value: str) -> object: + """The single cross-field lookup the classifier is expected to issue.""" + return call( + where={"OR": [{"sso_user_id": value}, {"user_email": {"equals": value, "mode": "insensitive"}}]}, + take=2, + ) + + +@pytest.mark.asyncio +async def test_resolve_group_member_ids_matches_sso_user_id(mocker, scim_upsert_user_enabled): + """An OIDC subject in a group payload must resolve to the existing user's + internal id instead of provisioning a placeholder.""" + prisma_client = _member_resolution_prisma( + mocker, + users=set(), + teams=set(), + sso_user_id_to_user_id={"member-sub": "sso-user"}, + ) + create_user_mock = mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2._create_user_if_not_exists", + AsyncMock(), + ) + + result = await _resolve_group_member_ids( + members=[SCIMMember(value="member-sub")], + created_via="scim_group_membership", + prisma_client=prisma_client, + ) + + create_user_mock.assert_not_called() + assert result.existing_member_ids == ["sso-user"] + assert result.created_users == [] + assert result.all_member_ids == ["sso-user"] + assert prisma_client.db.litellm_usertable.find_many.await_args_list == [_identity_lookup("member-sub")] + + +@pytest.mark.asyncio +async def test_resolve_group_member_ids_matches_user_email(mocker, scim_upsert_user_enabled): + """A group member email must resolve to the existing user's internal id + when the identity provider sends email rather than the user id.""" + prisma_client = _member_resolution_prisma( + mocker, + users=set(), + teams=set(), + email_to_user_id={"member@example.com": "email-user"}, + ) + create_user_mock = mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2._create_user_if_not_exists", + AsyncMock(), + ) + + result = await _resolve_group_member_ids( + members=[SCIMMember(value="member@example.com")], + created_via="scim_group_membership", + prisma_client=prisma_client, + ) + + create_user_mock.assert_not_called() + assert result.existing_member_ids == ["email-user"] + assert result.created_users == [] + assert result.all_member_ids == ["email-user"] + assert prisma_client.db.litellm_usertable.find_many.await_args_list == [_identity_lookup("member@example.com")] + + +@pytest.mark.parametrize( + "pushed", + ["MEMBER@EXAMPLE.COM", "Member@Example.com", " member@example.com "], + ids=["upper", "mixed", "padded"], +) +@pytest.mark.asyncio +async def test_resolve_group_member_ids_matches_user_email_as_the_write_path_would( + mocker, scim_upsert_user_enabled, pushed +): + """The member value must be compared the way the layer that would reject a + placeholder compares it. + + ``new_user`` refuses a duplicate email case-insensitively and after stripping, so + a lookup that is stricter than that resolves nothing, creates a placeholder, and + is refused by that same layer, which surfaces as a 500 on the whole group push. + """ + prisma_client = _member_resolution_prisma( + mocker, + users=set(), + teams=set(), + email_to_user_id={"member@example.com": "email-user"}, + ) + create_user_mock = mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2._create_user_if_not_exists", + AsyncMock(return_value=None), + ) + + result = await _resolve_group_member_ids( + members=[SCIMMember(value=pushed)], + created_via="scim_group_membership", + prisma_client=prisma_client, + ) + + create_user_mock.assert_not_called() + assert result.all_member_ids == ["email-user"] + + +@pytest.mark.parametrize( + "population", + [ + {"email_to_user_ids": {"duplicate@example.com": ("email-user-a", "email-user-b")}}, + {"email_to_user_ids": {"duplicate@example.com": ("email-user-a",), "DUPLICATE@EXAMPLE.COM": ("email-user-b",)}}, + { + "sso_user_id_to_user_id": {"duplicate@example.com": "sso-user"}, + "email_to_user_id": {"duplicate@example.com": "email-user"}, + }, + ], + ids=["same-email-twice", "emails-differing-only-in-case", "one-account-by-sso-another-by-email"], +) +@pytest.mark.asyncio +async def test_resolve_group_member_ids_rejects_a_value_naming_two_accounts( + mocker, scim_upsert_user_enabled, caplog, population +): + """A value that names two accounts names a real person we cannot identify, so + the write is refused rather than attributed to one of them. + + Every shape of collision is refused, not just two rows holding the same email + verbatim: rows whose emails differ only in case are one row to the layer that + rejects duplicates, and a value that is one account's SSO identity and another's + email would otherwise be handed to whichever field happened to be searched first. + + It must not fall through to placeholder creation. That path can only fail: the + placeholder carries ``user_email`` set to the member value, which the duplicate + email check rejects, and the recovery lookup that follows searches by ``user_id`` + and so misses the very rows that caused the collision. The operator's data problem + then surfaces as an HTTP 500 the identity provider retries forever. + """ + prisma_client = _member_resolution_prisma(mocker, users=set(), teams=set(), **population) + create_user_mock = mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2._create_user_if_not_exists", + AsyncMock(return_value=None), + ) + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + with pytest.raises(HTTPException) as exc_info: + await _resolve_group_member_ids( + members=[SCIMMember(value="duplicate@example.com")], + created_via="scim_group_membership", + prisma_client=prisma_client, + ) + + assert exc_info.value.status_code == 400 + assert "duplicate@example.com" in str(exc_info.value.detail) + assert "more than one" in str(exc_info.value.detail) + create_user_mock.assert_not_called() + assert any( + record.levelno >= logging.WARNING + and "duplicate@example.com" in record.getMessage() + and "more than one account" in record.getMessage() + for record in caplog.records + ) + + +@pytest.mark.asyncio +async def test_resolve_group_member_ids_does_not_fold_case_on_the_sso_identity(mocker, scim_upsert_user_enabled): + """An email and an SSO identity are not comparable the same way. + + OIDC defines ``sub`` as case-sensitive and nothing folds its case on the way in, + so two subjects differing only in case are two people. Folding it would hand the + group to an account the provider never named, which is the mis-grant the email + comparison is deliberately loose enough to avoid and this one is not. + """ + prisma_client = _member_resolution_prisma( + mocker, + users=set(), + teams=set(), + sso_user_id_to_user_id={"AbC-subject": "other-user"}, + ) + create_user_mock = mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2._create_user_if_not_exists", + AsyncMock(return_value=NewUserResponse(user_id="abc-subject", key="placeholder-key")), + ) + + result = await _resolve_group_member_ids( + members=[SCIMMember(value="abc-subject")], + created_via="scim_group_membership", + prisma_client=prisma_client, + ) + + assert result.existing_member_ids == [] + assert result.all_member_ids == ["abc-subject"] + create_user_mock.assert_awaited_once_with(user_id="abc-subject", created_via="scim_group_membership") + + +@pytest.mark.asyncio +async def test_resolve_group_member_ids_ambiguous_email_outranks_upsert_rejection(mocker, scim_upsert_user_disabled): + """Ambiguity does not depend on scim_upsert_user, so the operator gets the + actionable message on either setting rather than being told to create a user that + already exists twice.""" + prisma_client = _member_resolution_prisma( + mocker, + users=set(), + teams=set(), + email_to_user_ids={"duplicate@example.com": ("email-user-a", "email-user-b")}, + ) + + with pytest.raises(HTTPException) as exc_info: + await _resolve_group_member_ids( + members=[SCIMMember(value="duplicate@example.com")], + created_via="scim_group_membership", + prisma_client=prisma_client, + ) + + assert exc_info.value.status_code == 400 + assert "more than one" in str(exc_info.value.detail) + assert "does not exist" not in str(exc_info.value.detail) + + +@pytest.mark.asyncio +async def test_create_group_rejects_ambiguous_member_email(mocker, scim_upsert_user_enabled): + """The refusal reaches the endpoint, so the identity provider sees a 400 on the + group write rather than a 500 it will retry.""" + scim_group = SCIMGroup( + schemas=["urn:ietf:params:scim:schemas:core:2.0:Group"], + id="ambiguous-group", + displayName="Ambiguous Group", + members=[SCIMMember(value="duplicate@example.com")], + ) + mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2._get_prisma_client_or_raise_exception", + AsyncMock( + return_value=_member_resolution_prisma( + mocker, + users=set(), + teams=set(), + email_to_user_ids={"duplicate@example.com": ("email-user-a", "email-user-b")}, + ) + ), + ) + create_user_mock = mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2._create_user_if_not_exists", + AsyncMock(return_value=None), + ) + + with pytest.raises(ProxyException) as exc_info: + await create_group(group=scim_group) + + assert int(exc_info.value.code) == 400 + assert "duplicate@example.com" in str(exc_info.value.message) + create_user_mock.assert_not_called() + + +@pytest.mark.parametrize( + "removed_by", + ["member@example.com", "member-sub"], + ids=["by-email", "by-sso-subject"], +) +@pytest.mark.asyncio +async def test_process_group_patch_remove_by_the_id_the_directory_added_with( + mocker, scim_upsert_user_enabled, removed_by +): + """A directory removes people by the same id it added them with. + + Resolving on add and not on remove would let someone keep a team after the + directory took them out of the group: the roster holds the canonical user id, so + subtracting the email or the subject the request names would match nothing. + """ + patch_ops = SCIMPatchOp( + schemas=["urn:ietf:params:scim:api:messages:2.0:PatchOp"], + Operations=[SCIMPatchOperation(op="remove", path="members", value=[{"value": removed_by}])], + ) + existing_team = LiteLLM_TeamTable( + team_id="parent-group", + team_alias="Parent Group", + members=[], + members_with_roles=[Member(user_id="real-user", role="user"), Member(user_id="keep-user", role="user")], + ) + create_user_mock = mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2._create_user_if_not_exists", + AsyncMock(return_value=NewUserResponse(user_id="phantom-user", key="phantom-key")), + ) + + _, final_members, _ = await _process_group_patch_operations( + patch_ops=patch_ops, + existing_team=existing_team, + prisma_client=_member_resolution_prisma( + mocker, + users={"real-user", "keep-user"}, + teams=set(), + email_to_user_id={"member@example.com": "real-user"}, + sso_user_id_to_user_id={"member-sub": "real-user"}, + ), + ) + + create_user_mock.assert_not_called() + assert final_members == {"keep-user"} + + +@pytest.mark.asyncio +async def test_process_group_patch_remove_still_drops_a_placeholder_by_its_literal_id( + mocker, scim_upsert_user_enabled +): + """An earlier release put unmatched ids on the roster verbatim, so a remove has to + keep clearing the id as written even once it also resolves.""" + patch_ops = SCIMPatchOp( + schemas=["urn:ietf:params:scim:api:messages:2.0:PatchOp"], + Operations=[SCIMPatchOperation(op="remove", path="members", value=[{"value": "legacy@example.com"}])], + ) + existing_team = LiteLLM_TeamTable( + team_id="parent-group", + team_alias="Parent Group", + members=[], + members_with_roles=[Member(user_id="legacy@example.com", role="user"), Member(user_id="keep-user", role="user")], + ) + + _, final_members, _ = await _process_group_patch_operations( + patch_ops=patch_ops, + existing_team=existing_team, + prisma_client=_member_resolution_prisma(mocker, users={"keep-user"}, teams=set()), + ) + + assert final_members == {"keep-user"} + + +@pytest.mark.asyncio +async def test_process_group_patch_remove_when_the_id_turned_ambiguous_after_admission( + mocker, scim_upsert_user_enabled +): + """Ambiguity is a property of the table as it stands, not of the value. + + Someone admitted while their email was theirs alone must stay removable after a + second account takes that email. Resolving the removal against the whole table + would find two accounts, decline to pick, drop nobody, and still answer 200, + leaving the person the directory just removed holding the team. + """ + patch_ops = SCIMPatchOp( + schemas=["urn:ietf:params:scim:api:messages:2.0:PatchOp"], + Operations=[SCIMPatchOperation(op="remove", path="members", value=[{"value": "shared@example.com"}])], + ) + existing_team = LiteLLM_TeamTable( + team_id="parent-group", + team_alias="Parent Group", + members=[], + members_with_roles=[Member(user_id="admitted-user", role="user"), Member(user_id="keep-user", role="user")], + ) + + _, final_members, _ = await _process_group_patch_operations( + patch_ops=patch_ops, + existing_team=existing_team, + # the newcomer took the address but never joined the group + prisma_client=_member_resolution_prisma( + mocker, + users={"admitted-user", "keep-user"}, + teams=set(), + email_to_user_ids={"shared@example.com": ("admitted-user", "newcomer")}, + ), + ) + + assert final_members == {"keep-user"} + + +@pytest.mark.asyncio +async def test_process_group_patch_remove_refuses_a_value_naming_one_member_by_id_and_another_by_email( + mocker, scim_upsert_user_enabled +): + """One value must never revoke two people. + + A SCIM-provisioned account is keyed by its userName, so a canonical user id that + looks like an email is ordinary rather than exotic, and a second account can hold + that address as its email. Counting the id as written and the resolved accounts + separately makes each look singular, and the removal then takes both. + """ + patch_ops = SCIMPatchOp( + schemas=["urn:ietf:params:scim:api:messages:2.0:PatchOp"], + Operations=[SCIMPatchOperation(op="remove", path="members", value=[{"value": "shared@example.com"}])], + ) + existing_team = LiteLLM_TeamTable( + team_id="parent-group", + team_alias="Parent Group", + members=[], + members_with_roles=[ + Member(user_id="shared@example.com", role="user"), + Member(user_id="other-account", role="user"), + ], + ) + + with pytest.raises(HTTPException) as exc_info: + await _process_group_patch_operations( + patch_ops=patch_ops, + existing_team=existing_team, + prisma_client=_member_resolution_prisma( + mocker, + users={"shared@example.com", "other-account"}, + teams=set(), + email_to_user_id={"shared@example.com": "other-account"}, + ), + ) + + assert exc_info.value.status_code == 400 + assert "shared@example.com" in str(exc_info.value.detail) + assert "more than one member of this group" in str(exc_info.value.detail) + + +@pytest.mark.parametrize("position", [0, 1, 2], ids=["first", "middle", "last"]) +@pytest.mark.asyncio +async def test_process_group_patch_remove_finds_the_member_past_the_bounded_read( + mocker, scim_upsert_user_enabled, position +): + """A removal has to know *which* accounts a value names, not merely whether it + names several, so it reads them all. + + An add stops after two matches, which is all it needs to decide the value is + ambiguous. Reusing that bounded read here would silently drop the member whenever + the one on the roster sorted past the cap, which no fixture smaller than the cap + can show. The member is placed at each position so the test cannot pass by luck + of ordering. + """ + strangers = ["stranger-one", "stranger-two"] + sharers = tuple(strangers[:position] + ["admitted-user"] + strangers[position:]) + patch_ops = SCIMPatchOp( + schemas=["urn:ietf:params:scim:api:messages:2.0:PatchOp"], + Operations=[SCIMPatchOperation(op="remove", path="members", value=[{"value": "shared@example.com"}])], + ) + existing_team = LiteLLM_TeamTable( + team_id="parent-group", + team_alias="Parent Group", + members=[], + members_with_roles=[Member(user_id="admitted-user", role="user"), Member(user_id="keep-user", role="user")], + ) + + _, final_members, _ = await _process_group_patch_operations( + patch_ops=patch_ops, + existing_team=existing_team, + prisma_client=_member_resolution_prisma( + mocker, + users={"admitted-user", "keep-user"}, + teams=set(), + email_to_user_ids={"shared@example.com": sharers}, + ), + ) + + assert final_members == {"keep-user"} + + +@pytest.mark.asyncio +async def test_process_group_patch_remove_refuses_when_two_members_share_the_id(mocker, scim_upsert_user_enabled): + """When both accounts a value names are on the roster the removal is genuinely + undecidable, so it fails rather than reporting a removal it did not perform or + revoking a membership the directory did not name.""" + patch_ops = SCIMPatchOp( + schemas=["urn:ietf:params:scim:api:messages:2.0:PatchOp"], + Operations=[SCIMPatchOperation(op="remove", path="members", value=[{"value": "shared@example.com"}])], + ) + existing_team = LiteLLM_TeamTable( + team_id="parent-group", + team_alias="Parent Group", + members=[], + members_with_roles=[Member(user_id="member-a", role="user"), Member(user_id="member-b", role="user")], + ) + + with pytest.raises(HTTPException) as exc_info: + await _process_group_patch_operations( + patch_ops=patch_ops, + existing_team=existing_team, + prisma_client=_member_resolution_prisma( + mocker, + users={"member-a", "member-b"}, + teams=set(), + email_to_user_ids={"shared@example.com": ("member-a", "member-b")}, + ), + ) + + assert exc_info.value.status_code == 400 + assert "shared@example.com" in str(exc_info.value.detail) + assert "more than one member of this group" in str(exc_info.value.detail) + + + +@pytest.mark.asyncio +async def test_resolve_group_member_ids_exact_user_id_wins_when_it_names_nobody_else( + mocker, scim_upsert_user_enabled +): + """The canonical user id stays authoritative, including when the same account also + holds that value as its email, which is how a SCIM-provisioned account is keyed.""" + prisma_client = _member_resolution_prisma( + mocker, + users={"member-id"}, + teams=set(), + email_to_user_id={"member-id": "member-id"}, + ) + create_user_mock = mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2._create_user_if_not_exists", + AsyncMock(), + ) + + result = await _resolve_group_member_ids( + members=[SCIMMember(value="member-id")], + created_via="scim_group_membership", + prisma_client=prisma_client, + ) + + create_user_mock.assert_not_called() + assert result.existing_member_ids == ["member-id"] + assert result.all_member_ids == ["member-id"] + + +@pytest.mark.parametrize( + "population", + [ + {"sso_user_id_to_user_id": {"member-id": "someone-else"}}, + {"email_to_user_id": {"member-id": "someone-else"}}, + ], + ids=["another-account-by-sso", "another-account-by-email"], +) +@pytest.mark.asyncio +async def test_resolve_group_member_ids_refuses_a_user_id_that_names_another_account( + mocker, scim_upsert_user_enabled, caplog, population +): + """An exact user id is checked for collisions like every other match. + + Taking it on sight would hand the group to whichever account happened to be keyed + by the value. The placeholders this bug provisioned are exactly that shape, since + they are keyed by the very id the provider keeps pushing, so on a tenant that + already has them the real account can never win. Refusing names the problem + instead of silently landing on the placeholder again. + """ + prisma_client = _member_resolution_prisma(mocker, users={"member-id"}, teams=set(), **population) + create_user_mock = mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2._create_user_if_not_exists", + AsyncMock(return_value=None), + ) + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + with pytest.raises(HTTPException) as exc_info: + await _resolve_group_member_ids( + members=[SCIMMember(value="member-id")], + created_via="scim_group_membership", + prisma_client=prisma_client, + ) + + assert exc_info.value.status_code == 400 + assert "member-id" in str(exc_info.value.detail) + create_user_mock.assert_not_called() + assert any( + record.levelno >= logging.WARNING and "someone-else" in record.getMessage() for record in caplog.records + ) + + + +@pytest.mark.asyncio +async def test_resolve_group_member_ids_warns_before_creating_unmatched_placeholder( + mocker, scim_upsert_user_enabled, caplog +): + """An unmatched member still follows upsert behavior, but operators receive + a warning before the placeholder can leave an SSO user teamless.""" + prisma_client = _member_resolution_prisma(mocker, users=set(), teams=set()) + create_user_mock = mocker.patch( + "litellm.proxy.management_endpoints.scim.scim_v2._create_user_if_not_exists", + AsyncMock(return_value=NewUserResponse(user_id="placeholder", key="placeholder-key")), + ) + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + result = await _resolve_group_member_ids( + members=[SCIMMember(value="unmatched-id")], + created_via="scim_group_membership", + prisma_client=prisma_client, + ) + + create_user_mock.assert_awaited_once_with(user_id="unmatched-id", created_via="scim_group_membership") + assert result.existing_member_ids == [] + assert result.created_users == [NewUserResponse(user_id="placeholder", key="placeholder-key")] + assert result.all_member_ids == ["unmatched-id"] + assert any( + record.levelno >= logging.WARNING + and "unmatched-id" in record.getMessage() + and "matched no user by user_id, sso_user_id or user_email" in record.getMessage() + and "real account stays teamless" in record.getMessage() + for record in caplog.records + ) + + @pytest.mark.parametrize( "operation", [ @@ -4428,6 +5122,7 @@ async def test_get_groups_members_are_typed_as_users(mocker): mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( return_value=mocker.MagicMock(user_id="member-1", user_email="member-1@example.com") ) + mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=()) mocker.patch( "litellm.proxy.management_endpoints.scim.scim_v2._get_prisma_client_or_raise_exception", @@ -4662,6 +5357,7 @@ async def test_update_group_roster_failure_propagates(mocker): mock_prisma_client.db.litellm_teamtable.update = AsyncMock(return_value=existing_team) mock_prisma_client.db.litellm_usertable = mocker.MagicMock() mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=mocker.MagicMock()) + mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=()) mocker.patch( "litellm.proxy.management_endpoints.scim.scim_v2._get_prisma_client_or_raise_exception", @@ -4714,6 +5410,7 @@ async def test_resolve_group_member_ids_admits_member_created_concurrently(mocke prisma_client.db.litellm_usertable.find_unique = AsyncMock( side_effect=[None, LiteLLM_UserTable(user_id="raced-user")] ) + prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=()) mocker.patch( "litellm.proxy.management_endpoints.scim.scim_v2._create_user_if_not_exists", AsyncMock(return_value=None), diff --git a/tests/test_litellm/proxy/management_endpoints/test_cost_tracking_settings.py b/tests/test_litellm/proxy/management_endpoints/test_cost_tracking_settings.py index 7e83180bfcd..ea86731eba4 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_cost_tracking_settings.py +++ b/tests/test_litellm/proxy/management_endpoints/test_cost_tracking_settings.py @@ -683,3 +683,103 @@ class TestEstimateCostOnPremProvider: assert response.cost_per_request == pytest.approx(0.002) assert response.input_cost_per_token == pytest.approx(0.000001) assert response.output_cost_per_token == pytest.approx(0.000002) + + + + +class TestBlockRequestsForModelsWithoutPricing: + """Test suite for the block_requests_for_models_without_pricing toggle endpoints""" + + @pytest.mark.asyncio + async def test_get_reflects_in_memory_flag(self): + with patch.object(litellm, "block_requests_for_models_without_pricing", True): + response = client.get( + "/config/block_requests_for_models_without_pricing", + headers={"Authorization": "Bearer sk-1234"}, + ) + + assert response.status_code == 200 + assert response.json() == {"enabled": True} + + @pytest.mark.asyncio + async def test_patch_persists_and_updates_flag(self): + mock_proxy_config = AsyncMock() + mock_proxy_config.get_config = AsyncMock(return_value={"litellm_settings": {}}) + mock_proxy_config.save_config = AsyncMock() + + with ( + patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), + patch("litellm.proxy.proxy_server.proxy_config", mock_proxy_config), + patch("litellm.proxy.proxy_server.store_model_in_db", True), + patch.object(litellm, "block_requests_for_models_without_pricing", False), + ): + response = client.patch( + "/config/block_requests_for_models_without_pricing", + headers={"Authorization": "Bearer sk-1234"}, + json={"enabled": True}, + ) + + assert response.status_code == 200 + assert response.json() == {"enabled": True} + assert litellm.block_requests_for_models_without_pricing is True + + saved_config = mock_proxy_config.save_config.call_args.kwargs["new_config"] + assert saved_config["litellm_settings"]["block_requests_for_models_without_pricing"] is True + + def test_peer_workers_pick_up_persisted_flag_on_config_reload(self): + """A PATCH only mutates the flag on the worker that served it; peer workers must pick the + persisted value up when they reload litellm_settings from the DB.""" + from litellm.proxy.proxy_server import ProxyConfig + + with patch.object(litellm, "block_requests_for_models_without_pricing", False): + ProxyConfig()._update_config_fields( + current_config={}, + param_name="litellm_settings", + db_param_value={"block_requests_for_models_without_pricing": True}, + ) + + assert litellm.block_requests_for_models_without_pricing is True + + @pytest.mark.asyncio + @pytest.mark.parametrize("loads_config_overrides", [True, False]) + async def test_periodic_db_sync_applies_flag_to_peer_worker(self, loads_config_overrides): + """The ~10s reconcile loop runs _init_non_llm_objects_in_db on every worker; it must apply + the persisted flag so peers converge without a restart, including when supported_db_objects + leaves config_overrides out.""" + from types import SimpleNamespace + + from litellm.proxy.proxy_server import ProxyConfig + + config_record = SimpleNamespace( + param_value={"block_requests_for_models_without_pricing": True, "unsafe_key": "x"} + ) + with ( + patch.object(litellm, "block_requests_for_models_without_pricing", False), + patch.object( + ProxyConfig, + "_should_load_db_object", + side_effect=lambda object_type: loads_config_overrides and object_type == "config_overrides", + ), + patch.object(ProxyConfig, "_init_hashicorp_vault_config_override", AsyncMock()), + patch("litellm.proxy.proxy_server.get_config_param", AsyncMock(return_value=config_record)), + ): + await ProxyConfig()._init_non_llm_objects_in_db(prisma_client=MagicMock()) + + assert litellm.block_requests_for_models_without_pricing is True + assert not hasattr(litellm, "unsafe_key") + + @pytest.mark.asyncio + async def test_patch_requires_store_model_in_db(self): + with ( + patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), + patch("litellm.proxy.proxy_server.proxy_config", AsyncMock()), + patch("litellm.proxy.proxy_server.store_model_in_db", False), + ): + response = client.patch( + "/config/block_requests_for_models_without_pricing", + headers={"Authorization": "Bearer sk-1234"}, + json={"enabled": True}, + ) + + assert response.status_code == 500 + assert "error" in response.json()["detail"] diff --git a/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py index 11b7f4553ac..8a036d7e62f 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py @@ -668,6 +668,8 @@ def test_validate_sort_params(): """ Test that validate_sort_params returns None if sort_by is None """ + from fastapi import HTTPException + from litellm.proxy.management_endpoints.internal_user_endpoints import ( _validate_sort_params, ) @@ -676,7 +678,7 @@ def test_validate_sort_params(): assert _validate_sort_params(None, "desc") is None assert _validate_sort_params("user_id", "asc") == {"user_id": "asc"} assert _validate_sort_params("user_id", "desc") == {"user_id": "desc"} - with pytest.raises(Exception): + with pytest.raises(HTTPException): _validate_sort_params("user_id", "invalid") diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index 8e661af8daa..069cfa01178 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -10179,7 +10179,7 @@ async def test_update_key_creator_reassigned_key_blocked(monkeypatch): mock_request = MagicMock() mock_request.query_params = {} - with pytest.raises(Exception) as exc: + with pytest.raises(Exception, match='User can only create keys for themselves\\. Got') as exc: await update_key_fn( request=mock_request, data=UpdateKeyRequest(key=test_hashed_token, key_alias="hijacked"), @@ -11037,8 +11037,7 @@ class TestKeyAliasSkipValidationOnUnchanged: assert new_alias != existing_alias with pytest.raises(ProxyException): - if new_alias != existing_alias: - _validate_key_alias_format(new_alias) + _validate_key_alias_format(new_alias) @pytest.mark.asyncio async def test_update_key_changed_to_valid_alias_passes( diff --git a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py index 84dee5b05c5..01c0760bd27 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py @@ -2325,7 +2325,7 @@ class TestTemporaryMCPSessionEndpoints: "litellm.proxy.management_endpoints.mcp_management_endpoints.validate_and_normalize_mcp_server_payload", MagicMock(), ): - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='User does not have permission to create temporary mcp') as exc_info: await add_session_mcp_server( payload=payload, user_api_key_dict=non_admin, diff --git a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py index 7e4596d154b..42e96ad8659 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py @@ -140,7 +140,7 @@ class TestModelManagementAuthChecks: @pytest.mark.asyncio async def test_can_user_make_team_model_call_non_premium_fails(self): """Test that non-premium users cannot make team model calls""" - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='You must be a LiteLLM Enterprise user to use this feature\\.') as exc_info: ModelManagementAuthChecks.can_user_make_team_model_call( team_id="test_team", user_api_key_dict=self.admin_user, @@ -195,7 +195,7 @@ class TestModelManagementAuthChecks: ) prisma_client = MockPrismaClient(team_exists=True) - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='You must be a LiteLLM Enterprise user to use this feature\\.') as exc_info: await ModelManagementAuthChecks.allow_team_model_action( model_params=model_params, user_api_key_dict=self.admin_user, @@ -216,7 +216,7 @@ class TestModelManagementAuthChecks: ) prisma_client = MockPrismaClient(team_exists=False) - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match="Team id=nonexistent_team does not exist in db'\\}") as exc_info: await ModelManagementAuthChecks.allow_team_model_action( model_params=model_params, user_api_key_dict=self.admin_user, @@ -257,7 +257,7 @@ class TestModelManagementAuthChecks: ) prisma_client = MockPrismaClient(team_exists=True, user_admin=False) - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match="Team ID=test_team does not match the API key's team") as exc_info: await ModelManagementAuthChecks.can_user_make_model_call( model_params=model_params, user_api_key_dict=self.normal_user, @@ -1483,7 +1483,7 @@ class TestTeamModelUpdate: "litellm.proxy.proxy_server.premium_user", True, ): - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match="does not match the API key's team ID=None, OR you are") as exc_info: await _update_team_model_in_db( db_model=db_model, patch_data=patch_data, @@ -3256,7 +3256,7 @@ class TestPatchModelBlockedAuthGate: new=AsyncMock(return_value=None), ), ): - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match="Only proxy admins can change a model's blocked flag\\.") as exc_info: await patch_model( model_id="m1", patch_data=updateDeployment(blocked=True), diff --git a/tests/test_litellm/proxy/management_endpoints/test_ptu_model_settings.py b/tests/test_litellm/proxy/management_endpoints/test_ptu_model_settings.py index 92a34b5ee7c..a1c38d26b9d 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_ptu_model_settings.py +++ b/tests/test_litellm/proxy/management_endpoints/test_ptu_model_settings.py @@ -58,7 +58,7 @@ def test_model_info_accepts_valid_ptu_fields(): def test_model_info_rejects_non_positive_count(): - with pytest.raises(ValueError): + with pytest.raises(ValueError, match='value_error, input_value'): ModelInfo( id="x", team_id="t", @@ -69,7 +69,7 @@ def test_model_info_rejects_non_positive_count(): def test_model_info_rejects_negative_rate(): - with pytest.raises(ValueError): + with pytest.raises(ValueError, match='value_error, input_value'): ModelInfo( id="x", team_id="t", @@ -82,7 +82,7 @@ def test_model_info_rejects_negative_rate(): def test_model_info_rejects_a_count_beyond_the_cap(): """flat cost multiplies the count by a float, and an unbounded int overflows that conversion, which aborted the rollup for every team rather than skipping one model.""" - with pytest.raises(ValueError): + with pytest.raises(ValueError, match='validation error for ModelInfo'): ModelInfo(id="x", team_id="t", ptu_count=10**400, cost_per_ptu_per_hour=2.0) @@ -95,12 +95,12 @@ def test_model_info_accepts_a_count_at_the_cap(): def test_model_info_rejects_a_non_finite_rate(rate): """NaN compares False against every bound, so a bare `< 0` check let it through and the deployment then accrued a flat cost of nan.""" - with pytest.raises(ValueError): + with pytest.raises(ValueError, match='value_error, input_value'): ModelInfo(id="x", team_id="t", ptu_count=5, cost_per_ptu_per_hour=rate) def test_model_info_rejects_a_rate_beyond_the_cap(): - with pytest.raises(ValueError): + with pytest.raises(ValueError, match='validation error for ModelInfo'): ModelInfo(id="x", team_id="t", ptu_count=5, cost_per_ptu_per_hour=ModelInfo.MAX_COST_PER_PTU_PER_HOUR * 2) @@ -148,7 +148,7 @@ def test_validate_helper_passes_full_config(): def test_model_info_rejects_effective_to_before_from(): import datetime - with pytest.raises(ValueError): + with pytest.raises(ValueError, match='validation error for ModelInfo'): ModelInfo( id="x", team_id="t", @@ -186,7 +186,7 @@ def test_model_info_compares_mixed_naive_and_aware_timestamps(): ) assert info.ptu_effective_to is not None - with pytest.raises(ValueError): + with pytest.raises(ValueError, match='validation error for ModelInfo'): ModelInfo( id="x", team_id="t", @@ -698,7 +698,7 @@ class TestAddNewModelPtuGate: with ExitStack() as stack: for active_patch in patches: stack.enter_context(active_patch) - with pytest.raises(Exception) as exc: + with pytest.raises(Exception, match='PTU cost attribution is disabled, so ptu_count') as exc: await add_new_model(model_params=self._ptu_deployment("ptu-gate-model"), user_api_key_dict=admin) assert PTU_COST_ATTRIBUTION_ENV_VAR in str(exc.value) @@ -1273,7 +1273,7 @@ class TestPtuDeploymentsAreNotBilledPerToken: with ExitStack() as stack: for active_patch in patches: stack.enter_context(active_patch) - with pytest.raises(Exception) as exc: + with pytest.raises(Exception, match='A PTU deployment bills by reserved capacity, so') as exc: await add_new_model(model_params=deployment, user_api_key_dict=admin) assert "input_cost_per_token" in str(exc.value) diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py index 0767288d0bc..e39b09ae073 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -10253,6 +10253,65 @@ async def test_team_info_returns_model_aliases(): assert litellm_model_table.model_aliases == {"gpt-4o": "gpt-4o-team-1"} +@pytest.mark.asyncio +async def test_team_info_hydrates_member_emails_from_the_user_table(): + """/team/info must fill in emails missing from the members_with_roles snapshot. + + members_with_roles is written at add-time, so a member added by user_id alone + carries user_email=None forever. Without this join the Admin UI's member table + shows "-" for a user that has an email on their user row. A stored email is left + exactly as-is. + """ + from fastapi import Request + + from litellm.proxy.management_endpoints import team_endpoints + + team_row = LiteLLM_TeamTable( + team_id="team-1", + members_with_roles=[ + Member(user_id="no-email-on-roster", role="admin"), + Member(user_id="already-stored", user_email="stored@example.com", role="user"), + ], + ) + + mock_prisma = MagicMock() + mock_prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=team_row) + mock_prisma.get_data = AsyncMock(return_value=[]) + + find_many = AsyncMock( + return_value=[ + LiteLLM_UserTable( + user_id="no-email-on-roster", + user_email="real@example.com", + max_budget=None, + spend=0.0, + models=[], + ) + ] + ) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), + patch.object(team_endpoints, "get_all_team_memberships", AsyncMock(return_value=[])), + patch.object(team_endpoints, "UserRepository") as repo, + ): + repo.return_value.table.find_many = find_many + + response = await team_endpoints.team_info( + http_request=MagicMock(spec=Request), + team_id="team-1", + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + ) + + members = response["team_info"].members_with_roles + assert [(m.user_id, m.user_email) for m in members] == [ + ("no-email-on-roster", "real@example.com"), + ("already-stored", "stored@example.com"), + ] + # only the member actually missing an email is looked up + assert find_many.await_args.kwargs["where"] == {"user_id": {"in": ["no-email-on-roster"]}} + + @pytest.mark.asyncio async def test_update_model_table_clears_aliases_with_empty_map(): """``model_aliases={}`` on /team/update must persist an empty map (json.dumps({})) @@ -11369,6 +11428,141 @@ async def test_resolve_existing_member_user_ids_skips_the_query_when_no_user_ids repo.return_value.table.find_many.assert_not_awaited() +def _user_row(user_id: str, user_email: str | None) -> LiteLLM_UserTable: + return LiteLLM_UserTable( + user_id=user_id, user_email=user_email, max_budget=None, spend=0.0, models=[] + ) + + +@pytest.mark.asyncio +async def test_hydrate_member_emails_fills_in_emails_the_roster_snapshot_never_captured(): + """A member added by user_id alone has user_email=None on the stored roster entry. + + /team/info has to fill it in from the user row, or the UI renders "-" for a user + that plainly has an email. + """ + from litellm.proxy.management_endpoints.team_endpoints import _hydrate_member_emails + + find_many = AsyncMock(return_value=[_user_row("by-id", "found@example.com")]) + + with patch("litellm.proxy.management_endpoints.team_endpoints.UserRepository") as repo: + repo.return_value.table.find_many = find_many + + hydrated = await _hydrate_member_emails( + prisma_client=MagicMock(), + members=[Member(user_id="by-id", role="admin")], + ) + + assert [(m.user_id, m.user_email, m.role) for m in hydrated] == [("by-id", "found@example.com", "admin")] + find_many.assert_awaited_once() + assert find_many.await_args.kwargs["where"] == {"user_id": {"in": ["by-id"]}} + + +@pytest.mark.asyncio +async def test_hydrate_member_emails_never_overwrites_a_stored_email(): + """The snapshot wins wherever it has a value - hydration only fills blanks. + + Overwriting would be a real behavior change to /team/info; filling a null is not. + """ + from litellm.proxy.management_endpoints.team_endpoints import _hydrate_member_emails + + find_many = AsyncMock(return_value=[_user_row("has-email", "current@example.com")]) + + with patch("litellm.proxy.management_endpoints.team_endpoints.UserRepository") as repo: + repo.return_value.table.find_many = find_many + + hydrated = await _hydrate_member_emails( + prisma_client=MagicMock(), + members=[Member(user_id="has-email", user_email="stored@example.com", role="user")], + ) + + assert hydrated[0].user_email == "stored@example.com" + # nothing was missing, so no round-trip either + find_many.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_hydrate_member_emails_leaves_members_alone_when_the_user_row_has_no_email(): + """A user row with no email leaves the member as-is rather than inventing one.""" + from litellm.proxy.management_endpoints.team_endpoints import _hydrate_member_emails + + with patch("litellm.proxy.management_endpoints.team_endpoints.UserRepository") as repo: + repo.return_value.table.find_many = AsyncMock(return_value=[_user_row("no-email", None)]) + + hydrated = await _hydrate_member_emails( + prisma_client=MagicMock(), + members=[Member(user_id="no-email", role="user"), Member(user_email="e@example.com", role="user")], + ) + + assert [m.user_email for m in hydrated] == [None, "e@example.com"] + + +@pytest.mark.asyncio +async def test_hydrate_member_emails_skips_the_query_when_every_member_has_one(): + """No blanks means /team/info pays for no extra query.""" + from litellm.proxy.management_endpoints.team_endpoints import _hydrate_member_emails + + with patch("litellm.proxy.management_endpoints.team_endpoints.UserRepository") as repo: + repo.return_value.table.find_many = AsyncMock() + + hydrated = await _hydrate_member_emails( + prisma_client=MagicMock(), + members=[Member(user_id="a", user_email="a@example.com", role="user")], + ) + + assert hydrated[0].user_email == "a@example.com" + repo.return_value.table.find_many.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_update_team_members_list_stamps_email_for_a_member_added_by_user_id(): + """Identity resolution runs both ways, so new roster entries stop being born blank. + + Previously only user_id was backfilled (from email); a member added by user_id + was written with user_email=None forever. + """ + from litellm.proxy.management_endpoints.team_endpoints import ( + _update_team_members_list, + ) + + mock_team = MagicMock(spec=LiteLLM_TeamTable) + mock_team.members_with_roles = [] + + await _update_team_members_list( + data=TeamMemberAddRequest(team_id="test-team-123", member=Member(user_id="new-user-123", role="user")), + complete_team_data=mock_team, + updated_users=[_user_row("new-user-123", "new@example.com")], + ) + + assert len(mock_team.members_with_roles) == 1 + assert mock_team.members_with_roles[0].user_email == "new@example.com" + + +@pytest.mark.asyncio +async def test_update_team_members_list_stamps_email_for_each_member_in_a_bulk_add(): + """Same both-ways resolution for the list branch.""" + from litellm.proxy.management_endpoints.team_endpoints import ( + _update_team_members_list, + ) + + mock_team = MagicMock(spec=LiteLLM_TeamTable) + mock_team.members_with_roles = [] + + await _update_team_members_list( + data=TeamMemberAddRequest( + team_id="test-team-123", + member=[Member(user_id="u1", role="user"), Member(user_email="u2@example.com", role="admin")], + ), + complete_team_data=mock_team, + updated_users=[_user_row("u1", "u1@example.com"), _user_row("u2", "u2@example.com")], + ) + + assert [(m.user_id, m.user_email) for m in mock_team.members_with_roles] == [ + ("u1", "u1@example.com"), + ("u2", "u2@example.com"), + ] + + def test_pre_existing_user_ids_counts_ids_filled_in_by_member_resolution(): """An id the member-resolution step filled in came from a matched row, so it pre-existed. diff --git a/tests/test_litellm/proxy/management_endpoints/usage_endpoints/test_ai_usage_chat.py b/tests/test_litellm/proxy/management_endpoints/usage_endpoints/test_ai_usage_chat.py index e8a74e41dae..3a32b3cc128 100644 --- a/tests/test_litellm/proxy/management_endpoints/usage_endpoints/test_ai_usage_chat.py +++ b/tests/test_litellm/proxy/management_endpoints/usage_endpoints/test_ai_usage_chat.py @@ -458,7 +458,7 @@ class TestUsageAiChatServiceAccountGuard: _resolve_fetch_kwargs, ) - with pytest.raises(ValueError) as exc_info: + with pytest.raises(ValueError, match='Non-admin caller has user_id=None; refusing to issue an') as exc_info: _resolve_fetch_kwargs( fn_name="get_usage_data", fn_args={"start_date": "2025-01-01", "end_date": "2025-01-31"}, diff --git a/tests/test_litellm/proxy/management_helpers/test_team_metadata_validation.py b/tests/test_litellm/proxy/management_helpers/test_team_metadata_validation.py index d22c3db0f7e..1acb8e7e016 100644 --- a/tests/test_litellm/proxy/management_helpers/test_team_metadata_validation.py +++ b/tests/test_litellm/proxy/management_helpers/test_team_metadata_validation.py @@ -22,6 +22,7 @@ from litellm.proxy.management_helpers.team_metadata_validation import ( run_team_metadata_validation, validate_team_metadata_if_configured, ) +from pydantic import ValidationError def _registry_with(validator): @@ -572,12 +573,15 @@ async def test_http_validator_service_outage_fails_closed(monkeypatch, kind, exi monkeypatch.setenv("TEAM_METADATA_VALIDATION_SERVICE_URL", _closed_port_url()) with _configured(impls.validate_via_http): - with pytest.raises(ProxyException) as exc_info: + async def _drive(): if kind == "create": await _drive_create(metadata=request_payload) else: await _drive_update(kind, existing_metadata, request_payload) + with pytest.raises(ProxyException) as exc_info: + await _drive() + assert str(exc_info.value.code) == "503" assert DEFAULT_TEAM_METADATA_VALIDATION_UNAVAILABLE_MESSAGE in str(exc_info.value.message) @@ -634,7 +638,7 @@ def test_parse_schema_round_trips_fields_in_order(): ], ) def test_parse_schema_malformed_raises(raw): - with pytest.raises(Exception): + with pytest.raises(ValidationError): parse_team_metadata_schema(raw) @@ -667,7 +671,7 @@ async def test_non_callable_validator_is_rejected_with_clean_500(): def test_parse_schema_duplicate_error_lists_offending_keys(): - with pytest.raises(ValueError) as exc_info: + with pytest.raises(ValueError, match='team_metadata_schema contains duplicate keys: app_name') as exc_info: parse_team_metadata_schema( [{"key": "cost_center"}, {"key": "app_name"}, {"key": "cost_center"}, {"key": "app_name"}] ) diff --git a/tests/test_litellm/proxy/openai_files_endpoint/test_batch_guardrails.py b/tests/test_litellm/proxy/openai_files_endpoint/test_batch_guardrails.py index a05b8ae530c..6ce7af1e2ee 100644 --- a/tests/test_litellm/proxy/openai_files_endpoint/test_batch_guardrails.py +++ b/tests/test_litellm/proxy/openai_files_endpoint/test_batch_guardrails.py @@ -781,7 +781,7 @@ async def test_the_rewrite_closes_its_own_output_when_it_cannot_finish(): original_read = bg._read_spooled bg._read_spooled = _boom try: - with pytest.raises(OSError): + with pytest.raises(OSError, match='no space left on device'): rewrite_batch_input_file(source, result) finally: bg.tempfile.SpooledTemporaryFile = real diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py b/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py index f994fba371b..d3237f5f49d 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py @@ -2714,7 +2714,7 @@ class TestMilvusProxyRoute: None ) - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Vector store not found for missing-store') as exc_info: await milvus_proxy_route( endpoint="vectors/search", request=mock_request, @@ -2779,7 +2779,7 @@ class TestMilvusProxyRoute: mock_vector_store ) - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='api_base not found in vector store configuration for') as exc_info: await milvus_proxy_route( endpoint="vectors/search", request=mock_request, @@ -2988,7 +2988,7 @@ class TestOpenAIPassthroughRoute: "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.passthrough_endpoint_router.get_credentials", return_value=None, ): - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match="Required 'OPENAI_API_KEY' in environment to make") as exc_info: await openai_proxy_route( endpoint="v1/chat/completions", request=mock_request, @@ -3177,7 +3177,7 @@ class TestCursorProxyRoute: [], ), ): - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Cursor API key not found\\. Add Cursor credentials via') as exc_info: await cursor_proxy_route( endpoint="v0/agents", request=mock_request, diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py index 00097166c13..090acf2dbb0 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -350,7 +350,7 @@ async def test_pass_through_request_failure_handler(): mock_user_api_key_dict = MagicMock() # Call the function with a target that will trigger an HTTPError - with pytest.raises(Exception): + with pytest.raises(ProxyException): await pass_through_request( request=mock_request, target="http://test.com", @@ -1154,7 +1154,7 @@ async def test_pass_through_request_uses_resolved_timeout(): mock_user_api_key_dict = MagicMock() - with pytest.raises(Exception): + with pytest.raises(TypeError): await pass_through_request( request=mock_request, target="http://test.com", diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_passthrough_post_call_guardrails.py b/tests/test_litellm/proxy/pass_through_endpoints/test_passthrough_post_call_guardrails.py index 470179a0429..a48e9e9e17f 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_passthrough_post_call_guardrails.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_passthrough_post_call_guardrails.py @@ -18,6 +18,7 @@ from litellm.integrations.custom_guardrail import ( CustomGuardrail, ModifyResponseException, ) +from litellm.proxy._types import ProxyException _PT_MOD = "litellm.proxy.pass_through_endpoints.pass_through_endpoints" _COLLECT = "litellm.proxy.pass_through_endpoints.passthrough_guardrails.PassthroughGuardrailHandler.collect_guardrails" @@ -251,7 +252,7 @@ class TestPassthroughPostCallGuardrails: ) with _common_patches(mock_proxy_logging, mock_response): - with pytest.raises(Exception): + with pytest.raises(ProxyException): await pass_through_request( request=_make_mock_request(), target="https://example.com/v1/generateContent", diff --git a/tests/test_litellm/proxy/policy_engine/test_policy_versioning.py b/tests/test_litellm/proxy/policy_engine/test_policy_versioning.py index ebebfde5cd3..b6633779326 100644 --- a/tests/test_litellm/proxy/policy_engine/test_policy_versioning.py +++ b/tests/test_litellm/proxy/policy_engine/test_policy_versioning.py @@ -157,7 +157,7 @@ class TestUpdatePolicyDraftOnly: prod_row = _make_row(policy_id="pid-1", version_status="production") prisma.db.litellm_policytable.find_unique = AsyncMock(return_value=prod_row) - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Error updating policy in DB: Only draft versions can be') as exc_info: await registry.update_policy_in_db( policy_id="pid-1", policy_request=PolicyUpdateRequest(description="new"), @@ -341,7 +341,7 @@ class TestUpdateVersionStatus: draft = _make_row(policy_id="d-1", version_status="draft") prisma.db.litellm_policytable.find_unique = AsyncMock(return_value=draft) - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Error updating version status: Cannot promote draft') as exc_info: await registry.update_version_status( policy_id="d-1", new_status="production", diff --git a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py index d8047cca728..47f01fe096d 100644 --- a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py +++ b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py @@ -28,6 +28,7 @@ from litellm.proxy.proxy_server import ( ) from .conftest import normalize +from pydantic import ValidationError # --------------------------------------------------------------------------- # _is_remote_module_url @@ -393,7 +394,7 @@ def test_ProxyConfig__load_yaml_file_returns_parsed_dict(tmp_path): def test_ProxyConfig__load_yaml_file_raises_on_missing_file(): pc = ProxyConfig() - with pytest.raises(Exception): + with pytest.raises(Exception, match="Error loading yaml file"): pc._load_yaml_file("/no/such/file.yaml") @@ -418,7 +419,7 @@ async def test_ProxyConfig__get_config_from_file_loads_yaml(tmp_path): @pytest.mark.asyncio async def test_ProxyConfig__get_config_from_file_missing_path_raises(): pc = ProxyConfig() - with pytest.raises(Exception): + with pytest.raises(Exception, match="Config file not found"): await pc._get_config_from_file(config_file_path="/no/such/file.yaml") @@ -476,7 +477,7 @@ async def test_ProxyConfig_save_config_invalid_path_raises(monkeypatch): monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", False) monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {}) pc = ProxyConfig() - with pytest.raises(Exception): + with pytest.raises(FileNotFoundError): await pc.save_config({"x": 1}) @@ -641,7 +642,7 @@ def test_ProxyConfig__get_team_config_returns_match(): def test_ProxyConfig__get_team_config_missing_team_id_raises(): pc = ProxyConfig() - with pytest.raises(Exception): + with pytest.raises(Exception, match="team_id missing from team"): pc._get_team_config(team_id="t1", all_teams_config=[{"no_id_field": True}]) @@ -671,7 +672,7 @@ def test_ProxyConfig_load_team_config_no_settings_returns_empty(): assert out == {} # Error-style: a misconfigured team list without team_id raises. pc.config = {"litellm_settings": {"default_team_settings": [{"no_id": True}]}} - with pytest.raises(Exception): + with pytest.raises(Exception, match="team_id missing from team"): pc.load_team_config(team_id="anything") @@ -698,7 +699,7 @@ def test_ProxyConfig__init_cache_sets_litellm_cache(monkeypatch): def test_ProxyConfig__init_cache_invalid_params_raises(): pc = ProxyConfig() - with pytest.raises(Exception): + with pytest.raises(AttributeError): pc._init_cache(cache_params={"type": "this-cache-type-does-not-exist"}) @@ -765,7 +766,7 @@ async def test_ProxyConfig_get_config_missing_file_raises(monkeypatch): monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", False) monkeypatch.delenv("LITELLM_CONFIG_BUCKET_NAME", raising=False) pc = ProxyConfig() - with pytest.raises(Exception): + with pytest.raises(Exception, match="Config file not found"): await pc.get_config(config_file_path="/no/such/path.yaml") @@ -1041,7 +1042,7 @@ def test_ProxyConfig_load_credential_list_returns_items(): def test_ProxyConfig_load_credential_list_invalid_entry_raises(): pc = ProxyConfig() - with pytest.raises(Exception): + with pytest.raises(ValidationError): pc.load_credential_list({"credential_list": [{"missing_required": True}]}) @@ -1381,7 +1382,7 @@ async def test_ProxyConfig_load_config_missing_file_raises(monkeypatch): monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", False) monkeypatch.delenv("LITELLM_CONFIG_BUCKET_NAME", raising=False) pc = ProxyConfig() - with pytest.raises(Exception): + with pytest.raises(Exception, match="Config file not found"): await pc.load_config(router=None, config_file_path="/no/file.yaml") @@ -1492,7 +1493,7 @@ async def test_ProxyConfig__init_non_llm_configs_empty_config(): async def test_ProxyConfig__init_non_llm_configs_premium_invalid_worker_registry_raises(monkeypatch): monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True) pc = ProxyConfig() - with pytest.raises(Exception): + with pytest.raises(ValidationError): await pc._init_non_llm_configs( config={"worker_registry": [{"totally": "invalid"}]}, config_file_path=None, @@ -1503,7 +1504,7 @@ async def test_ProxyConfig__init_non_llm_configs_premium_invalid_worker_registry async def test_ProxyConfig__init_non_llm_configs_worker_registry_requires_premium(monkeypatch): monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", False) pc = ProxyConfig() - with pytest.raises(ValueError) as exc_info: + with pytest.raises(ValueError, match='Trying to use `worker_registry`You must be a LiteLLM') as exc_info: await pc._init_non_llm_configs( config={ "worker_registry": [ @@ -1572,7 +1573,7 @@ async def test_ProxyConfig__init_policy_engine_none_config_noop(): # None config returns early without raising. await pc._init_policy_engine(config=None, prisma_client=None, llm_router=None) # Error-style: invalid policies value should raise. - with pytest.raises(Exception): + with pytest.raises(AttributeError): await pc._init_policy_engine( config={"policies": "not-a-list"}, prisma_client=None, @@ -1601,7 +1602,7 @@ def test_ProxyConfig__load_alerting_settings_noop_when_no_alerting(): def test_ProxyConfig__load_alerting_settings_invalid_alerting_raises(): pc = ProxyConfig() - with pytest.raises(Exception): + with pytest.raises(RuntimeError): # alerting must be iterable — int triggers an error. pc._load_alerting_settings({"alerting": 12345}) @@ -1768,7 +1769,7 @@ def test_ProxyConfig_initialize_secret_manager_none_noop(): def test_ProxyConfig_initialize_secret_manager_invalid_kms_raises(): pc = ProxyConfig() - with pytest.raises(ValueError): + with pytest.raises(ValueError, match='Invalid Key Management System selected'): pc.initialize_secret_manager(key_management_system="not-a-real-kms") @@ -1823,7 +1824,7 @@ async def test_ProxyConfig__delete_deployment_invalid_models_raises(monkeypatch) fake_router.get_model_ids = MagicMock(return_value=[]) monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", fake_router) pc = ProxyConfig() - with pytest.raises(Exception): + with pytest.raises(AttributeError): # Non-model objects without expected attrs trigger an error. await pc._delete_deployment(db_models=[{"not_a_model": True}]) @@ -2721,7 +2722,7 @@ async def test_ProxyConfig__update_general_settings_none_input_noop(): result = await pc._update_general_settings(db_general_settings=None) assert result is None # Error-style: dict() will fail on non-mapping non-None input. - with pytest.raises(Exception): + with pytest.raises(TypeError): await pc._update_general_settings(db_general_settings=12345) # type: ignore[arg-type] @@ -2743,7 +2744,7 @@ def test_ProxyConfig__update_config_fields_merges_dict(): def test_ProxyConfig__update_config_fields_invalid_param_raises(): pc = ProxyConfig() - with pytest.raises(Exception): + with pytest.raises(TypeError): # Missing required arg. pc._update_config_fields(current_config={}, param_name="general_settings") # type: ignore[call-arg] diff --git a/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py b/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py index 03d228cc732..72e7b1c18d6 100644 --- a/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py +++ b/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py @@ -298,6 +298,37 @@ def test_nvidia_riva_provider_fields(): assert fields_by_key["nvcf_function_id"]["required"] is False +def test_cognition_provider_fields(): + """Cognition must be selectable in the Add Model flow (LIT-5348). + + The dropdown is driven entirely by /public/providers/fields, so without an + entry here admins have to fall back to the generic OpenAI-compatible route, + which is exactly the provider identity mix-up this feature removes. + """ + app_instance = FastAPI() + app_instance.include_router(router) + test_client = TestClient(app_instance) + + response = test_client.get("/public/providers/fields") + assert response.status_code == 200 + providers = response.json() + + cognition = next((p for p in providers if p["provider"] == "Cognition"), None) + assert cognition is not None, "Cognition provider entry not found" + + assert cognition["provider_display_name"] == "Cognition" + assert cognition["litellm_provider"] == "cognition" + assert cognition["default_model_placeholder"].startswith("cognition/") + + fields_by_key = {f["key"]: f for f in cognition["credential_fields"]} + + assert fields_by_key["api_key"]["required"] is True + assert fields_by_key["api_key"]["field_type"] == "password" + + assert fields_by_key["api_base"]["field_type"] == "text" + assert fields_by_key["api_base"]["required"] is False + + def test_google_ai_studio_provider_fields_expose_api_base(): """The Google AI Studio (gemini) credential form must let admins set a custom api_base so they can point at a Gemini-compatible gateway (e.g. a self-hosted diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py index 15a3e6609f0..b2ec500d045 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py @@ -3263,7 +3263,7 @@ async def test_provider_budget_over(disable_budget_sync): model_list=MODEL_LIST, ) - with pytest.raises(Exception) as e: + with pytest.raises(Exception, match='No deployments available - crossed budget: Exceeded budget') as e: await router.acompletion( model="azure-gpt-4o", messages=[{"role": "user", "content": "Hello, world!"}], @@ -5096,7 +5096,7 @@ def test_resolve_spend_report_scope_missing_caller_value_400(): @pytest.mark.parametrize("bad_column", ["metadata", "end_user", "evil; DROP TABLE", ""]) def test_scoped_spend_report_sql_rejects_unknown_column(bad_column): - with pytest.raises(ValueError): + with pytest.raises(ValueError, match='Unsupported spend report scope column'): spend_management_endpoints._scoped_spend_report_sql(scope_column=bad_column) diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py index c7db1ef1792..a980fed746c 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py @@ -7,7 +7,6 @@ from datetime import timezone from typing import Any, Final, cast import pytest -from fastapi.testclient import TestClient sys.path.insert( 0, os.path.abspath("../../../..") @@ -29,8 +28,8 @@ from litellm.proxy.spend_tracking.spend_tracking_utils import ( _get_response_for_spend_logs_payload, _get_spend_logs_metadata, _get_vector_store_request_for_spend_logs_payload, - _hash_api_key_for_spend_log, _is_master_key, + _redact_logged_api_key, _redact_prompt_leaks_in_error_string, _sanitize_error_information_for_spend_logs, _sanitize_guardrail_information_for_spend_logs, @@ -39,6 +38,7 @@ from litellm.proxy.spend_tracking.spend_tracking_utils import ( get_logging_payload, get_spend_logs_id, ) +from litellm.proxy.utils import hash_token from litellm.types.utils import ( StandardLoggingHiddenParams, StandardLoggingMetadata, @@ -888,8 +888,6 @@ def test_get_logging_payload_api_key_preserved_when_standard_logging_payload_is_ assert payload["model"] == "openai/gpt-4.1" assert payload["user"] == "test_user" - print(f"✅ Test passed! api_key preserved: {payload['api_key']}") - @pytest.mark.asyncio @patch("litellm.proxy.proxy_server.master_key", "sk-master-key") @@ -1037,18 +1035,6 @@ async def test_api_key_preserved_through_failure_hook_to_database(): assert payload.get("model") == "gpt-3.5-turbo" assert payload.get("user") == "test_user" - print("\n" + "=" * 80) - print("✅ CRITICAL E2E TEST PASSED") - print("=" * 80) - print(f"Token: {data['token']}") - print(f"Payload api_key: {payload_api_key}") - print(f"Match: {data['token'] == payload_api_key}") - print("=" * 80) - print("Production incident bug is FIXED and protected:") - print("- Failed requests preserve api_key through entire flow") - print("- Both SpendLogs AND DailyUserSpend will have correct api_key") - print("=" * 80 + "\n") - @patch("litellm.proxy.proxy_server.master_key", None) @patch("litellm.proxy.proxy_server.general_settings", {}) @@ -2591,6 +2577,219 @@ def test_sanitize_error_information_redacts_pydantic_assignment_form( assert REDACTED_BY_LITELM_STRING in sanitized["error_message"] +# ── _redact_logged_api_key unit tests ────────────────────────────────────── + + +def test_redact_logged_api_key_none_returns_none(): + assert _redact_logged_api_key(None) is None + + +def test_redact_logged_api_key_empty_string_returns_none(): + assert _redact_logged_api_key("") is None + + +def test_redact_logged_api_key_sk_key_is_hashed(): + raw = "sk-1234secret" + result = _redact_logged_api_key(raw) + assert result == hash_token(raw) + assert result is not None + assert not result.startswith("sk-") + assert len(result) == 64 + + +def test_redact_logged_api_key_bearer_sk_equals_sk_hash(): + raw = "sk-1234secret" + result_plain = _redact_logged_api_key(raw) + result_bearer = _redact_logged_api_key(f"Bearer {raw}") + assert result_bearer == result_plain + + +def test_redact_logged_api_key_bearer_case_insensitive(): + raw = "sk-1234secret" + result_lower = _redact_logged_api_key(f"bearer {raw}") + result_upper = _redact_logged_api_key(f"BEARER {raw}") + expected = hash_token(raw) + assert result_lower == expected + assert result_upper == expected + + +def test_redact_logged_api_key_non_sk_raw_key_is_hashed(): + raw = "anthropic-raw-key-xyz" + result = _redact_logged_api_key(raw) + assert result is not None + assert result != raw + assert len(result) == 64 + assert result == hash_token(raw) + + +def test_redact_logged_api_key_already_valid_sha256_passes_through_with_flag(): + already_hashed = hash_token("sk-some-key") + assert len(already_hashed) == 64 + result = _redact_logged_api_key(already_hashed, already_redacted=True) + assert result == already_hashed + assert hash_token(already_hashed) != result # no double-hash + + +def test_redact_logged_api_key_sha256_without_flag_is_hashed(): + already_hashed = hash_token("sk-some-key") + assert len(already_hashed) == 64 + result = _redact_logged_api_key(already_hashed) + assert result is not None + assert result != already_hashed + assert len(result) == 64 + assert result == hash_token(already_hashed) + + +def test_redact_logged_api_key_long_opaque_token_is_hashed(): + raw = "x1" * 450 + assert len(raw) == 900 + result = _redact_logged_api_key(raw) + assert result is not None + assert result != raw + assert raw not in result + assert len(result) == 64 + assert result == hash_token(raw) + + +def test_redact_logged_api_key_hashed_jwt_passes_through(): + jwt_hash = "hashed-jwt-" + "a" * 64 + result = _redact_logged_api_key(jwt_hash, already_redacted=True) + assert result == jwt_hash + + +def test_redact_logged_api_key_hashed_jwt_shape_without_provenance_is_hashed(): + lookalike = "hashed-jwt-" + "a" * 64 + result = _redact_logged_api_key(lookalike) + assert result == hash_token(lookalike) + assert result != lookalike + + +def test_redact_logged_api_key_hashed_jwt_trailing_newline_is_hashed(): + trailing = "hashed-jwt-" + "a" * 64 + "\n" + result = _redact_logged_api_key(trailing, already_redacted=True) + assert result == hash_token(trailing) + assert result != trailing + + +def test_redact_logged_api_key_hashed_jwt_short_suffix_is_hashed(): + short_jwt = "hashed-jwt-tooshort" + result = _redact_logged_api_key(short_jwt) + assert result is not None + assert result != short_jwt + assert len(result) == 64 + assert result == hash_token(short_jwt) + + +def test_redact_logged_api_key_master_key_alias_passes_through(): + from litellm.constants import LITELLM_PROXY_MASTER_KEY_ALIAS + + result = _redact_logged_api_key(LITELLM_PROXY_MASTER_KEY_ALIAS, already_redacted=True) + assert result == LITELLM_PROXY_MASTER_KEY_ALIAS + + +def test_redact_logged_api_key_master_key_alias_without_provenance_is_hashed(): + from litellm.constants import LITELLM_PROXY_MASTER_KEY_ALIAS + + result = _redact_logged_api_key(LITELLM_PROXY_MASTER_KEY_ALIAS) + assert result == hash_token(LITELLM_PROXY_MASTER_KEY_ALIAS) + assert result != LITELLM_PROXY_MASTER_KEY_ALIAS + + +def test_get_spend_logs_metadata_keeps_master_key_alias_readable(): + from litellm.constants import LITELLM_PROXY_MASTER_KEY_ALIAS + + meta = _get_spend_logs_metadata( + { + "user_api_key": LITELLM_PROXY_MASTER_KEY_ALIAS, + "user_api_key_hash": LITELLM_PROXY_MASTER_KEY_ALIAS, + } + ) + assert meta["user_api_key"] == LITELLM_PROXY_MASTER_KEY_ALIAS + + +def test_redact_logged_api_key_bearer_only_returns_none(): + # "bearer " with nothing after stripping is equivalent to no key + assert _redact_logged_api_key("bearer ") is None + assert _redact_logged_api_key("Bearer ") is None + assert _redact_logged_api_key("BEARER ") is None + + +# ── _get_spend_logs_metadata key-hash invariant tests ───────────────────── + + +def test_get_spend_logs_metadata_sk_key_hashed(): + raw = "sk-1234secret" + meta = _get_spend_logs_metadata({"user_api_key": raw}) + assert meta["user_api_key"] == hash_token(raw) + assert meta["user_api_key"] is not None + result = meta["user_api_key"] + assert result is not None + assert not result.startswith("sk-") + assert len(result) == 64 + + +def test_get_spend_logs_metadata_bearer_sk_key_hashed_same_as_plain(): + raw = "sk-1234secret" + meta_plain = _get_spend_logs_metadata({"user_api_key": raw}) + meta_bearer = _get_spend_logs_metadata({"user_api_key": f"Bearer {raw}"}) + assert meta_bearer["user_api_key"] == meta_plain["user_api_key"] + + +def test_get_spend_logs_metadata_non_sk_raw_key_hashed(): + raw = "anthropic-raw-key-xyz" + meta = _get_spend_logs_metadata({"user_api_key": raw}) + result = meta["user_api_key"] + assert result is not None + assert result != raw + assert len(result) == 64 + + +def test_get_spend_logs_metadata_already_hashed_unchanged_with_provenance(): + already_hashed = hash_token("sk-some-key") + meta = _get_spend_logs_metadata( + {"user_api_key": already_hashed, "user_api_key_hash": already_hashed} + ) + assert meta["user_api_key"] == already_hashed + assert hash_token(already_hashed) != meta["user_api_key"] # no double-hash + + +def test_get_spend_logs_metadata_already_hashed_no_provenance_is_rehashed(): + already_hashed = hash_token("sk-some-key") + meta = _get_spend_logs_metadata({"user_api_key": already_hashed}) + assert meta["user_api_key"] != already_hashed + assert meta["user_api_key"] == hash_token(already_hashed) + + +def test_get_spend_logs_metadata_provenance_bypass_requires_hash_match(): + already_hashed = hash_token("sk-some-key") + different_hash = hash_token("sk-other-key") + meta = _get_spend_logs_metadata( + {"user_api_key": already_hashed, "user_api_key_hash": different_hash} + ) + assert meta["user_api_key"] == hash_token(already_hashed) + + +def test_get_spend_logs_metadata_hashed_jwt_unchanged(): + jwt_hash = "hashed-jwt-" + "b" * 64 + meta = _get_spend_logs_metadata({"user_api_key": jwt_hash, "user_api_key_hash": jwt_hash}) + assert meta["user_api_key"] == jwt_hash + + +def test_get_spend_logs_metadata_hashed_jwt_shape_without_provenance_is_hashed(): + lookalike = "hashed-jwt-" + "b" * 64 + meta = _get_spend_logs_metadata({"user_api_key": lookalike}) + assert meta["user_api_key"] == hash_token(lookalike) + assert meta["user_api_key"] != lookalike + + +def test_get_spend_logs_metadata_none_key_is_none(): + meta = _get_spend_logs_metadata({"user_api_key": None}) + assert meta["user_api_key"] is None + + +# ── get_logging_payload key-hash invariant tests ─────────────────────────── + + def test_get_logging_payload_uses_recovered_combined_usage_on_failure(): """A request that fails mid-stream has no usable response_obj usage, but the streaming handler recovers the usage from the chunks already delivered and @@ -2747,44 +2946,107 @@ def test_get_logging_payload_cache_hit_keeps_raw_litellm_call_id(): assert json.loads(payload["metadata"])["litellm_call_id"] != payload["request_id"] -class TestHashApiKeyForSpendLog: +class TestSpendLogKeyRedaction: """Regression: plaintext API keys with Bearer prefix were stored in SpendLogs for failed requests (LIT-4121)""" def test_bearer_prefixed_sk_key_is_hashed(self): raw = "Bearer sk-WLi4iRn4JmbVlTaYw12IOA" - result = _hash_api_key_for_spend_log(raw) + result = _redact_logged_api_key(raw) + assert result is not None assert not result.startswith("Bearer") assert not result.startswith("sk-") assert len(result) == 64 def test_bare_sk_key_is_hashed(self): raw = "sk-WLi4iRn4JmbVlTaYw12IOA" - result = _hash_api_key_for_spend_log(raw) + result = _redact_logged_api_key(raw) + assert result is not None assert not result.startswith("sk-") assert len(result) == 64 def test_bearer_lowercase_is_handled(self): raw = "bearer sk-WLi4iRn4JmbVlTaYw12IOA" - result = _hash_api_key_for_spend_log(raw) + result = _redact_logged_api_key(raw) + assert result is not None assert not result.startswith("bearer") assert not result.startswith("sk-") assert len(result) == 64 def test_already_hashed_key_unchanged(self): hashed = "bcfe8173f5447f10be0e7fb37aaa8b97829d5c9e0498232152f9d123456789ab" - assert _hash_api_key_for_spend_log(hashed) == hashed + assert _redact_logged_api_key(hashed, already_redacted=True) == hashed - def test_bearer_prefixed_non_sk_key_strips_prefix(self): + def test_bearer_prefixed_non_sk_key_is_hashed(self): raw = "Bearer some-other-token-format" - result = _hash_api_key_for_spend_log(raw) - assert result == "some-other-token-format" + result = _redact_logged_api_key(raw) + assert result == hash_token("some-other-token-format") + assert result is not None assert not result.startswith("Bearer") def test_bearer_and_bare_produce_same_hash(self): bare = "sk-WLi4iRn4JmbVlTaYw12IOA" bearer = "Bearer sk-WLi4iRn4JmbVlTaYw12IOA" - assert _hash_api_key_for_spend_log(bare) == _hash_api_key_for_spend_log(bearer) + assert _redact_logged_api_key(bare) == _redact_logged_api_key(bearer) + + +@patch("litellm.proxy.proxy_server.master_key", None) +@patch("litellm.proxy.proxy_server.general_settings", {}) +def test_get_logging_payload_non_sk_raw_key_both_fields_hashed(): + raw = "anthropic-raw-key-xyz" + kwargs = { + "model": "openai/gpt-4.1", + "messages": [{"role": "user", "content": "Hello"}], + "call_type": "acompletion", + "litellm_params": { + "metadata": { + "user_api_key": raw, + "user_api_key_user_id": "test_user", + "user_api_key_team_id": "test_team", + } + }, + } + payload = get_logging_payload( + kwargs=kwargs, + response_obj=Exception("error"), + start_time=datetime.datetime.now(timezone.utc), + end_time=datetime.datetime.now(timezone.utc), + ) + + assert payload["api_key"] != raw + assert len(payload["api_key"]) == 64 + + parsed_meta = json.loads(payload["metadata"]) + assert parsed_meta["user_api_key"] != raw + assert parsed_meta["user_api_key"] is not None + assert len(parsed_meta["user_api_key"]) == 64 + + +def test_get_logging_payload_keeps_master_key_alias_readable(): + from litellm.constants import LITELLM_PROXY_MASTER_KEY_ALIAS + + kwargs = { + "model": "openai/gpt-4.1", + "messages": [{"role": "user", "content": "Hello"}], + "call_type": "acompletion", + "litellm_params": { + "metadata": { + "user_api_key": LITELLM_PROXY_MASTER_KEY_ALIAS, + "user_api_key_hash": LITELLM_PROXY_MASTER_KEY_ALIAS, + "user_api_key_user_id": "test_user", + } + }, + } + payload = get_logging_payload( + kwargs=kwargs, + response_obj=Exception("error"), + start_time=datetime.datetime.now(timezone.utc), + end_time=datetime.datetime.now(timezone.utc), + ) + + assert payload["api_key"] == LITELLM_PROXY_MASTER_KEY_ALIAS + parsed_meta = json.loads(payload["metadata"]) + assert parsed_meta["user_api_key"] == LITELLM_PROXY_MASTER_KEY_ALIAS @patch("litellm.proxy.proxy_server.master_key", None) @@ -3297,3 +3559,69 @@ def test_get_logging_payload_failed_request_without_standard_logging_payload_lea assert payload["model_group"] == "" assert payload["api_base"] == "" assert payload["custom_llm_provider"] == "" + + +@patch("litellm.proxy.proxy_server.master_key", None) +@patch("litellm.proxy.proxy_server.general_settings", {}) +def test_get_logging_payload_empty_key_slp_none_is_empty_string_not_none_literal(): + kwargs = { + "model": "openai/gpt-4.1", + "messages": [{"role": "user", "content": "Hello"}], + "call_type": "acompletion", + "litellm_params": { + "metadata": { + "user_api_key_user_id": "test_user", + } + }, + } + payload = get_logging_payload( + kwargs=kwargs, + response_obj=Exception("error"), + start_time=datetime.datetime.now(timezone.utc), + end_time=datetime.datetime.now(timezone.utc), + ) + + assert payload["api_key"] == "", ( + f"Expected empty string but got {payload['api_key']!r}; " + "dropping _redact_logged_api_key's 'or \"\"' guard would yield 'None' here" + ) + + +def test_get_spend_logs_metadata_sibling_fields_preserved(): + raw = "anthropic-raw-key-xyz" + meta = _get_spend_logs_metadata( + { + "user_api_key": raw, + "user_api_key_alias": "my-alias", + "user_api_key_team_id": "team-123", + } + ) + assert meta["user_api_key"] == hash_token(raw) + assert meta["user_api_key_alias"] == "my-alias" + assert meta["user_api_key_team_id"] == "team-123" + + +def test_redact_logged_api_key_partial_sha256_is_hashed(): + partial_hex = "a" * 63 + result = _redact_logged_api_key(partial_hex) + assert result is not None + assert result != partial_hex + assert len(result) == 64 + assert result == hash_token(partial_hex) + + +def test_redact_logged_api_key_bearer_already_hashed_passes_through_with_flag(): + already_hashed = hash_token("sk-some-key") + assert len(already_hashed) == 64 + result = _redact_logged_api_key(f"Bearer {already_hashed}", already_redacted=True) + assert result == already_hashed + assert hash_token(already_hashed) != result + + +def test_redact_logged_api_key_bearer_sha256_without_flag_is_hashed(): + already_hashed = hash_token("sk-some-key") + assert len(already_hashed) == 64 + result = _redact_logged_api_key(f"Bearer {already_hashed}") + assert result is not None + assert result != already_hashed + assert result == hash_token(already_hashed) diff --git a/tests/test_litellm/proxy/test_budget_reservation.py b/tests/test_litellm/proxy/test_budget_reservation.py index 133b53bb18d..2388654bf4b 100644 --- a/tests/test_litellm/proxy/test_budget_reservation.py +++ b/tests/test_litellm/proxy/test_budget_reservation.py @@ -2412,10 +2412,13 @@ async def test_streaming_cancel_before_any_chunk_reconciles_to_input_cost( generator, streaming_logging_obj = _drive_streaming_cancel(valid_token, cancel_before_chunk) received = [] - with pytest.raises(asyncio.CancelledError): + async def _drain(): async for chunk in generator: received.append(chunk) + with pytest.raises(asyncio.CancelledError): + await _drain() + assert received == [] # no chunk delivered, but the provider already received the input, so the # reservation is reconciled to the input cost (0.5), not refunded to zero @@ -2444,10 +2447,13 @@ async def test_streaming_cancel_after_chunk_keeps_reservation( generator, streaming_logging_obj = _drive_streaming_cancel(valid_token, cancel_after_chunk) received = [] - with pytest.raises(asyncio.CancelledError): + async def _drain(): async for chunk in generator: received.append(chunk) + with pytest.raises(asyncio.CancelledError): + await _drain() + assert received == ["data: chunk\n\n"] # a consumed stream must NOT be refunded assert counter_cache.in_memory_cache.get_cache( @@ -2508,10 +2514,13 @@ async def test_streaming_cancel_in_slow_path_before_yield_refunds(spend_counter_ received = [] # include_cost_in_streaming_usage forces fast_path off, so the hook above runs with patch.object(litellm, "include_cost_in_streaming_usage", True, create=True): - with pytest.raises(asyncio.CancelledError): + async def _drain(): async for chunk in generator: received.append(chunk) + with pytest.raises(asyncio.CancelledError): + await _drain() + assert received == [] # cancellation happened before any chunk reached the client, but the # provider already received the input -> reconcile to the input cost (0.5) diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index 510fb977a61..716fba370df 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -435,7 +435,7 @@ class TestProxyBaseLLMRequestProcessing: # Test with invalid header value (should raise ValueError when converting to float) headers_with_invalid = {"x-litellm-stream-timeout": "invalid"} - with pytest.raises(ValueError): + with pytest.raises(ValueError, match="could not convert string to float: 'invalid"): LiteLLMProxyRequestSetup._get_stream_timeout_from_request(headers_with_invalid) @pytest.mark.asyncio @@ -5519,6 +5519,196 @@ class TestStreamingClientDisconnectBilling: proxy_logging_obj._arelease_max_parallel_requests_on_disconnect.assert_awaited_once() + async def _bill_and_collect_success_event(self, prepare=None, request_data=None): + recorder = _RecordingSuccessLogger() + original_callbacks = litellm.callbacks + litellm.callbacks = [recorder] + try: + response = await self._start_partial_stream() + if prepare is not None: + prepare(response) + billed = await _bill_partial_streamed_spend_on_disconnect( + {"litellm_logging_obj": response.logging_obj, **(request_data or {})}, response + ) + assert billed is True + for _ in range(50): + if recorder.success_events: + break + await asyncio.sleep(0.1) + finally: + litellm.callbacks = original_callbacks + assert len(recorder.success_events) == 1 + return recorder.success_events[0] + + @pytest.mark.asyncio + async def test_disconnect_billing_prices_alias_restamped_chunks_at_real_model(self): + assert "openai/my-public-alias" not in litellm.model_cost + + def restamp_chunks_to_alias(response): + for chunk in response.chunks: + chunk.model = "my-public-alias" + + event = await self._bill_and_collect_success_event(restamp_chunks_to_alias) + + assert event["response_obj"].model == "gpt-4o-mini" + standard_logging_object = event["kwargs"]["standard_logging_object"] + assert standard_logging_object["response_cost"] > 0.0 + + @pytest.mark.asyncio + async def test_disconnect_billing_prices_a_partly_restamped_chunk_list_at_real_model(self): + """ + A chunk that carries usage is stored as a copy before the proxy restamps the + one it forwards, so an aliased stream can reach billing with its first chunk + still on the deployment model and the rest on the client's name. + """ + assert "openai/my-public-alias" not in litellm.model_cost + + def restamp_only_the_chunks_the_proxy_forwarded(response): + for chunk in response.chunks[1:]: + chunk.model = "my-public-alias" + + event = await self._bill_and_collect_success_event( + restamp_only_the_chunks_the_proxy_forwarded, + request_data={"model": "my-public-alias"}, + ) + + assert event["response_obj"].model == "gpt-4o-mini" + standard_logging_object = event["kwargs"]["standard_logging_object"] + assert standard_logging_object["response_cost"] > 0.0 + + @pytest.mark.asyncio + async def test_disconnect_billing_keeps_the_model_azure_model_router_picked(self): + def restamp_like_azure_model_router(response): + response.chunks[0].model = "azure-model-router" + for chunk in response.chunks[1:]: + chunk.model = "gpt-4.1-nano-2025-04-14" + + event = await self._bill_and_collect_success_event( + restamp_like_azure_model_router, + request_data={"model": "azure-model-router"}, + ) + + assert event["response_obj"].model == "gpt-4.1-nano-2025-04-14" + standard_logging_object = event["kwargs"]["standard_logging_object"] + assert standard_logging_object["response_cost"] > 0.0 + + @pytest.mark.asyncio + async def test_disconnect_billing_keeps_the_routed_model_when_request_data_model_was_rewritten(self): + """ + Pre-call processing rewrites request_data["model"] for aliasing and routing, so the + routed model on the later chunks can end up matching it. Only the name the client + sent says whether the proxy restamped this stream. + """ + + def restamp_like_azure_model_router(response): + response.chunks[0].model = "azure-model-router" + for chunk in response.chunks[1:]: + chunk.model = "gpt-4.1-nano-2025-04-14" + + event = await self._bill_and_collect_success_event( + restamp_like_azure_model_router, + request_data={ + "model": "gpt-4.1-nano-2025-04-14", + "_litellm_client_requested_model": "azure-model-router", + }, + ) + + assert event["response_obj"].model == "gpt-4.1-nano-2025-04-14" + standard_logging_object = event["kwargs"]["standard_logging_object"] + assert standard_logging_object["response_cost"] > 0.0 + + @pytest.mark.asyncio + async def test_disconnect_billing_backfills_missing_cache_fields(self): + event = await self._bill_and_collect_success_event() + + usage = event["response_obj"].usage + assert getattr(usage, "cache_creation_input_tokens", None) == 0 + assert getattr(usage, "cache_read_input_tokens", None) == 0 + assert usage.prompt_tokens_details is not None + assert usage.prompt_tokens_details.cached_tokens == 0 + + @pytest.mark.asyncio + async def test_disconnect_billing_carries_up_openai_style_cached_tokens(self): + from litellm.types.utils import ( + Delta, + ModelResponseStream, + PromptTokensDetailsWrapper, + StreamingChoices, + Usage, + ) + + def append_openai_style_cached_usage_chunk(response): + response.chunks.append( + ModelResponseStream( + id=response.chunks[0].id, + model="gpt-4o-mini", + object="chat.completion.chunk", + choices=[ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta(content=" and more", role="assistant"), + ) + ], + usage=Usage( + prompt_tokens=1000, + completion_tokens=10, + total_tokens=1010, + prompt_tokens_details=PromptTokensDetailsWrapper( + cached_tokens=500 + ), + ), + ) + ) + + event = await self._bill_and_collect_success_event( + append_openai_style_cached_usage_chunk + ) + + usage = event["response_obj"].usage + assert getattr(usage, "cache_read_input_tokens", None) == 500 + assert getattr(usage, "cache_creation_input_tokens", None) == 0 + + @pytest.mark.asyncio + async def test_disconnect_billing_keeps_cache_values_recovered_from_chunks(self): + from litellm.types.utils import ( + Delta, + ModelResponseStream, + StreamingChoices, + Usage, + ) + + def append_usage_chunk(response): + response.chunks.append( + ModelResponseStream( + id=response.chunks[0].id, + model="gpt-4o-mini", + object="chat.completion.chunk", + choices=[ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta(content=" and more", role="assistant"), + ) + ], + usage=Usage( + prompt_tokens=40, + completion_tokens=5, + total_tokens=45, + cache_read_input_tokens=7, + cache_creation_input_tokens=3, + ), + ) + ) + + event = await self._bill_and_collect_success_event(append_usage_chunk) + + usage = event["response_obj"].usage + assert getattr(usage, "cache_read_input_tokens", None) == 7 + assert getattr(usage, "cache_creation_input_tokens", None) == 3 + assert usage.prompt_tokens_details is not None + assert usage.prompt_tokens_details.cached_tokens == 7 + def _apply_stream_usage_tracking( data: dict, @@ -6082,6 +6272,254 @@ class TestInjectCostIntoUsageDict: injected = json.loads(result.split("\n")[0].split("data:", 1)[1].strip()) assert injected["usage"]["cost"] == pytest.approx(self._expected_cost("gpt-4o-mini", 11, 4)) + def test_message_delta_cost_charges_the_non_cached_input_tokens(self): + """Anthropic reports ``input_tokens`` excluding cache tokens, so reading it as the whole + prompt total drops the non-cached input from the bill on every cache hit.""" + model = "claude-haiku-4-5" + pricing = litellm.model_cost[model] + event = { + "type": "message_delta", + "delta": {"stop_reason": "end_turn"}, + "usage": { + "input_tokens": 14, + "output_tokens": 8, + "cache_read_input_tokens": 3202, + "cache_creation_input_tokens": 0, + }, + } + + result = ProxyBaseLLMRequestProcessing._inject_cost_into_usage_dict(event, model) + + assert result is not None + expected = ( + 14 * pricing["input_cost_per_token"] + + 3202 * pricing["cache_read_input_token_cost"] + + 8 * pricing["output_cost_per_token"] + ) + dropped_input = expected - 14 * pricing["input_cost_per_token"] + assert result["usage"]["cost"] == pytest.approx(expected) + assert result["usage"]["cost"] > dropped_input + + def test_message_delta_prices_1h_cache_creation_above_the_5m_rate(self): + """The ``cache_creation`` 5m/1h split has to survive into ``prompt_tokens_details``, + otherwise a 1h write is billed at the cheaper 5m rate.""" + model = "claude-haiku-4-5" + pricing = litellm.model_cost[model] + event = { + "type": "message_delta", + "delta": {"stop_reason": "end_turn"}, + "usage": { + "input_tokens": 14, + "output_tokens": 8, + "cache_read_input_tokens": 0, + "cache_creation_input_tokens": 2000, + "cache_creation": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 2000}, + }, + } + + result = ProxyBaseLLMRequestProcessing._inject_cost_into_usage_dict(event, model) + + assert result is not None + base = 14 * pricing["input_cost_per_token"] + 8 * pricing["output_cost_per_token"] + expected_1h = base + 2000 * pricing["cache_creation_input_token_cost_above_1hr"] + flat_5m = base + 2000 * pricing["cache_creation_input_token_cost"] + assert expected_1h != pytest.approx(flat_5m) + assert result["usage"]["cost"] == pytest.approx(expected_1h) + + def test_message_delta_prices_through_the_logging_obj_so_custom_pricing_applies(self): + """Costing by model name alone yields sticker price, so a deployment with a negotiated + discount streamed a ``usage.cost`` that disagreed with the callback's ``response_cost``.""" + + class _StubLoggingObj: + def __init__(self, cost): + self._cost = cost + self.captured_result = None + + def _response_cost_calculator(self, result): + self.captured_result = result + return self._cost + + model = "claude-haiku-4-5" + discounted_cost = 0.00099 + stub = _StubLoggingObj(discounted_cost) + event = { + "type": "message_delta", + "delta": {"stop_reason": "end_turn"}, + "usage": { + "input_tokens": 14, + "output_tokens": 8, + "cache_read_input_tokens": 3202, + "cache_creation_input_tokens": 500, + "cache_creation": {"ephemeral_5m_input_tokens": 100, "ephemeral_1h_input_tokens": 400}, + }, + } + + result = ProxyBaseLLMRequestProcessing._inject_cost_into_usage_dict(event, model, stub) + + assert result is not None + assert result["usage"]["cost"] == discounted_cost + assert result["usage"]["cost"] != pytest.approx(self._expected_cost(model, 14 + 500 + 3202, 8)) + usage = stub.captured_result.usage + assert usage.prompt_tokens == 14 + 500 + 3202 + details = usage.prompt_tokens_details.cache_creation_token_details + assert details.ephemeral_5m_input_tokens == 100 + assert details.ephemeral_1h_input_tokens == 400 + + def test_message_delta_falls_back_to_model_pricing_when_the_logging_obj_returns_no_cost(self): + class _StubLoggingObj: + def _response_cost_calculator(self, result): + return None + + model = "claude-haiku-4-5" + pricing = litellm.model_cost[model] + event = { + "type": "message_delta", + "delta": {"stop_reason": "end_turn"}, + "usage": {"input_tokens": 14, "output_tokens": 8, "cache_read_input_tokens": 3202}, + } + + result = ProxyBaseLLMRequestProcessing._inject_cost_into_usage_dict(event, model, _StubLoggingObj()) + + assert result is not None + assert result["usage"]["cost"] == pytest.approx( + 14 * pricing["input_cost_per_token"] + + 3202 * pricing["cache_read_input_token_cost"] + + 8 * pricing["output_cost_per_token"] + ) + + def test_message_delta_falls_back_to_model_pricing_when_the_logging_obj_raises(self): + """A pricing failure mid-stream must not break the frame, so the raise falls back to + model-name pricing rather than propagating into the response body.""" + + class _StubLoggingObj: + def _response_cost_calculator(self, result): + raise ValueError("no pricing for this deployment") + + model = "claude-haiku-4-5" + pricing = litellm.model_cost[model] + event = { + "type": "message_delta", + "delta": {"stop_reason": "end_turn"}, + "usage": {"input_tokens": 14, "output_tokens": 8, "cache_read_input_tokens": 3202}, + } + + result = ProxyBaseLLMRequestProcessing._inject_cost_into_usage_dict(event, model, _StubLoggingObj()) + + assert result is not None + assert result["usage"]["cost"] == pytest.approx( + 14 * pricing["input_cost_per_token"] + + 3202 * pricing["cache_read_input_token_cost"] + + 8 * pricing["output_cost_per_token"] + ) + + def test_pricing_a_frame_leaves_the_real_logging_obj_unchanged(self): + """Pricing runs against the live logging object, and the pass-through handlers never + recompute cost_breakdown, so a frame-derived breakdown would reach the spend log.""" + from litellm.litellm_core_utils.litellm_logging import ( + Logging as LiteLLMLoggingObj, + ) + from litellm.types.utils import ModelResponse, Usage + + logging_obj = LiteLLMLoggingObj( + model="claude-haiku-4-5", + messages=[{"role": "user", "content": "test"}], + stream=True, + call_type="completion", + start_time=None, + litellm_call_id="lit4902-breakdown-test", + function_id="lit4902-breakdown-test", + ) + logging_obj.update_environment_variables(litellm_params={}, optional_params={}) + logging_obj.model_call_details["custom_llm_provider"] = "anthropic" + assert logging_obj.cost_breakdown is None + + model_response = ModelResponse( + usage=Usage(prompt_tokens=3216, completion_tokens=8, total_tokens=3224) + ) + cost = ProxyBaseLLMRequestProcessing._logging_obj_cost_or_none(model_response, logging_obj) + + assert cost is not None and cost > 0 + assert logging_obj.cost_breakdown is None + assert "response_cost_failure_debug_information" not in logging_obj.model_call_details + + def test_pricing_a_frame_restores_a_breakdown_the_request_already_had(self): + from litellm.litellm_core_utils.litellm_logging import ( + Logging as LiteLLMLoggingObj, + ) + from litellm.types.utils import ModelResponse, Usage + + logging_obj = LiteLLMLoggingObj( + model="claude-haiku-4-5", + messages=[{"role": "user", "content": "test"}], + stream=True, + call_type="completion", + start_time=None, + litellm_call_id="lit4902-breakdown-restore", + function_id="lit4902-breakdown-restore", + ) + logging_obj.update_environment_variables(litellm_params={}, optional_params={}) + logging_obj.model_call_details["custom_llm_provider"] = "anthropic" + logging_obj.set_cost_breakdown( + input_cost=0.5, output_cost=0.25, total_cost=0.75, cost_for_built_in_tools_cost_usd_dollar=0.0 + ) + existing = logging_obj.cost_breakdown + + model_response = ModelResponse( + usage=Usage(prompt_tokens=3216, completion_tokens=8, total_tokens=3224) + ) + ProxyBaseLLMRequestProcessing._logging_obj_cost_or_none(model_response, logging_obj) + + assert logging_obj.cost_breakdown is existing + assert logging_obj.cost_breakdown["total_cost"] == 0.75 + + def test_openai_chunk_prices_through_the_logging_obj_so_custom_pricing_applies(self): + """The chat.completion.chunk path rides the same pricer, so a discounted deployment + streaming /v1/chat/completions gets its negotiated price instead of sticker.""" + + class _StubLoggingObj: + def __init__(self, cost): + self._cost = cost + self.captured_result = None + + def _response_cost_calculator(self, result): + self.captured_result = result + return self._cost + + discounted_cost = 0.00031 + stub = _StubLoggingObj(discounted_cost) + event = { + "id": "chatcmpl-1", + "object": "chat.completion.chunk", + "choices": [], + "usage": {"prompt_tokens": 1000, "completion_tokens": 100, "total_tokens": 1100}, + } + + result = ProxyBaseLLMRequestProcessing._inject_cost_into_usage_dict(event, "gpt-4o-mini", stub) + + assert result is not None + assert result["usage"]["cost"] == discounted_cost + assert result["usage"]["cost"] != pytest.approx(self._expected_cost("gpt-4o-mini", 1000, 100)) + usage = stub.captured_result.usage + assert usage.prompt_tokens == 1000 + assert usage.completion_tokens == 100 + + def test_openai_chunk_falls_back_to_model_pricing_when_the_logging_obj_returns_no_cost(self): + class _StubLoggingObj: + def _response_cost_calculator(self, result): + return None + + event = { + "id": "chatcmpl-1", + "object": "chat.completion.chunk", + "choices": [], + "usage": {"prompt_tokens": 11, "completion_tokens": 4, "total_tokens": 15}, + } + + result = ProxyBaseLLMRequestProcessing._inject_cost_into_usage_dict(event, "gpt-4o-mini", _StubLoggingObj()) + + assert result is not None + assert result["usage"]["cost"] == pytest.approx(self._expected_cost("gpt-4o-mini", 11, 4)) + class TestProcessChunkWithCostInjection: def test_complete_usage_frame_chunk_is_injected(self, monkeypatch): @@ -6116,6 +6554,31 @@ class TestProcessChunkWithCostInjection: assert ProxyBaseLLMRequestProcessing._process_chunk_with_cost_injection(chunk, "gpt-4o-mini") == chunk + def test_message_delta_frame_is_priced_with_the_logging_obj(self, monkeypatch): + """Pins that the logging object reaches the pricer through the byte-frame entry point, + which is how the proxy actually calls this on a streamed Messages API request.""" + monkeypatch.setattr(litellm, "include_cost_in_streaming_usage", True) + + class _StubLoggingObj: + def _response_cost_calculator(self, result): + return 0.00042 + + chunk = ( + b"event: message_delta\n" + b'data: {"type":"message_delta","delta":{"stop_reason":"end_turn"},' + b'"usage":{"input_tokens":14,"output_tokens":8,"cache_read_input_tokens":3202}}\n\n' + ) + + result = ProxyBaseLLMRequestProcessing._process_chunk_with_cost_injection( + chunk, "claude-haiku-4-5", _StubLoggingObj() + ) + + assert result != chunk + data_line = next(ln for ln in result.decode("utf-8").splitlines() if ln.startswith("data:")) + payload = json.loads(data_line.split("data:", 1)[1].strip()) + assert payload["usage"]["cost"] == 0.00042 + assert payload["usage"]["cache_read_input_tokens"] == 3202 + # --------------------------------------------------------------------------- # SSE keepalive during the time-to-first-token (issue #34819) diff --git a/tests/test_litellm/proxy/test_enforce_user_param.py b/tests/test_litellm/proxy/test_enforce_user_param.py index 6891123e70e..1001372aeb5 100644 --- a/tests/test_litellm/proxy/test_enforce_user_param.py +++ b/tests/test_litellm/proxy/test_enforce_user_param.py @@ -56,7 +56,7 @@ class TestEnforceUserParamPostGetFiltering: new_callable=AsyncMock, return_value=True, ): - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match="user' param not passed in\\. 'enforce_user_param'=True") as exc_info: await common_checks( request_body=request_body, team_object=None, @@ -175,7 +175,7 @@ class TestEnforceUserParamPostGetFiltering: new_callable=AsyncMock, return_value=True, ): - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match="user' param not passed in\\. 'enforce_user_param'=True") as exc_info: await common_checks( request_body=request_body, team_object=None, @@ -405,7 +405,7 @@ class TestEnforceUserParamEdgeCases: new_callable=AsyncMock, return_value=True, ): - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match="user' param not passed in\\. 'enforce_user_param'=True") as exc_info: await common_checks( request_body=request_body, team_object=None, diff --git a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py index b1071150f3b..636974d5deb 100644 --- a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py +++ b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py @@ -2279,12 +2279,12 @@ def test_get_num_retries_from_request(): # Test case 7: Header present with invalid value (should raise ValueError when int() is called) headers_with_invalid = {"x-litellm-num-retries": "invalid"} - with pytest.raises(ValueError): + with pytest.raises(ValueError, match='invalid literal for int\\(\\) with base'): LiteLLMProxyRequestSetup._get_num_retries_from_request(headers_with_invalid) # Test case 8: Header present with float string (should raise ValueError when int() is called) headers_with_float = {"x-litellm-num-retries": "3.5"} - with pytest.raises(ValueError): + with pytest.raises(ValueError, match='invalid literal for int\\(\\) with base'): LiteLLMProxyRequestSetup._get_num_retries_from_request(headers_with_float) # Test case 9: Header present with negative number @@ -2324,7 +2324,7 @@ def test_get_keepalive_seconds_from_request(): # Header present with invalid value raises ValueError, matching the other # x-litellm-* numeric header helpers (_get_timeout_from_request, etc.) - with pytest.raises(ValueError): + with pytest.raises(ValueError, match="could not convert string to float: 'not-a-number"): LiteLLMProxyRequestSetup._get_keepalive_seconds_from_request( {"x-litellm-keepalive-seconds": "not-a-number"} ) diff --git a/tests/test_litellm/proxy/test_proxy_logging_hook_detection.py b/tests/test_litellm/proxy/test_proxy_logging_hook_detection.py index 133156f9321..542572e1e56 100644 --- a/tests/test_litellm/proxy/test_proxy_logging_hook_detection.py +++ b/tests/test_litellm/proxy/test_proxy_logging_hook_detection.py @@ -291,7 +291,7 @@ async def test_post_call_stream_guardrail_blocks_anthropic_messages_stream(monke yield chunk delivered = [] - with pytest.raises(HTTPException) as exc_info: + async def _drain(): async for chunk in proxy_logging.async_post_call_streaming_iterator_hook( response=fake_stream(), user_api_key_dict=UserAPIKeyAuth(api_key="sk-1234", request_route="/v1/messages"), @@ -299,6 +299,9 @@ async def test_post_call_stream_guardrail_blocks_anthropic_messages_stream(monke ): delivered.append(chunk) + with pytest.raises(HTTPException) as exc_info: + await _drain() + detail = exc_info.value.detail assert detail["guardrail_name"] == "output-filter" assert detail["keyword"] == "zebra" @@ -411,7 +414,7 @@ async def test_post_call_stream_guardrail_reroutes_inherited_apply_guardrail(mon yield chunk delivered = [] - with pytest.raises(HTTPException) as exc_info: + async def _drain(): async for chunk in proxy_logging.async_post_call_streaming_iterator_hook( response=fake_stream(), user_api_key_dict=UserAPIKeyAuth(api_key="sk-1234", request_route="/v1/messages"), @@ -419,6 +422,9 @@ async def test_post_call_stream_guardrail_reroutes_inherited_apply_guardrail(mon ): delivered.append(chunk) + with pytest.raises(HTTPException) as exc_info: + await _drain() + assert exc_info.value.detail["keyword"] == "zebra" assert delivered == [] diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 75aa716bb85..83e9095c8ec 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -1507,7 +1507,7 @@ def test_team_info_masking(): "langfuse_public_key": "public-test-key", } - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match="secr\\*\\*\\*\\*\\*\\*\\*-key', 'langfuse_public_key':") as exc_info: proxy_config._get_team_config( team_id="test_dev", all_teams_config=[team1_info], diff --git a/tests/test_litellm/proxy/test_proxy_utils.py b/tests/test_litellm/proxy/test_proxy_utils.py index 07877514b69..fe79ef25da6 100644 --- a/tests/test_litellm/proxy/test_proxy_utils.py +++ b/tests/test_litellm/proxy/test_proxy_utils.py @@ -1584,7 +1584,7 @@ async def test_prisma_health_check_failure_redacts_database_credentials(caplog): client._report_health_check_failure = AsyncMock() with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): - with pytest.raises(Exception): + with pytest.raises(Exception, match="could not connect to"): await PrismaClient.health_check(client) emitted = [record.getMessage() for record in caplog.records if record.name == "LiteLLM Proxy"] diff --git a/tests/test_litellm/proxy/test_route_llm_request.py b/tests/test_litellm/proxy/test_route_llm_request.py index 1e716f7c148..23e0bbfb3ee 100644 --- a/tests/test_litellm/proxy/test_route_llm_request.py +++ b/tests/test_litellm/proxy/test_route_llm_request.py @@ -169,7 +169,7 @@ async def test_route_request_proxy_admin_can_call_all_team_scoped_deployments_wi ) ) - with pytest.raises(litellm.BadRequestError, match="multiple teams"): + async def _route_and_await(): ambiguous_call = await route_request( data=data, llm_router=router, @@ -179,6 +179,9 @@ async def test_route_request_proxy_admin_can_call_all_team_scoped_deployments_wi ) await ambiguous_call + with pytest.raises(litellm.BadRequestError, match="multiple teams"): + await _route_and_await() + router.add_deployment( Deployment( model_name="team-azure", diff --git a/tests/test_litellm/proxy/test_spend_log_cleanup.py b/tests/test_litellm/proxy/test_spend_log_cleanup.py index ce0b6b755cc..bf1538183ab 100644 --- a/tests/test_litellm/proxy/test_spend_log_cleanup.py +++ b/tests/test_litellm/proxy/test_spend_log_cleanup.py @@ -82,10 +82,10 @@ def test_spend_log_cleanup_cron_scheduling(): assert trigger_weekly is not None # Invalid cron expression should raise ValueError - with pytest.raises(ValueError): + with pytest.raises(ValueError, match='Wrong number of fields; got'): CronTrigger.from_crontab("invalid cron") - with pytest.raises(ValueError): + with pytest.raises(ValueError, match='is higher than the maximum value'): CronTrigger.from_crontab("60 25 * * *") # Invalid minute and hour diff --git a/tests/test_litellm/proxy/test_team_org_move.py b/tests/test_litellm/proxy/test_team_org_move.py index 2dc961bec85..064e9de550e 100644 --- a/tests/test_litellm/proxy/test_team_org_move.py +++ b/tests/test_litellm/proxy/test_team_org_move.py @@ -97,7 +97,7 @@ class TestValidateTeamOrgChange: team = _make_team(member_ids=["sso-user-001"]) org = _make_org(members=[]) - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Cannot move team to organization\\. Team has user_id') as exc_info: validate_team_org_change( team=team, organization=org, llm_router=router, is_proxy_admin=False ) diff --git a/tests/test_litellm/proxy/utils/helpers/test_team_configs.py b/tests/test_litellm/proxy/utils/helpers/test_team_configs.py index 0e0906892b0..185d4d26ff4 100644 --- a/tests/test_litellm/proxy/utils/helpers/test_team_configs.py +++ b/tests/test_litellm/proxy/utils/helpers/test_team_configs.py @@ -66,7 +66,7 @@ def test_is_valid_team_configs_short_circuits_when_team_id_none(): def test_is_valid_team_configs_raises_on_model_not_in_team_models(): team_config = {"models": ["gpt-4o"]} request_data = {"model": "claude-haiku"} - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='claude-haiku\\. Valid models for team are') as exc_info: _is_valid_team_configs( team_id="team-1", team_config=team_config, diff --git a/tests/test_litellm/proxy/utils/prisma_and_spend/test_proxy_update_spend.py b/tests/test_litellm/proxy/utils/prisma_and_spend/test_proxy_update_spend.py index 7057a112c83..93c99c7fd04 100644 --- a/tests/test_litellm/proxy/utils/prisma_and_spend/test_proxy_update_spend.py +++ b/tests/test_litellm/proxy/utils/prisma_and_spend/test_proxy_update_spend.py @@ -561,7 +561,7 @@ async def test_update_spend_logs_does_not_requeue_non_transport_failures( proxy_logging.failure_handler = AsyncMock() mock_prisma_client.spend_log_transactions = [] - with pytest.raises(ValueError): + with pytest.raises(ValueError, match="bad payload"): await ProxyUpdateSpend.update_spend_logs( n_retry_times=1, prisma_client=mock_prisma_client, diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_callback_capabilities_class.py b/tests/test_litellm/proxy/utils/proxy_logging/test_callback_capabilities_class.py index 9452e8042bd..75a91177f00 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_callback_capabilities_class.py +++ b/tests/test_litellm/proxy/utils/proxy_logging/test_callback_capabilities_class.py @@ -181,7 +181,7 @@ def test_has_streaming_callbacks_error_when_resolution_fails(monkeypatch): "get_custom_logger_compatible_class", lambda *a, **kw: (_ for _ in ()).throw(ValueError("nope")), ) - with pytest.raises(ValueError): + with pytest.raises(ValueError, match="nope"): ProxyLogging.has_streaming_callbacks() diff --git a/tests/test_litellm/repositories/test_unit_of_work.py b/tests/test_litellm/repositories/test_unit_of_work.py index c270a570ad9..1ebfd917e36 100644 --- a/tests/test_litellm/repositories/test_unit_of_work.py +++ b/tests/test_litellm/repositories/test_unit_of_work.py @@ -59,11 +59,14 @@ async def test_updates_across_tables_share_one_batch_and_commit_once(): async def test_raising_inside_block_skips_commit(): batch = FakeBatch() - with pytest.raises(RuntimeError, match="boom"): + async def _blow_up_mid_transaction(): async with spend_reset_unit_of_work(lambda: batch) as uow: uow.keys.queue_spend_reset(token="tok-1", budget_reset_at=None) raise RuntimeError("boom") + with pytest.raises(RuntimeError, match="boom"): + await _blow_up_mid_transaction() + assert batch.commit_count == 0 @@ -119,9 +122,12 @@ async def test_budget_cascade_raising_inside_block_skips_commit(): the tier is still due on the next tick.""" batch = FakeBatch() - with pytest.raises(RuntimeError, match="boom"): + async def _blow_up_mid_transaction(): async with budget_cascade_unit_of_work(lambda: batch) as uow: uow.team_memberships.queue_spend_zero(where={"budget_id": {"in": ["budget-1"]}}) raise RuntimeError("boom") + with pytest.raises(RuntimeError, match="boom"): + await _blow_up_mid_transaction() + assert batch.commit_count == 0 diff --git a/tests/test_litellm/responses/test_responses_websocket_all_providers.py b/tests/test_litellm/responses/test_responses_websocket_all_providers.py index 2d523bfdeb3..e8333214ea8 100644 --- a/tests/test_litellm/responses/test_responses_websocket_all_providers.py +++ b/tests/test_litellm/responses/test_responses_websocket_all_providers.py @@ -84,7 +84,7 @@ class TestResponsesAPIWebSocketSupport: def test_azure_websocket_url_requires_api_base(self): config = AzureOpenAIResponsesAPIConfig() - with pytest.raises(ValueError): + with pytest.raises(ValueError, match='api_base is required for Azure WebSocket'): config.get_websocket_url(api_base=None, litellm_params={}) def test_azure_model_not_in_websocket_url(self): diff --git a/tests/test_litellm/responses/test_streaming_iterator_error_events.py b/tests/test_litellm/responses/test_streaming_iterator_error_events.py index 3b87246ebdb..321abe4cc6d 100644 --- a/tests/test_litellm/responses/test_streaming_iterator_error_events.py +++ b/tests/test_litellm/responses/test_streaming_iterator_error_events.py @@ -208,9 +208,12 @@ async def test_async_iterator_error_after_first_chunk_carries_generated_content( ) chunks = [] - with pytest.raises(MidStreamFallbackError) as exc_info: + async def _drain(): async for chunk in iterator: chunks.append(chunk) + + with pytest.raises(MidStreamFallbackError) as exc_info: + await _drain() assert len(chunks) == 2 assert exc_info.value.status_code == 500 assert exc_info.value.is_pre_first_chunk is False diff --git a/tests/test_litellm/router_strategy/adaptive_router/test_bandit.py b/tests/test_litellm/router_strategy/adaptive_router/test_bandit.py index ab322f0fb37..78390cc1193 100644 --- a/tests/test_litellm/router_strategy/adaptive_router/test_bandit.py +++ b/tests/test_litellm/router_strategy/adaptive_router/test_bandit.py @@ -100,7 +100,7 @@ def test_score_combines_quality_and_cost(): def test_pick_best_empty_dict_raises(): - with pytest.raises(ValueError): + with pytest.raises(ValueError, match='pick_best called with no models'): pick_best({}, {}) diff --git a/tests/test_litellm/router_strategy/test_router_tag_routing.py b/tests/test_litellm/router_strategy/test_router_tag_routing.py index 73491490b14..60b1166de73 100644 --- a/tests/test_litellm/router_strategy/test_router_tag_routing.py +++ b/tests/test_litellm/router_strategy/test_router_tag_routing.py @@ -656,7 +656,7 @@ async def test_negation_all_excluded_raises(): enable_tag_filtering=True, ) - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Not allowed to access model due to tags configuration\\.') as exc_info: await router.acompletion( model="gpt-4", messages=[{"role": "user", "content": "hi"}], @@ -699,7 +699,7 @@ async def test_negation_ban_only_cannot_escape_default_pool(): # Sending only "!default" must NOT route to the paid deployment. # The base pool for ban-only is the default pool; banning the only # default deployment should raise rather than falling through to paid. - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Not allowed to access model due to tags configuration\\.') as exc_info: await router.acompletion( model="gpt-4", messages=[{"role": "user", "content": "hi"}], @@ -969,7 +969,7 @@ async def test_negation_exhausts_entire_fallback_chain(): enable_tag_filtering=True, ) - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Not allowed to access model due to tags configuration\\.') as exc_info: await router.acompletion( model="primary", messages=[{"role": "user", "content": "hi"}], @@ -1719,7 +1719,7 @@ async def test_required_and_unmatched_raises_by_default(): enable_tag_filtering=True, ) - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Not allowed to access model due to tags configuration\\.') as exc_info: await router.acompletion( model="gpt-4", messages=[{"role": "user", "content": "hi"}], @@ -1751,7 +1751,7 @@ async def test_required_and_combined_with_positive_unmatched_raises_by_default() enable_tag_filtering=True, ) - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Not allowed to access model due to tags configuration\\.') as exc_info: await router.acompletion( model="gpt-4", messages=[{"role": "user", "content": "hi"}], @@ -1973,7 +1973,7 @@ async def test_negation_combined_with_positive_unmatched_raises_by_default(): enable_tag_filtering=True, ) - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Not allowed to access model due to tags configuration\\.') as exc_info: await router.acompletion( model="gpt-4", messages=[{"role": "user", "content": "hi"}], @@ -2131,7 +2131,7 @@ async def test_mixed_constraint_survivor_unmatched_by_positive_tag_raises_by_def enable_tag_filtering=True, ) - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Not allowed to access model due to tags configuration\\.') as exc_info: await router.acompletion( model="gpt-4", messages=[{"role": "user", "content": "hi"}], @@ -2224,7 +2224,7 @@ async def test_allow_fail_open_denied_when_request_includes_unknown_tag(): enable_tag_filtering=True, ) - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Not allowed to access model due to tags configuration\\.') as exc_info: await router.acompletion( model="gpt-4", messages=[{"role": "user", "content": "hi"}], @@ -2538,7 +2538,7 @@ async def test_plain_tag_exhaustion_with_universal_default_tag_raises_by_default "litellm.router._async_get_cooldown_deployments", new=AsyncMock(return_value=["quality-high-1", "quality-high-2"]), ): - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Not allowed to access model due to tags configuration\\.') as exc_info: await router.acompletion( model="gpt-4", messages=[{"role": "user", "content": "hi"}], @@ -2767,7 +2767,7 @@ async def test_allow_fail_open_raises_when_inherited_constraint_alone_is_unsatis # allow_fail_open unset. router = _eu_region_router() - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Not allowed to access model due to tags configuration\\.') as exc_info: await router.acompletion( model="chat", messages=[{"role": "user", "content": "hi"}], @@ -2941,7 +2941,7 @@ async def test_tagged_request_direct_to_plain_group_still_rejected(): # tag filtering must reject exactly as before. router = _tagged_marker_router() - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Not allowed to access model due to tags configuration\\.') as exc_info: await router.acompletion( model="gemini-flash", messages=[{"role": "user", "content": "hi"}], @@ -2962,7 +2962,7 @@ async def test_caller_forged_consumption_stamp_is_neutralized_by_the_hook(): # tag filtering runs. router = _tagged_marker_router() - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Not allowed to access model due to tags configuration\\.') as exc_info: await router.acompletion( model="gemini-flash", messages=[{"role": "user", "content": "hi"}], @@ -2984,7 +2984,7 @@ async def test_inherited_constraint_still_applies_to_the_routed_tier(): # ®ion:eu comes from key/team policy (present in inherited_tags): # consuming the router-selecting "route" tag must not also discard the # inherited requirement, so a tier without the tag still raises... - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Not allowed to access model due to tags configuration\\.') as exc_info: await _tagged_marker_router().acompletion( model="gpt4o", messages=[{"role": "user", "content": "hi"}], diff --git a/tests/test_litellm/sandbox/test_e2b_sandbox.py b/tests/test_litellm/sandbox/test_e2b_sandbox.py index cc5b12156a1..e01b9120416 100644 --- a/tests/test_litellm/sandbox/test_e2b_sandbox.py +++ b/tests/test_litellm/sandbox/test_e2b_sandbox.py @@ -293,7 +293,7 @@ async def test_public_lifecycle_create_run_delete(): @pytest.mark.asyncio async def test_unsupported_provider_raises(): - with pytest.raises(ValueError): + with pytest.raises(ValueError, match="not-a-provider' is not a valid SandboxProviders"): await litellm.acreate_sandbox(provider="not-a-provider") diff --git a/tests/test_litellm/secret_managers/test_base_secret_manager.py b/tests/test_litellm/secret_managers/test_base_secret_manager.py index cba6a99ab7f..e1ccb91c381 100644 --- a/tests/test_litellm/secret_managers/test_base_secret_manager.py +++ b/tests/test_litellm/secret_managers/test_base_secret_manager.py @@ -32,7 +32,7 @@ from litellm.secret_managers.base_secret_manager import raise_if_unsafe_secret_n ], ) def test_raise_if_unsafe_secret_name_rejects_traversal_and_line_breaks(secret_name): - with pytest.raises(ValueError): + with pytest.raises(ValueError, match='Invalid secret_name'): raise_if_unsafe_secret_name(secret_name) diff --git a/tests/test_litellm/secret_managers/test_custom_secret_manager.py b/tests/test_litellm/secret_managers/test_custom_secret_manager.py index 1f4f9a47671..0426c5973cc 100644 --- a/tests/test_litellm/secret_managers/test_custom_secret_manager.py +++ b/tests/test_litellm/secret_managers/test_custom_secret_manager.py @@ -243,17 +243,17 @@ def test_minimal_custom_secret_manager(): assert value == "sync-TEST_KEY-value" # Write should raise NotImplementedError - with pytest.raises(NotImplementedError) as exc_info: - import asyncio + import asyncio + with pytest.raises(NotImplementedError) as exc_info: asyncio.run(secret_manager.async_write_secret("KEY", "value")) assert "Write operations are not implemented" in str(exc_info.value) # Delete should raise NotImplementedError - with pytest.raises(NotImplementedError) as exc_info: - import asyncio + import asyncio + with pytest.raises(NotImplementedError) as exc_info: asyncio.run(secret_manager.async_delete_secret("KEY")) assert "Delete operations are not implemented" in str(exc_info.value) diff --git a/tests/test_litellm/test_bedrock_batch_pricing.py b/tests/test_litellm/test_bedrock_batch_pricing.py new file mode 100644 index 00000000000..856085ec253 --- /dev/null +++ b/tests/test_litellm/test_bedrock_batch_pricing.py @@ -0,0 +1,43 @@ +import json +from pathlib import Path + +import pytest + +PRICING_FILES = ( + "model_prices_and_context_window.json", + "litellm/model_prices_and_context_window_backup.json", +) + +BEDROCK_BATCH_MODELS = ( + "qwen.qwen3-235b-a22b-2507-v1:0", + "anthropic.claude-haiku-4-5-20251001-v1:0", + "apac.anthropic.claude-haiku-4-5-20251001-v1:0", + "au.anthropic.claude-haiku-4-5-20251001-v1:0", + "eu.anthropic.claude-haiku-4-5-20251001-v1:0", + "global.anthropic.claude-haiku-4-5-20251001-v1:0", + "jp.anthropic.claude-haiku-4-5-20251001-v1:0", + "us.anthropic.claude-haiku-4-5-20251001-v1:0", + "anthropic.claude-sonnet-4-5-20250929-v1:0", + "au.anthropic.claude-sonnet-4-5-20250929-v1:0", + "claude-sonnet-4-5-20250929-v1:0", + "eu.anthropic.claude-sonnet-4-5-20250929-v1:0", + "global.anthropic.claude-sonnet-4-5-20250929-v1:0", + "jp.anthropic.claude-sonnet-4-5-20250929-v1:0", + "us.anthropic.claude-sonnet-4-5-20250929-v1:0", +) + + +@pytest.mark.parametrize("pricing_file", PRICING_FILES) +@pytest.mark.parametrize("model", BEDROCK_BATCH_MODELS) +def test_bedrock_batch_pricing_is_half_of_on_demand( + pricing_file: str, model: str +) -> None: + model_cost_map = json.loads((Path(__file__).parents[2] / pricing_file).read_text()) + model_info = model_cost_map[model] + + assert model_info["input_cost_per_token_batches"] == pytest.approx( + model_info["input_cost_per_token"] / 2 + ) + assert model_info["output_cost_per_token_batches"] == pytest.approx( + model_info["output_cost_per_token"] / 2 + ) diff --git a/tests/test_litellm/test_dashscope_image_generation.py b/tests/test_litellm/test_dashscope_image_generation.py index af95e2ca6b4..c9f0df4febb 100644 --- a/tests/test_litellm/test_dashscope_image_generation.py +++ b/tests/test_litellm/test_dashscope_image_generation.py @@ -17,6 +17,7 @@ from litellm.llms.dashscope.image_generation.transformation import ( ) from litellm.types.utils import ImageObject, ImageResponse from litellm.utils import get_llm_provider +from litellm.llms.base_llm.chat.transformation import BaseLLMException # --------------------------------------------------------------------------- @@ -247,7 +248,7 @@ class TestDashScopeImageGenerationConfig: "message": "Size not supported", } - with pytest.raises(Exception): + with pytest.raises(BaseLLMException): self.cfg.transform_image_generation_response( model="qwen-image-2.0", raw_response=mock_resp, @@ -268,7 +269,7 @@ class TestDashScopeImageGenerationConfig: "message": "Size not supported", } - with pytest.raises(Exception): + with pytest.raises(BaseLLMException): self.cfg.transform_image_generation_response( model="qwen-image-2.0", raw_response=mock_resp, diff --git a/tests/test_litellm/test_daybreak_model_metadata.py b/tests/test_litellm/test_daybreak_model_metadata.py new file mode 100644 index 00000000000..d04cca3c077 --- /dev/null +++ b/tests/test_litellm/test_daybreak_model_metadata.py @@ -0,0 +1,52 @@ +import json +from pathlib import Path + +import pytest + +REPO_ROOT = Path(__file__).parents[2] +MAIN_PATH = REPO_ROOT / "model_prices_and_context_window.json" +BACKUP_PATH = REPO_ROOT / "litellm" / "model_prices_and_context_window_backup.json" + +DAYBREAK_MODELS = ( + "gpt-5.6-cyber", + "daybreak-red-latest", + "daybreak-blue-latest", +) +BLUE_ALIAS = "daybreak-blue-latest" +BLUE_SNAPSHOT = "gpt-5.6-sol" + + +def _load(path): + with open(path) as f: + return json.load(f) + + +@pytest.mark.parametrize("model", DAYBREAK_MODELS) +def test_daybreak_capability_contract(model): + info = _load(MAIN_PATH).get(model) + assert info is not None, f"{model} missing from model_prices_and_context_window.json" + + assert info["litellm_provider"] == "openai" + assert info["mode"] == "chat" + assert info["supported_endpoints"] == ["/v1/chat/completions", "/v1/responses"] + + assert info["supports_computer_use"] is True + assert info["supports_parallel_function_calling"] is True + assert info["supports_function_calling"] is True + assert info["supports_reasoning"] is True + assert info["supports_vision"] is True + + +def test_blue_alias_matches_its_snapshot_computer_use(): + cost_map = _load(MAIN_PATH) + + assert cost_map[BLUE_ALIAS]["supports_computer_use"] is True + assert cost_map[BLUE_SNAPSHOT]["supports_computer_use"] is True + + +@pytest.mark.parametrize("model", (*DAYBREAK_MODELS, BLUE_SNAPSHOT)) +def test_backup_matches_main(model): + main_cost = _load(MAIN_PATH) + backup_cost = _load(BACKUP_PATH) + + assert backup_cost.get(model) == main_cost.get(model), f"{model} differs between main and backup model cost maps" diff --git a/tests/test_litellm/test_gemini_3_1_flash_lite_image_pricing.py b/tests/test_litellm/test_gemini_3_1_flash_lite_image_pricing.py new file mode 100644 index 00000000000..67d6b9e76cf --- /dev/null +++ b/tests/test_litellm/test_gemini_3_1_flash_lite_image_pricing.py @@ -0,0 +1,147 @@ +"""Pricing entry for ``gemini-3.1-flash-lite-image`` (Google's Nano Banana 2 Lite). + +Google publishes: $0.25/1M input, $1.50/1M text output, and $30/1M image-output +tokens for the Lite image model (https://cloud.google.com/vertex-ai/generative-ai/pricing). +A 1K image is ~1120 output image tokens => ~$0.0336 / image. + +Without this entry, ``completion_cost`` raises "model isn't mapped yet" and Vertex +generateContent pass-through cost tracking silently logs $0. These tests pin the +values in both the primary price map and the ``litellm/`` backup, and verify +``get_model_info`` / ``completion_cost`` surface them. +""" + +import json +import os + +import litellm +from litellm import completion_cost +from litellm.types.utils import CompletionTokensDetailsWrapper, ModelResponse, Usage + +VARIANTS = [ + "gemini-3.1-flash-lite-image", + "gemini/gemini-3.1-flash-lite-image", + "vertex_ai/gemini-3.1-flash-lite-image", +] + +EXPECTED = { + "input_cost_per_token": 2.5e-07, + "output_cost_per_token": 1.5e-06, + "output_cost_per_image_token": 3e-05, + "mode": "image_generation", +} + +EXPECTED_CAPABILITIES = { + "max_output_tokens": 4096, + "max_tokens": 4096, + "supports_response_schema": False, + "supports_reasoning": True, +} + +EXPECTED_PER_ROUTE = { + "gemini-3.1-flash-lite-image": { + "supports_prompt_caching": True, + "supports_function_calling": False, + }, + "vertex_ai/gemini-3.1-flash-lite-image": { + "supports_prompt_caching": True, + "supports_function_calling": False, + }, + "gemini/gemini-3.1-flash-lite-image": { + "supports_prompt_caching": False, + "supports_function_calling": True, + "input_cost_per_token_batches": 1.25e-07, + "output_cost_per_token_batches": 7.5e-07, + }, +} + + +def _load_json(path: str) -> dict: + with open(path, encoding="utf-8") as f: + return json.load(f) + + +def _backup_path() -> str: + return os.path.join( + os.path.dirname(litellm.__file__), + "model_prices_and_context_window_backup.json", + ) + + +def _main_path() -> str: + return os.path.join( + os.path.dirname(__file__), "..", "..", "model_prices_and_context_window.json" + ) + + +class TestGeminiFlashLiteImagePricingData: + """Both price maps must carry Google's published Nano Banana 2 Lite costs.""" + + def test_present_in_both_maps(self): + main = _load_json(_main_path()) + backup = _load_json(_backup_path()) + for key in VARIANTS: + for label, data in (("main", main), ("backup", backup)): + assert key in data, f"{key} missing from {label} JSON" + entry = data[key] + for field, value in EXPECTED.items(): + assert entry[field] == value, f"{key} {field} in {label}: {entry.get(field)} != {value}" + + def test_capabilities_match_model_cards(self): + main = _load_json(_main_path()) + backup = _load_json(_backup_path()) + for key in VARIANTS: + expected = {**EXPECTED_CAPABILITIES, **EXPECTED_PER_ROUTE[key]} + for label, data in (("main", main), ("backup", backup)): + entry = data[key] + for field, value in expected.items(): + assert entry[field] == value, f"{key} {field} in {label}: {entry.get(field)} != {value}" + + def test_grounding_fields_absent(self): + """Grounding with Google Search is unsupported on Lite, so no search pricing.""" + for path in (_main_path(), _backup_path()): + data = _load_json(path) + for key in VARIANTS: + for field in ( + "supports_web_search", + "search_context_cost_per_query", + "web_search_billing_unit", + ): + assert field not in data[key], f"{key} should not define {field}" + + def test_image_output_pricing_consistent(self): + """1120 image-output tokens * output_cost_per_image_token == output_cost_per_image.""" + backup = _load_json(_backup_path()) + entry = backup["gemini-3.1-flash-lite-image"] + assert round(1120 * entry["output_cost_per_image_token"], 6) == entry["output_cost_per_image"] + + +class TestGeminiFlashLiteImageModelInfo: + """``get_model_info`` and ``completion_cost`` must report the new costs.""" + + def test_get_model_info_and_cost(self): + original = litellm.model_cost + try: + litellm.model_cost = _load_json(_backup_path()) + info = litellm.get_model_info("gemini-3.1-flash-lite-image") + assert info["input_cost_per_token"] == EXPECTED["input_cost_per_token"] + assert info["output_cost_per_token"] == EXPECTED["output_cost_per_token"] + + resp = ModelResponse() + resp.model = "gemini-3.1-flash-lite-image" + resp.usage = Usage( + prompt_tokens=7, + completion_tokens=1120, + total_tokens=1127, + completion_tokens_details=CompletionTokensDetailsWrapper( + image_tokens=1120, text_tokens=0 + ), + ) + cost = completion_cost( + completion_response=resp, + model="gemini-3.1-flash-lite-image", + custom_llm_provider="vertex_ai", + ) + expected_cost = 1120 * 3e-05 + 7 * 2.5e-07 + assert abs(cost - expected_cost) < 1e-6, f"unexpected cost {cost}" + finally: + litellm.model_cost = original diff --git a/tests/test_litellm/test_get_blog_posts.py b/tests/test_litellm/test_get_blog_posts.py index 32edc5423d0..241dce23633 100644 --- a/tests/test_litellm/test_get_blog_posts.py +++ b/tests/test_litellm/test_get_blog_posts.py @@ -12,6 +12,7 @@ from litellm.litellm_core_utils.get_blog_posts import ( GetBlogPosts, get_blog_posts, ) +from xml.etree import ElementTree SAMPLE_RSS = """\ @@ -71,7 +72,7 @@ def test_parse_rss_to_posts_multiple(): def test_parse_rss_to_posts_invalid_xml(): - with pytest.raises(Exception): + with pytest.raises(ElementTree.ParseError): GetBlogPosts.parse_rss_to_posts("not xml") diff --git a/tests/test_litellm/test_github_close_low_quality_prs.py b/tests/test_litellm/test_github_close_low_quality_prs.py index e3b653dde64..2a891ca72f5 100644 --- a/tests/test_litellm/test_github_close_low_quality_prs.py +++ b/tests/test_litellm/test_github_close_low_quality_prs.py @@ -697,7 +697,7 @@ class TestListOpenItemsNoCap: def test_list_open_items_rejects_unknown_kind(self, closer_module): shared = self._shared(closer_module) - with pytest.raises(ValueError): + with pytest.raises(ValueError, match="kind must be 'pr' or 'issue', got 'both"): shared.list_open_items("both", repo="o/r", fields="number") def test_fetch_open_prs_delegates_with_no_cap(self, closer_module, monkeypatch): diff --git a/tests/test_litellm/test_github_triage_with_llm.py b/tests/test_litellm/test_github_triage_with_llm.py index 96b77e80457..ddffb978b48 100644 --- a/tests/test_litellm/test_github_triage_with_llm.py +++ b/tests/test_litellm/test_github_triage_with_llm.py @@ -665,11 +665,11 @@ class TestParseVerdict: assert triage_module.parse_verdict(raw)["verdict"] == "pass" def test_should_raise_for_unparseable_text(self, triage_module): - with pytest.raises(ValueError): + with pytest.raises(ValueError, match='could not extract JSON from LLM response: not even close to'): triage_module.parse_verdict("not even close to json") def test_should_raise_for_empty(self, triage_module): - with pytest.raises(ValueError): + with pytest.raises(ValueError, match='empty LLM response'): triage_module.parse_verdict("") diff --git a/tests/test_litellm/test_mistral_small_4_0_model_metadata.py b/tests/test_litellm/test_mistral_small_4_0_model_metadata.py new file mode 100644 index 00000000000..0442321ba0b --- /dev/null +++ b/tests/test_litellm/test_mistral_small_4_0_model_metadata.py @@ -0,0 +1,49 @@ +import json +from pathlib import Path + +import pytest + +REPO_ROOT = Path(__file__).parents[2] +MAIN_PATH = REPO_ROOT / "model_prices_and_context_window.json" +BACKUP_PATH = REPO_ROOT / "litellm" / "model_prices_and_context_window_backup.json" + +SMALL_4_0_MODELS = ( + "mistral/mistral-small-latest", + "mistral/mistral-small-2603", +) + + +def _load(path): + with open(path) as f: + return json.load(f) + + +@pytest.mark.parametrize("model", SMALL_4_0_MODELS) +def test_small_4_0_specs(model): + info = _load(MAIN_PATH).get(model) + assert info is not None, f"{model} missing from model_prices_and_context_window.json" + + assert info["litellm_provider"] == "mistral" + assert info["mode"] == "chat" + + assert info["input_cost_per_token"] == 1.5e-07 + assert info["output_cost_per_token"] == 6e-07 + + assert info["max_input_tokens"] == 262144 + assert info["max_output_tokens"] == 262144 + assert info["max_tokens"] == 262144 + + assert info["supports_reasoning"] is True + assert info["supports_vision"] is True + assert info["supports_function_calling"] is True + assert info["supports_response_schema"] is True + assert info["supports_tool_choice"] is True + assert info["supports_assistant_prefill"] is True + + +@pytest.mark.parametrize("model", SMALL_4_0_MODELS) +def test_backup_matches_main(model): + main_cost = _load(MAIN_PATH) + backup_cost = _load(BACKUP_PATH) + + assert backup_cost.get(model) == main_cost.get(model), f"{model} differs between main and backup model cost maps" diff --git a/tests/test_litellm/test_project_tags_pydantic.py b/tests/test_litellm/test_project_tags_pydantic.py index b3f58df2325..c04cf2c686b 100644 --- a/tests/test_litellm/test_project_tags_pydantic.py +++ b/tests/test_litellm/test_project_tags_pydantic.py @@ -1,5 +1,6 @@ import pytest from litellm.proxy._types import NewProjectRequest, UpdateProjectRequest +from pydantic import ValidationError def test_new_project_request_tags(): @@ -21,11 +22,11 @@ def test_update_project_request_tags(): def test_new_project_request_invalid_tags_type(): # tags must be a list — a string should raise a ValidationError - with pytest.raises(Exception): + with pytest.raises(ValidationError): NewProjectRequest(project_id="test_proj", team_id="team_1", tags="not-a-list") def test_update_project_request_invalid_tags_type(): # tags must be a list — a string should raise a ValidationError - with pytest.raises(Exception): + with pytest.raises(ValidationError): UpdateProjectRequest(project_id="test_proj", tags="not-a-list") diff --git a/tests/test_litellm/test_redact_string_in_error_paths.py b/tests/test_litellm/test_redact_string_in_error_paths.py index d01a9da6617..1c4d91397d1 100644 --- a/tests/test_litellm/test_redact_string_in_error_paths.py +++ b/tests/test_litellm/test_redact_string_in_error_paths.py @@ -234,7 +234,7 @@ class TestRouterFallbackFailureTracebackRedaction: raise ValueError(f"primary deployment failed api_key={secret}") except ValueError as original_exception: with caplog.at_level(logging.DEBUG, logger="LiteLLM Router"): - with pytest.raises(Exception): + with pytest.raises(ValueError, match='primary deployment failed api_key=sk-testsecretvalu'): await router.async_function_with_fallbacks_common_utils( e=original_exception, disable_fallbacks=False, diff --git a/tests/test_litellm/test_redis.py b/tests/test_litellm/test_redis.py index 3aa4bc58f13..c645a67ef84 100644 --- a/tests/test_litellm/test_redis.py +++ b/tests/test_litellm/test_redis.py @@ -131,7 +131,7 @@ def test_get_redis_url_from_environment_missing_host_port(monkeypatch): monkeypatch.delenv("REDIS_PORT", raising=False) # Call the function and expect a ValueError - with pytest.raises(ValueError) as excinfo: + with pytest.raises(ValueError, match="Either 'REDIS_URL' or both 'REDIS_HOST' and 'REDIS_PORT") as excinfo: get_redis_url_from_environment() # Check the error message @@ -149,7 +149,7 @@ def test_get_redis_url_from_environment_missing_port(monkeypatch): monkeypatch.setenv("REDIS_HOST", "redis-server") # Call the function and expect a ValueError - with pytest.raises(ValueError) as excinfo: + with pytest.raises(ValueError, match="Either 'REDIS_URL' or both 'REDIS_HOST' and 'REDIS_PORT") as excinfo: get_redis_url_from_environment() # Check the error message diff --git a/tests/test_litellm/test_retrieve_batch_bedrock_dispatch.py b/tests/test_litellm/test_retrieve_batch_bedrock_dispatch.py index 9df18a9f0f0..aa057b7bc73 100644 --- a/tests/test_litellm/test_retrieve_batch_bedrock_dispatch.py +++ b/tests/test_litellm/test_retrieve_batch_bedrock_dispatch.py @@ -23,6 +23,7 @@ import pytest sys.path.insert(0, os.path.abspath("../..")) import litellm # noqa: E402 +import openai ASYNC_INVOKE_ARN = "arn:aws:bedrock:us-west-2:123456789012:async-invoke/abc123def456" MIJ_ARN = "arn:aws:bedrock:us-west-2:123456789012:model-invocation-job/abc1234567" @@ -134,7 +135,7 @@ def test_unrelated_bedrock_arn_falls_through_to_provider_config(mock_handlers): # Use a plausible-but-unsupported Bedrock ARN family. unrelated_arn = "arn:aws:bedrock:us-west-2:123456789012:provisioned-model/xyz" - with pytest.raises(Exception): + with pytest.raises(litellm.BadRequestError): # Will raise because no provider_config exists for this path — # that's fine, we just need to assert neither bedrock handler ran # before the failure. @@ -152,7 +153,7 @@ def test_non_bedrock_id_skips_bedrock_dispatch_entirely(mock_handlers): block — they belong to other providers' retrieve flows.""" async_invoke, mij, _ = mock_handlers - with pytest.raises(Exception): + with pytest.raises(openai.OpenAIError): litellm.retrieve_batch( batch_id="batch_abc123", custom_llm_provider="openai", diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 4a06a5dfb2e..a47525749f6 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -367,7 +367,7 @@ async def test_arouter_with_tags_and_fallbacks(): enable_tag_filtering=True, ) - with pytest.raises(Exception): + with pytest.raises(litellm.InternalServerError): response = await router.acompletion( model="gpt-3.5-turbo", messages=[{"role": "user", "content": "Hello, world!"}], @@ -953,7 +953,7 @@ async def test_arouter_filter_team_based_models(): assert result is not None # FAILS - with pytest.raises(Exception) as e: + with pytest.raises(Exception, match='No deployments available for selected model, Try again in') as e: result = await router.acompletion( model="gpt-3.5-turbo", messages=[{"role": "user", "content": "Hello, world!"}], @@ -1225,7 +1225,7 @@ def test_add_invalid_provider_to_router(): ], ) - with pytest.raises(Exception) as e: + with pytest.raises(Exception, match='Unsupported provider - vertex_ai_eu') as e: router.add_deployment( Deployment( model_name="vertex_ai/*", @@ -1320,7 +1320,7 @@ async def test_router_ageneric_api_call_with_fallbacks_helper(): with patch.object(router, "async_get_available_deployment") as mock_get_deployment: mock_get_deployment.side_effect = Exception("No deployment available") - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='No deployment available') as exc_info: await router._ageneric_api_call_with_fallbacks_helper( model="gpt-3.5-turbo", original_generic_function=mock_generic_function, @@ -1394,7 +1394,7 @@ async def test_router_ageneric_api_call_with_fallbacks_helper(): with patch.object( router, "async_routing_strategy_pre_call_checks" ) as mock_pre_call_checks: - with pytest.raises(Exception) as exc_info: + with pytest.raises(Exception, match='Mock failure') as exc_info: await router._ageneric_api_call_with_fallbacks_helper( model="gpt-3.5-turbo", original_generic_function=mock_failing_function, @@ -1999,10 +1999,13 @@ async def test_acompletion_streaming_iterator(): # Collect streamed chunks — the first chunk succeeds, then the error re-raises collected_chunks = [] - with pytest.raises(MidStreamFallbackError): + async def _drain(): async for chunk in result: collected_chunks.append(chunk) + with pytest.raises(MidStreamFallbackError): + await _drain() + assert len(collected_chunks) == 1, "one chunk yielded before the error" print("✓ MidStreamFallbackError re-raised correctly when content was already generated") @@ -3734,7 +3737,7 @@ def test_count_pre_call_check_tokens_across_api_surfaces(): assert string_input_tokens > 0 assert list_input_tokens > 0 - with pytest.raises(ValueError): + with pytest.raises(ValueError, match='Either messages or input must be provided to count tokens'): router._count_pre_call_check_tokens(messages=None, input=None) @@ -5557,10 +5560,13 @@ async def test_acompletion_streaming_iterator_does_not_log_success_on_terminal_f initial_kwargs=dict(initial_kwargs), ) collected = [] - with pytest.raises(MidStreamFallbackError): + async def _drain(): async for chunk in result: collected.append(chunk) + with pytest.raises(MidStreamFallbackError): + await _drain() + assert len(collected) == 1 logging_obj.dispatch_success_handlers.assert_not_called() @@ -5580,10 +5586,13 @@ async def test_acompletion_streaming_iterator_does_not_log_success_on_terminal_f initial_kwargs=dict(initial_kwargs), ) collected = [] - with pytest.raises(MidStreamFallbackError): + async def _drain(): async for chunk in result: collected.append(chunk) + with pytest.raises(MidStreamFallbackError): + await _drain() + assert len(collected) == 1, "only the partial chunk before the error" mock_fallback.assert_not_called() logging_obj.dispatch_success_handlers.assert_not_called() @@ -5707,7 +5716,7 @@ async def test_team_scoped_model_fallback_cross_team_blocked(): fallbacks=[{"primary-model": ["fallback-model"]}], ) - with pytest.raises(Exception): + with pytest.raises(litellm.InternalServerError): await router.acompletion( model="primary-model", messages=[{"role": "user", "content": "Hello"}], diff --git a/tests/test_litellm/test_router_model_cost_isolation.py b/tests/test_litellm/test_router_model_cost_isolation.py index 5bb854c12e0..dc210f900bf 100644 --- a/tests/test_litellm/test_router_model_cost_isolation.py +++ b/tests/test_litellm/test_router_model_cost_isolation.py @@ -1706,7 +1706,7 @@ def test_an_incomplete_reservation_is_refused_rather_than_served(dropped): state the operator was trying to leave.""" incomplete = {k: v for k, v in _PTU_MODEL_INFO.items() if k != dropped} - with pytest.raises(ValueError) as raised: + with pytest.raises(ValueError, match="PTU configuration on model 'gpt") as raised: _ptu_router(model_info=incomplete, litellm_params={"input_cost_per_token": 5e-06}) assert "gpt-4o-ptu" in str(raised.value) @@ -1726,7 +1726,7 @@ def test_the_refusal_reason_is_the_one_the_model_endpoint_answers_with(dropped, incomplete = {k: v for k, v in _PTU_MODEL_INFO.items() if k != dropped} assert ptu_config_error(incomplete) == expected - with pytest.raises(ValueError) as raised: + with pytest.raises(ValueError, match="PTU configuration on model 'gpt") as raised: _ptu_router(model_info=incomplete) assert expected in str(raised.value) diff --git a/tests/test_litellm/test_router_weighted_failover.py b/tests/test_litellm/test_router_weighted_failover.py index 0115638e1fe..162312a8c67 100644 --- a/tests/test_litellm/test_router_weighted_failover.py +++ b/tests/test_litellm/test_router_weighted_failover.py @@ -391,7 +391,7 @@ async def test_no_failover_when_flag_off(): # enable_weighted_failover defaults to False ) - with pytest.raises(Exception): + with pytest.raises(litellm.InternalServerError): await router.acompletion( model="test-model", messages=[{"role": "user", "content": "hi"}], @@ -515,7 +515,7 @@ async def test_failover_exhausted_raises_original_error_class(): enable_weighted_failover=True, ) - with pytest.raises(Exception): + with pytest.raises(litellm.InternalServerError): await router.acompletion( model="test-model", messages=[{"role": "user", "content": "hi"}], @@ -648,7 +648,7 @@ async def test_failover_skipped_for_non_simple_shuffle(): enable_weighted_failover=True, ) - with pytest.raises(Exception): + with pytest.raises(litellm.InternalServerError): await router.acompletion( model="test-model", messages=[{"role": "user", "content": "hi"}], diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index 041c60e0ba6..075b455e4b5 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -4370,7 +4370,7 @@ class TestVertexEmbeddingEncodingFormat: assert "encoding_format" not in optional_params def test_encoding_format_base64_still_rejected_without_drop_params(self): - with pytest.raises(Exception) as excinfo: + with pytest.raises(Exception, match='To drop these, set `litellm\\.drop_params=True` or for proxy') as excinfo: litellm.utils.get_optional_params_embeddings( model="gemini-embedding-001", encoding_format="base64", diff --git a/tests/test_litellm/test_video_generation.py b/tests/test_litellm/test_video_generation.py index 3d0472ef96e..117ca72c34f 100644 --- a/tests/test_litellm/test_video_generation.py +++ b/tests/test_litellm/test_video_generation.py @@ -151,7 +151,7 @@ class TestVideoGeneration: "video_generation_handler", side_effect=Exception("API Error"), ): - with pytest.raises(Exception): + with pytest.raises(litellm.APIConnectionError): video_generation(prompt="Test video", model="sora-2") def test_video_generation_provider_config(self): @@ -739,7 +739,7 @@ class TestVideoGeneration: "video_status_handler", side_effect=Exception("API Error"), ): - with pytest.raises(Exception): + with pytest.raises(litellm.APIConnectionError): video_status(video_id="test_video_id", model="sora-2") def test_video_status_request_transformation(self): diff --git a/tests/test_litellm/types/test_router.py b/tests/test_litellm/types/test_router.py index 5ce5eca4954..accd3b32a0d 100644 --- a/tests/test_litellm/types/test_router.py +++ b/tests/test_litellm/types/test_router.py @@ -87,5 +87,5 @@ def test_pricing_strings_are_coerced_to_float(): def test_invalid_pricing_is_rejected(): - with pytest.raises(ValueError): + with pytest.raises(ValueError, match='validation error for ModelInfo'): ModelInfo(id="x", input_cost_per_token="free") diff --git a/tests/test_litellm/videos/test_main.py b/tests/test_litellm/videos/test_main.py index a04a89ded99..38667e93eee 100644 --- a/tests/test_litellm/videos/test_main.py +++ b/tests/test_litellm/videos/test_main.py @@ -321,7 +321,7 @@ def test_get_character__mock_response_short_circuits(seams): def test_unsupported_provider_raises_without_dispatch(seams): seams.get_config.return_value = None - with pytest.raises(Exception): + with pytest.raises(litellm.APIConnectionError): videos_main.video_status(video_id=AZURE_VIDEO_ID) seams.handler.video_status_handler.assert_not_called() diff --git a/tests/test_ratelimit.py b/tests/test_ratelimit.py index 0469ded3f42..121dfbd99b7 100644 --- a/tests/test_ratelimit.py +++ b/tests/test_ratelimit.py @@ -149,19 +149,26 @@ def test_async_rate_limit( router: Router = router_factory(rpm, tpm, routing_strategy) print(f"router: {router.model_list}") - with pytest.raises(expected_exception) as excinfo: # asserts correct type raised - if sync_mode: - results = sync_call(router, list_of_messages) - else: - results = asyncio.run(async_call(router, list_of_messages)) + received = [] + + def _send_and_check(): + results = ( + sync_call(router, list_of_messages) + if sync_mode + else asyncio.run(async_call(router, list_of_messages)) + ) + received.extend(results) print(results) if len([i for i in results if i is not None]) != num_try_send: # since not all results got returned, raise rate limit error raise ValueError("No deployments available for selected model") raise ExpectNoException + with pytest.raises(expected_exception) as excinfo: # asserts correct type raised + _send_and_check() + print(expected_exception, excinfo) if expected_exception is ValueError: assert "No deployments available for selected model" in str(excinfo.value) else: - assert len([i for i in results if i is not None]) == num_try_send + assert len([i for i in received if i is not None]) == num_try_send diff --git a/tests/vector_store_tests/test_ragflow_vector_store.py b/tests/vector_store_tests/test_ragflow_vector_store.py index cb4cfd75c1f..46751b64cce 100644 --- a/tests/vector_store_tests/test_ragflow_vector_store.py +++ b/tests/vector_store_tests/test_ragflow_vector_store.py @@ -16,6 +16,7 @@ from tests.vector_store_tests.base_vector_store_test import BaseVectorStoreTest from litellm.llms.ragflow.vector_stores.transformation import RAGFlowVectorStoreConfig from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.types.vector_stores import VectorStoreCreateOptionalRequestParams +from litellm.llms.base_llm.chat.transformation import BaseLLMException class TestRAGFlowVectorStore(BaseVectorStoreTest): @@ -233,7 +234,7 @@ class TestRAGFlowVectorStore(BaseVectorStoreTest): "message": "Dataset name 'test-dataset' already exists", } - with pytest.raises(Exception): # Should raise BaseLLMException + with pytest.raises(BaseLLMException): config.transform_create_vector_store_response(mock_response) def test_transform_create_vector_store_response_missing_id(self): diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/cost_tracking_settings.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/cost_tracking_settings.tsx index 3fce928ed67..52b66edcbe5 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/cost_tracking_settings.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/cost_tracking_settings.tsx @@ -11,6 +11,7 @@ import { } from "@/components/ui/alert-dialog"; import { Button } from "@/components/ui/button"; import { Collapsible, CollapsibleContent, CollapsibleTrigger } from "@/components/ui/collapsible"; +import { Switch } from "@/components/ui/switch"; import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs"; import { CostTrackingSettingsProps } from "./types"; import ProviderDiscountTable from "./provider_discount_table"; @@ -22,6 +23,7 @@ import { DocsMenu } from "@/components/HelpLink"; import HowItWorks from "./how_it_works"; import { useDiscountConfig } from "./use_discount_config"; import { useMarginConfig } from "./use_margin_config"; +import { useBlockUnpricedConfig } from "./use_block_unpriced_config"; import { fetchAvailableModels, ModelGroup } from "@/components/llm_calls/fetch_models"; import { Dialog, DialogContent, DialogHeader, DialogTitle } from "@/components/ui/dialog"; @@ -86,9 +88,16 @@ const CostTrackingSettings: React.FC = ({ userID, use handleMarginChange, } = useMarginConfig({ accessToken }); + const { + blockUnpriced, + isUpdating: isUpdatingBlockUnpriced, + fetchBlockUnpriced, + setBlockUnpriced, + } = useBlockUnpricedConfig({ accessToken }); + useEffect(() => { if (accessToken) { - Promise.all([fetchDiscountConfig(), fetchMarginConfig()]).finally(() => { + Promise.all([fetchDiscountConfig(), fetchMarginConfig(), fetchBlockUnpriced()]).finally(() => { setIsFetching(false); }); @@ -103,7 +112,7 @@ const CostTrackingSettings: React.FC = ({ userID, use }; loadModels(); } - }, [accessToken, fetchDiscountConfig, fetchMarginConfig]); + }, [accessToken, fetchDiscountConfig, fetchMarginConfig, fetchBlockUnpriced]); const handleAddProvider = async () => { const success = await addProvider(selectedProvider, newDiscount); @@ -301,7 +310,35 @@ const CostTrackingSettings: React.FC = ({ userID, use )} - {/* Accordion 3: Pricing Calculator - Available to all roles */} + {/* Accordion 3: Block Unpriced Models - Only for proxy admins */} + {isProxyAdmin && ( + + + +
+
+
+

Block requests for models without pricing

+

+ When enabled, a request whose resolved model has no cost mapping is rejected with a 403 so an + admin can add pricing for it. Off by default +

+
+ setBlockUnpriced(checked)} + /> +
+
+ + + )} + + {/* Accordion 4: Pricing Calculator - Available to all roles */} ({ + apiClient: { + get: vi.fn(), + patch: vi.fn(), + }, +})); + +const ENDPOINT = "/config/block_requests_for_models_without_pricing"; + +describe("useBlockUnpricedConfig", () => { + beforeEach(() => { + vi.clearAllMocks(); + }); + + describe("fetchBlockUnpriced", () => { + it("reflects the enabled flag returned by the proxy", async () => { + vi.mocked(apiClient.get).mockResolvedValueOnce({ enabled: true }); + + const { result } = renderHook(() => useBlockUnpricedConfig({ accessToken: "test-token" })); + + await act(async () => { + await result.current.fetchBlockUnpriced(); + }); + + expect(apiClient.get).toHaveBeenCalledWith(ENDPOINT, { accessToken: "test-token" }); + expect(result.current.blockUnpriced).toBe(true); + }); + + it("surfaces a toast when the fetch throws", async () => { + const error = new Error("Network error"); + vi.mocked(apiClient.get).mockRejectedValueOnce(error); + + const { result } = renderHook(() => useBlockUnpricedConfig({ accessToken: "test-token" })); + + await act(async () => { + await result.current.fetchBlockUnpriced(); + }); + + expect(toast.fromError).toHaveBeenCalledWith(error); + expect(result.current.blockUnpriced).toBe(false); + }); + + it("does nothing without an access token", async () => { + const { result } = renderHook(() => useBlockUnpricedConfig({ accessToken: null })); + + await act(async () => { + await result.current.fetchBlockUnpriced(); + }); + + expect(apiClient.get).not.toHaveBeenCalled(); + }); + }); + + describe("setBlockUnpriced", () => { + it("persists the new value and confirms it with a toast", async () => { + vi.mocked(apiClient.patch).mockResolvedValueOnce({ enabled: true }); + + const { result } = renderHook(() => useBlockUnpricedConfig({ accessToken: "test-token" })); + + await act(async () => { + await result.current.setBlockUnpriced(true); + }); + + expect(apiClient.patch).toHaveBeenCalledWith(ENDPOINT, { + accessToken: "test-token", + body: { enabled: true }, + }); + expect(result.current.blockUnpriced).toBe(true); + expect(toast.success).toHaveBeenCalledWith(expect.stringMatching(/will now be blocked/i)); + expect(result.current.isUpdating).toBe(false); + }); + + it("confirms turning the block back off", async () => { + vi.mocked(apiClient.patch).mockResolvedValueOnce({ enabled: false }); + + const { result } = renderHook(() => useBlockUnpricedConfig({ accessToken: "test-token" })); + + await act(async () => { + await result.current.setBlockUnpriced(false); + }); + + expect(result.current.blockUnpriced).toBe(false); + expect(toast.success).toHaveBeenCalledWith(expect.stringMatching(/now allowed/i)); + }); + + it("surfaces the proxy error and leaves the flag unchanged when the update fails", async () => { + const error = new Error("Set `'STORE_MODEL_IN_DB='True'` in your env to enable this feature."); + vi.mocked(apiClient.patch).mockRejectedValueOnce(error); + + const { result } = renderHook(() => useBlockUnpricedConfig({ accessToken: "test-token" })); + + await act(async () => { + await result.current.setBlockUnpriced(true); + }); + + expect(toast.fromError).toHaveBeenCalledWith(error); + expect(toast.success).not.toHaveBeenCalled(); + expect(result.current.blockUnpriced).toBe(false); + expect(result.current.isUpdating).toBe(false); + }); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/use_block_unpriced_config.ts b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/use_block_unpriced_config.ts new file mode 100644 index 00000000000..4bf9d5ddacc --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/use_block_unpriced_config.ts @@ -0,0 +1,65 @@ +import { useState, useCallback } from "react"; +import { apiClient } from "@/components/networking"; +import { toast } from "@/lib/toast"; + +export interface UseBlockUnpricedConfigProps { + accessToken: string | null; +} + +export interface UseBlockUnpricedConfigReturn { + blockUnpriced: boolean; + isUpdating: boolean; + fetchBlockUnpriced: () => Promise; + setBlockUnpriced: (enabled: boolean) => Promise; +} + +interface BlockUnpricedResponse { + enabled: boolean; +} + +const ENDPOINT = "/config/block_requests_for_models_without_pricing"; + +export function useBlockUnpricedConfig({ accessToken }: UseBlockUnpricedConfigProps): UseBlockUnpricedConfigReturn { + const [blockUnpriced, setBlockUnpricedState] = useState(false); + const [isUpdating, setIsUpdating] = useState(false); + + const fetchBlockUnpriced = useCallback(async () => { + if (!accessToken) return; + try { + const data = await apiClient.get(ENDPOINT, { accessToken }); + setBlockUnpricedState(Boolean(data?.enabled)); + } catch (error) { + console.error("Error fetching block-unpriced-models setting:", error); + toast.fromError(error); + } + }, [accessToken]); + + const setBlockUnpriced = useCallback( + async (enabled: boolean) => { + if (!accessToken) return; + setIsUpdating(true); + try { + const data = await apiClient.patch(ENDPOINT, { accessToken, body: { enabled } }); + setBlockUnpricedState(Boolean(data?.enabled)); + toast.success( + enabled + ? "Requests for models without pricing will now be blocked" + : "Requests for models without pricing are now allowed", + ); + } catch (error) { + console.error("Error updating block-unpriced-models setting:", error); + toast.fromError(error); + } finally { + setIsUpdating(false); + } + }, + [accessToken], + ); + + return { + blockUnpriced, + isUpdating, + fetchBlockUnpriced, + setBlockUnpriced, + }; +} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/projects/useCreateProject.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/projects/useCreateProject.test.ts index 110a704725a..2fdb999e6cf 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/projects/useCreateProject.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/projects/useCreateProject.test.ts @@ -67,7 +67,12 @@ describe("useCreateProject", () => { const { result } = renderHook(() => useCreateProject(), { wrapper: makeWrapper(queryClient), }); - const params: ProjectCreateParams = { team_id: "team-1", project_alias: "New Project" }; + const params: ProjectCreateParams = { + team_id: "team-1", + project_alias: "New Project", + model_itpm_limit: { "gpt-4": 150 }, + model_otpm_limit: { "gpt-4": 250 }, + }; const data = await result.current.mutateAsync(params); expect(data).toEqual(mockProject); const [url, init] = (global.fetch as any).mock.calls[0]; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/projects/useCreateProject.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/projects/useCreateProject.ts index 2e67e626936..d1d16f08867 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/projects/useCreateProject.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/projects/useCreateProject.ts @@ -16,6 +16,8 @@ export interface ProjectCreateParams { metadata?: Record; model_rpm_limit?: Record; model_tpm_limit?: Record; + model_itpm_limit?: Record; + model_otpm_limit?: Record; } // ── Fetch function ─────────────────────────────────────────────────────────── diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/projects/useUpdateProject.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/projects/useUpdateProject.test.ts index 9e752ac098a..bf3add2d8c6 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/projects/useUpdateProject.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/projects/useUpdateProject.test.ts @@ -68,17 +68,25 @@ describe("useUpdateProject", () => { const { result } = renderHook(() => useUpdateProject(), { wrapper: makeWrapper(queryClient), }); + const params = { + project_alias: "Updated Name", + model_itpm_limit: { "gpt-4": 150 }, + model_otpm_limit: { "gpt-4": 250 }, + }; + const expectedBody = { + project_id: "proj-1", + project_alias: "Updated Name", + model_itpm_limit: { "gpt-4": 150 }, + model_otpm_limit: { "gpt-4": 250 }, + }; const data = await result.current.mutateAsync({ projectId: "proj-1", - params: { project_alias: "Updated Name" }, + params, }); expect(data).toEqual(updated); const [url, init] = (global.fetch as any).mock.calls[0]; expect(url).toContain("/project/update"); - expect(JSON.parse(init.body)).toMatchObject({ - project_id: "proj-1", - project_alias: "Updated Name", - }); + expect(JSON.parse(init.body)).toMatchObject(expectedBody); }); it("should invalidate project queries on success", async () => { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/projects/useUpdateProject.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/projects/useUpdateProject.ts index 6d8c2d9d4f8..8e6bad04a28 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/projects/useUpdateProject.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/projects/useUpdateProject.ts @@ -16,6 +16,8 @@ export interface ProjectUpdateParams { metadata?: Record; model_rpm_limit?: Record; model_tpm_limit?: Record; + model_itpm_limit?: Record; + model_otpm_limit?: Record; } // ── Fetch function ─────────────────────────────────────────────────────────── diff --git a/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/CreateProjectModal.integration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/CreateProjectModal.integration.test.tsx index c7d0d00057c..883e3173d98 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/CreateProjectModal.integration.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/CreateProjectModal.integration.test.tsx @@ -193,11 +193,15 @@ describe("CreateProjectModal submit payload", () => { fireEvent.change(screen.getByPlaceholderText("Model name (e.g. gpt-4)"), { target: { value: "gpt-4" } }); fireEvent.change(screen.getByPlaceholderText("TPM Limit"), { target: { value: "100" } }); fireEvent.change(screen.getByPlaceholderText("RPM Limit"), { target: { value: "20" } }); + fireEvent.change(screen.getByPlaceholderText("Input TPM Limit"), { target: { value: "60" } }); + fireEvent.change(screen.getByPlaceholderText("Output TPM Limit"), { target: { value: "40" } }); await submit(user); await waitFor(() => expect(mutate).toHaveBeenCalled()); expect(params().model_tpm_limit).toStrictEqual({ "gpt-4": 100 }); expect(params().model_rpm_limit).toStrictEqual({ "gpt-4": 20 }); + expect(params().model_itpm_limit).toStrictEqual({ "gpt-4": 60 }); + expect(params().model_otpm_limit).toStrictEqual({ "gpt-4": 40 }); }); it("sends metadata pairs as an object", async () => { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/CreateProjectModal.tsx b/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/CreateProjectModal.tsx index 14bcc40bfea..c923af02c4a 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/CreateProjectModal.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/CreateProjectModal.tsx @@ -10,7 +10,7 @@ import { UiLoadingSpinner } from "@/components/ui/ui-loading-spinner"; import { useCreateProject, ProjectCreateParams } from "@/app/(dashboard)/hooks/projects/useCreateProject"; import { ProjectBaseForm } from "./ProjectBaseForm"; import { emptyProjectFormValues, projectFormSchema } from "./projectFormSchema"; -import { buildProjectApiParams } from "./projectFormUtils"; +import { buildProjectCreateParams } from "./projectFormUtils"; import { Dialog, DialogContent, DialogHeader, DialogTitle } from "@/components/ui/dialog"; interface CreateProjectModalProps { @@ -25,7 +25,7 @@ function CreateProjectForm({ onClose }: { onClose: () => void }) { const handleSubmit = form.handleSubmit((values) => { const params: ProjectCreateParams = { - ...buildProjectApiParams(values), + ...buildProjectCreateParams(values), team_id: values.team_id, }; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/EditProjectModal.integration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/EditProjectModal.integration.test.tsx index 9abea27ceda..5b84aa15dd1 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/EditProjectModal.integration.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/EditProjectModal.integration.test.tsx @@ -52,6 +52,8 @@ const project: ProjectResponse = { guardrails: ["pii-guard"], model_rpm_limit: { "gpt-4": 20 }, model_tpm_limit: { "gpt-4": 100 }, + model_itpm_limit: { "gpt-4": 60 }, + model_otpm_limit: { "gpt-4": 40 }, }, models: ["gpt-4"], spend: 10, @@ -120,6 +122,8 @@ describe("EditProjectModal submit payload", () => { guardrails: ["pii-guard"], model_rpm_limit: { "gpt-4": 20 }, model_tpm_limit: { "gpt-4": 100 }, + model_itpm_limit: { "gpt-4": 60 }, + model_otpm_limit: { "gpt-4": 40 }, metadata: { owner: "platform" }, team_id: "team-1", }); @@ -200,6 +204,77 @@ describe("EditProjectModal submit payload", () => { expect(variables().params).not.toHaveProperty("guardrails"); expect(variables().params).not.toHaveProperty("model_rpm_limit"); expect(variables().params).not.toHaveProperty("model_tpm_limit"); + expect(variables().params).not.toHaveProperty("model_itpm_limit"); + expect(variables().params).not.toHaveProperty("model_otpm_limit"); expect(variables().params).not.toHaveProperty("metadata"); }); + + it("sends empty limit maps once the model limit row is removed, so the stored limits are cleared", async () => { + const user = setup(); + renderModal(); + await screen.findByDisplayValue("My Project"); + + await user.click(screen.getByText("Advanced Settings")); + await screen.findByText("Model-Specific Limits"); + await user.click(screen.getByRole("button", { name: "Remove model limit 1" })); + await save(user); + + await waitFor(() => expect(mutate).toHaveBeenCalled()); + expect(variables().params.model_itpm_limit).toStrictEqual({}); + expect(variables().params.model_otpm_limit).toStrictEqual({}); + expect(variables().params.model_tpm_limit).toStrictEqual({}); + expect(variables().params.model_rpm_limit).toStrictEqual({}); + expect(variables().params.metadata).toStrictEqual({ owner: "platform" }); + }); + + it("sends an empty input TPM map when only that field is blanked on a row that keeps its other limits", async () => { + const user = setup(); + renderModal(); + await screen.findByDisplayValue("My Project"); + + await user.click(screen.getByText("Advanced Settings")); + await screen.findByText("Model-Specific Limits"); + await user.clear(screen.getByLabelText("Input TPM Limit")); + await save(user); + + await waitFor(() => expect(mutate).toHaveBeenCalled()); + expect(variables().params.model_itpm_limit).toStrictEqual({}); + expect(variables().params.model_otpm_limit).toStrictEqual({ "gpt-4": 40 }); + expect(variables().params.model_tpm_limit).toStrictEqual({ "gpt-4": 100 }); + }); + + it("sends an empty metadata object once the last metadata row is removed", async () => { + const user = setup(); + renderModal({ ...project, metadata: { owner: "platform" } } as unknown as ProjectResponse); + await screen.findByDisplayValue("My Project"); + + await user.click(screen.getByText("Advanced Settings")); + await screen.findByText("Metadata"); + await user.click(screen.getByRole("button", { name: "Remove metadata pair 1" })); + await save(user); + + await waitFor(() => expect(mutate).toHaveBeenCalled()); + expect(variables().params.metadata).toStrictEqual({}); + }); + + it("round-trips input and output-only model limits from project metadata", async () => { + const user = setup(); + renderModal({ + ...project, + metadata: { + model_itpm_limit: { "input-model": 150 }, + model_otpm_limit: { "output-model": 250 }, + }, + } as unknown as ProjectResponse); + await screen.findByDisplayValue("My Project"); + + await user.click(screen.getByText("Advanced Settings")); + await screen.findByText("Model-Specific Limits"); + await save(user); + + await waitFor(() => expect(mutate).toHaveBeenCalled()); + expect(variables().params.model_itpm_limit).toStrictEqual({ "input-model": 150 }); + expect(variables().params.model_otpm_limit).toStrictEqual({ "output-model": 250 }); + expect(variables().params.metadata).toStrictEqual({}); + }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/EditProjectModal.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/EditProjectModal.test.tsx index c8de4644d91..9cf6ed24453 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/EditProjectModal.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/EditProjectModal.test.tsx @@ -1,7 +1,7 @@ import { describe, it, expect, vi, beforeEach } from "vitest"; import userEvent from "@testing-library/user-event"; import { renderWithProviders, screen } from "../../../../../../tests/test-utils"; -import { EditProjectModal } from "./EditProjectModal"; +import { EditProjectModal, toFormValues } from "./EditProjectModal"; import { ProjectResponse } from "@/app/(dashboard)/hooks/projects/useProjects"; const mockMutate = vi.fn(); @@ -72,4 +72,21 @@ describe("EditProjectModal", () => { renderWithProviders(); expect(screen.getByTestId("project-base-form")).toBeInTheDocument(); }); + + it("should prefill input and output TPM limits and keep them out of metadata", () => { + const values = toFormValues({ + ...mockProject, + metadata: { + model_itpm_limit: { "input-model": 150 }, + model_otpm_limit: { "output-model": 250 }, + owner: "platform", + }, + }); + + expect(values.modelLimits).toEqual([ + { model: "input-model", rpm: undefined, tpm: undefined, itpm: 150, otpm: undefined }, + { model: "output-model", rpm: undefined, tpm: undefined, itpm: undefined, otpm: 250 }, + ]); + expect(values.metadata).toEqual([{ key: "owner", value: "platform" }]); + }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/EditProjectModal.tsx b/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/EditProjectModal.tsx index 307ba88e0f3..5c1a443d6c5 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/EditProjectModal.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/EditProjectModal.tsx @@ -11,7 +11,7 @@ import { ProjectResponse } from "@/app/(dashboard)/hooks/projects/useProjects"; import { useUpdateProject, ProjectUpdateParams } from "@/app/(dashboard)/hooks/projects/useUpdateProject"; import { ProjectBaseForm } from "./ProjectBaseForm"; import { projectFormSchema, type ProjectFormValues } from "./projectFormSchema"; -import { buildProjectApiParams } from "./projectFormUtils"; +import { buildProjectUpdateParams } from "./projectFormUtils"; import { Dialog, DialogContent, DialogHeader, DialogTitle } from "@/components/ui/dialog"; interface EditProjectModalProps { @@ -21,18 +21,35 @@ interface EditProjectModalProps { onSuccess?: () => void; } -const INTERNAL_METADATA_KEYS = new Set(["model_rpm_limit", "model_tpm_limit", "guardrails"]); +const INTERNAL_METADATA_KEYS = new Set([ + "model_rpm_limit", + "model_tpm_limit", + "model_itpm_limit", + "model_otpm_limit", + "guardrails", +]); -const toFormValues = (project: ProjectResponse): ProjectFormValues => { +export const toFormValues = (project: ProjectResponse): ProjectFormValues => { const metadataObj = (project.metadata ?? {}) as Record; const rpmLimits = (metadataObj.model_rpm_limit ?? {}) as Record; const tpmLimits = (metadataObj.model_tpm_limit ?? {}) as Record; + const itpmLimits = (metadataObj.model_itpm_limit ?? {}) as Record; + const otpmLimits = (metadataObj.model_otpm_limit ?? {}) as Record; const guardrails = (Array.isArray(metadataObj.guardrails) ? metadataObj.guardrails : []) as string[]; - const modelLimits = Array.from(new Set([...Object.keys(rpmLimits), ...Object.keys(tpmLimits)])).map((model) => ({ + const modelLimits = Array.from( + new Set([ + ...Object.keys(rpmLimits), + ...Object.keys(tpmLimits), + ...Object.keys(itpmLimits), + ...Object.keys(otpmLimits), + ]), + ).map((model) => ({ model, rpm: rpmLimits[model], tpm: tpmLimits[model], + itpm: itpmLimits[model], + otpm: otpmLimits[model], })); const metadata = Object.entries(metadataObj) @@ -69,7 +86,7 @@ function EditProjectForm({ project, onClose, onSuccess }: Omit { expect(screen.getByText("Guardrails")).toBeInTheDocument(); }); }); + + it("should show combined, input, and output TPM limit inputs for a model row", async () => { + const user = userEvent.setup(); + renderWithProviders(); + await user.click(screen.getByText("Advanced Settings")); + await user.click(screen.getByRole("button", { name: /add model limit/i })); + + expect(screen.getByPlaceholderText("TPM Limit")).toBeInTheDocument(); + expect(screen.getByPlaceholderText("Input TPM Limit")).toBeInTheDocument(); + expect(screen.getByPlaceholderText("Output TPM Limit")).toBeInTheDocument(); + expect(screen.getByLabelText("TPM Limit")).toBeInTheDocument(); + expect(screen.getByLabelText("Input TPM Limit")).toBeInTheDocument(); + expect(screen.getByLabelText("Output TPM Limit")).toBeInTheDocument(); + }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/ProjectBaseForm.tsx b/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/ProjectBaseForm.tsx index d4b7e71dfe9..a18b5ecfb98 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/ProjectBaseForm.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/ProjectBaseForm.tsx @@ -49,6 +49,13 @@ export function ProjectBaseForm({ form, advancedOpen, onAdvancedOpenChange }: Pr const modelLimits = useFieldArray({ control: form.control, name: "modelLimits" }); const metadata = useFieldArray({ control: form.control, name: "metadata" }); + const emptyModelLimit: NonNullable[number] = { + model: "", + tpm: undefined, + rpm: undefined, + itpm: undefined, + otpm: undefined, + }; const teamIdValue = useWatch({ control: form.control, name: "team_id" }); const isBlocked = useWatch({ control: form.control, name: "isBlocked" }); @@ -262,13 +269,16 @@ export function ProjectBaseForm({ form, advancedOpen, onAdvancedOpenChange }: Pr

Model-Specific Limits

{modelLimits.fields.map((field, index) => ( -
- +
+ {({ ref, ...control }) => ( )} - + {({ ref, value, onChange, ...control }) => ( )} - + {({ ref, value, onChange, ...control }) => ( )} + + {({ ref, value, onChange, ...control }) => ( + onChange(toOptionalNumber(event.target.value))} + /> + )} + + + {({ ref, value, onChange, ...control }) => ( + onChange(toOptionalNumber(event.target.value))} + /> + )} +