From a1514efa210c60c00809b21d2906503b0c452cc8 Mon Sep 17 00:00:00 2001 From: michelligabriele Date: Mon, 27 Jul 2026 17:07:58 +0200 Subject: [PATCH 01/48] fix(vector_stores): S3 Vectors search router bypass + rag query config drop + UI error swallow --- .../azure_ai/vector_stores/transformation.py | 2 + .../base_llm/vector_store/transformation.py | 4 + .../bedrock/vector_stores/transformation.py | 2 + litellm/llms/custom_httpx/llm_http_handler.py | 7 + .../gemini/vector_stores/transformation.py | 2 + .../milvus/vector_stores/transformation.py | 2 + .../openai/vector_stores/transformation.py | 2 + .../pg_vector/vector_stores/transformation.py | 2 + .../ragflow/vector_stores/transformation.py | 2 + .../vector_stores/transformation.py | 32 ++- .../vector_stores/rag_api/transformation.py | 2 + .../search_api/transformation.py | 2 + litellm/proxy/rag_endpoints/endpoints.py | 14 ++ litellm/rag/main.py | 14 +- litellm/router.py | 8 + litellm/vector_stores/main.py | 9 +- .../test_s3_vectors_transformation.py | 189 +++++++++++++++++- .../proxy/rag_endpoints/test_rag_endpoints.py | 103 ++++++++++ tests/test_litellm/rag/test_main.py | 90 +++++++++ tests/test_litellm/test_router.py | 55 +++++ tests/test_litellm/vector_stores/test_main.py | 77 +++++++ .../_components/VectorStoreTester.test.tsx | 25 ++- .../_components/VectorStoreTester.tsx | 8 +- .../src/components/networking.tsx | 2 +- 24 files changed, 628 insertions(+), 27 deletions(-) create mode 100644 tests/test_litellm/vector_stores/test_main.py diff --git a/litellm/llms/azure_ai/vector_stores/transformation.py b/litellm/llms/azure_ai/vector_stores/transformation.py index da6a4a93cd8..bd3eaeee989 100644 --- a/litellm/llms/azure_ai/vector_stores/transformation.py +++ b/litellm/llms/azure_ai/vector_stores/transformation.py @@ -19,6 +19,7 @@ from litellm.types.vector_stores import ( if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj + from litellm.router import Router LiteLLMLoggingObj = _LiteLLMLoggingObj else: @@ -92,6 +93,7 @@ class AzureAIVectorStoreConfig(BaseVectorStoreConfig, BaseAzureLLM): litellm_logging_obj: LiteLLMLoggingObj, litellm_params: dict, extra_body: Optional[Dict[str, Any]] = None, + router: Optional["Router"] = None, ) -> Tuple[str, Dict[str, Any]]: """ Transform search request for Azure AI Search API diff --git a/litellm/llms/base_llm/vector_store/transformation.py b/litellm/llms/base_llm/vector_store/transformation.py index b222e3dd160..9a0e401b527 100644 --- a/litellm/llms/base_llm/vector_store/transformation.py +++ b/litellm/llms/base_llm/vector_store/transformation.py @@ -16,6 +16,7 @@ from litellm.types.vector_stores import ( if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj + from litellm.router import Router from ..chat.transformation import BaseLLMException as _BaseLLMException @@ -56,6 +57,7 @@ class BaseVectorStoreConfig: litellm_logging_obj: LiteLLMLoggingObj, litellm_params: dict, extra_body: Optional[Dict[str, Any]] = None, + router: Optional["Router"] = None, ) -> Tuple[str, Dict]: pass @@ -68,6 +70,7 @@ class BaseVectorStoreConfig: litellm_logging_obj: LiteLLMLoggingObj, litellm_params: dict, extra_body: Optional[Dict[str, Any]] = None, + router: Optional["Router"] = None, ) -> Tuple[str, Dict]: """ Optional async version of transform_search_vector_store_request. @@ -83,6 +86,7 @@ class BaseVectorStoreConfig: litellm_logging_obj=litellm_logging_obj, litellm_params=litellm_params, extra_body=extra_body, + router=router, ) @abstractmethod diff --git a/litellm/llms/bedrock/vector_stores/transformation.py b/litellm/llms/bedrock/vector_stores/transformation.py index c1b124caec1..7a6a0eb6d84 100644 --- a/litellm/llms/bedrock/vector_stores/transformation.py +++ b/litellm/llms/bedrock/vector_stores/transformation.py @@ -27,6 +27,7 @@ from litellm.types.vector_stores import ( if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.router import Router else: LiteLLMLoggingObj = Any @@ -196,6 +197,7 @@ class BedrockVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM): litellm_logging_obj: LiteLLMLoggingObj, litellm_params: dict, extra_body: Optional[Dict[str, Any]] = None, + router: Optional["Router"] = None, ) -> Tuple[str, Dict]: if isinstance(query, list): query = " ".join(query) diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index ec1301e5923..ec701fbe87e 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -167,6 +167,7 @@ if TYPE_CHECKING: AnthropicMessagesStreamingResponse, ) from litellm.llms.base_llm.passthrough.transformation import BasePassthroughConfig + from litellm.router import Router from litellm.types.llms.openai_evals import ( CancelEvalResponse, CancelRunResponse, @@ -9409,6 +9410,7 @@ class BaseLLMHTTPHandler: timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, _is_async: bool = False, + router: Optional["Router"] = None, ) -> VectorStoreSearchResponse: if client is None or not isinstance(client, AsyncHTTPHandler): async_httpx_client = get_async_httpx_client( @@ -9443,6 +9445,7 @@ class BaseLLMHTTPHandler: litellm_logging_obj=logging_obj, litellm_params=dict(litellm_params), extra_body=extra_body, + router=router, ) else: ( @@ -9456,6 +9459,7 @@ class BaseLLMHTTPHandler: litellm_logging_obj=logging_obj, litellm_params=dict(litellm_params), extra_body=extra_body, + router=router, ) all_optional_params: Dict[str, Any] = dict(litellm_params) all_optional_params.update(vector_store_search_optional_params or {}) @@ -9507,6 +9511,7 @@ class BaseLLMHTTPHandler: timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, _is_async: bool = False, + router: Optional["Router"] = None, ) -> Union[VectorStoreSearchResponse, Coroutine[Any, Any, VectorStoreSearchResponse]]: if _is_async: return self.async_vector_store_search_handler( @@ -9521,6 +9526,7 @@ class BaseLLMHTTPHandler: extra_body=extra_body, timeout=timeout, client=client, + router=router, ) if client is None or not isinstance(client, HTTPHandler): @@ -9551,6 +9557,7 @@ class BaseLLMHTTPHandler: litellm_logging_obj=logging_obj, litellm_params=dict(litellm_params), extra_body=extra_body, + router=router, ) all_optional_params: Dict[str, Any] = dict(litellm_params) diff --git a/litellm/llms/gemini/vector_stores/transformation.py b/litellm/llms/gemini/vector_stores/transformation.py index f98cb0e5b0c..5aba5752a44 100644 --- a/litellm/llms/gemini/vector_stores/transformation.py +++ b/litellm/llms/gemini/vector_stores/transformation.py @@ -31,6 +31,7 @@ from litellm.types.vector_stores import ( if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.router import Router else: LiteLLMLoggingObj = Any @@ -111,6 +112,7 @@ class GeminiVectorStoreConfig(BaseVectorStoreConfig): litellm_logging_obj: LiteLLMLoggingObj, litellm_params: dict, extra_body: Optional[Dict[str, Any]] = None, + router: Optional["Router"] = None, ) -> Tuple[str, Dict]: """ Transform search request to Gemini's generateContent format. diff --git a/litellm/llms/milvus/vector_stores/transformation.py b/litellm/llms/milvus/vector_stores/transformation.py index a53075ba1d6..589063cd188 100644 --- a/litellm/llms/milvus/vector_stores/transformation.py +++ b/litellm/llms/milvus/vector_stores/transformation.py @@ -19,6 +19,7 @@ from litellm.types.vector_stores import ( if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj + from litellm.router import Router LiteLLMLoggingObj = _LiteLLMLoggingObj else: @@ -123,6 +124,7 @@ class MilvusVectorStoreConfig(BaseVectorStoreConfig): litellm_logging_obj: LiteLLMLoggingObj, litellm_params: dict, extra_body: Optional[Dict[str, Any]] = None, + router: Optional["Router"] = None, ) -> Tuple[str, Dict[str, Any]]: """ Transform search request for Azure AI Search API diff --git a/litellm/llms/openai/vector_stores/transformation.py b/litellm/llms/openai/vector_stores/transformation.py index 6ccf8e271e5..9ab1568a375 100644 --- a/litellm/llms/openai/vector_stores/transformation.py +++ b/litellm/llms/openai/vector_stores/transformation.py @@ -21,6 +21,7 @@ from litellm.utils import add_openai_metadata if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj + from litellm.router import Router LiteLLMLoggingObj = _LiteLLMLoggingObj else: @@ -99,6 +100,7 @@ class OpenAIVectorStoreConfig(BaseVectorStoreConfig): litellm_logging_obj: LiteLLMLoggingObj, litellm_params: dict, extra_body: Optional[Dict[str, Any]] = None, + router: Optional["Router"] = None, ) -> Tuple[str, Dict]: encoded_vector_store_id = encode_url_path_segment(vector_store_id, field_name="vector_store_id") url = f"{api_base}/{encoded_vector_store_id}/search" diff --git a/litellm/llms/pg_vector/vector_stores/transformation.py b/litellm/llms/pg_vector/vector_stores/transformation.py index b58b6e7f498..116f79c834f 100644 --- a/litellm/llms/pg_vector/vector_stores/transformation.py +++ b/litellm/llms/pg_vector/vector_stores/transformation.py @@ -8,6 +8,7 @@ from litellm.types.vector_stores import VectorStoreSearchOptionalRequestParams if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.router import Router else: LiteLLMLoggingObj = Any @@ -80,6 +81,7 @@ class PGVectorStoreConfig(OpenAIVectorStoreConfig): litellm_logging_obj: LiteLLMLoggingObj, litellm_params: dict, extra_body: Optional[Dict[str, Any]] = None, + router: Optional["Router"] = None, ) -> Tuple[str, Dict]: encoded_vector_store_id = encode_url_path_segment(vector_store_id, field_name="vector_store_id") url = f"{api_base}/{encoded_vector_store_id}/search" diff --git a/litellm/llms/ragflow/vector_stores/transformation.py b/litellm/llms/ragflow/vector_stores/transformation.py index d8bdd981425..332ed7f0c6b 100644 --- a/litellm/llms/ragflow/vector_stores/transformation.py +++ b/litellm/llms/ragflow/vector_stores/transformation.py @@ -17,6 +17,7 @@ from litellm.types.vector_stores import ( if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.router import Router else: LiteLLMLoggingObj = Any @@ -92,6 +93,7 @@ class RAGFlowVectorStoreConfig(BaseVectorStoreConfig): litellm_logging_obj: LiteLLMLoggingObj, litellm_params: dict, extra_body: Optional[Dict[str, Any]] = None, + router: Optional["Router"] = None, ) -> Tuple[str, Dict]: """RAGFlow vector stores are management-only, search is not supported.""" raise NotImplementedError("RAGFlow vector stores support dataset management only, not search/retrieval") diff --git a/litellm/llms/s3_vectors/vector_stores/transformation.py b/litellm/llms/s3_vectors/vector_stores/transformation.py index b31e6f4511a..a999db21dbe 100644 --- a/litellm/llms/s3_vectors/vector_stores/transformation.py +++ b/litellm/llms/s3_vectors/vector_stores/transformation.py @@ -1,8 +1,8 @@ -import re from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union import httpx +from litellm.caching._embedding_router import resolve_embedding_router from litellm.llms.base_llm.vector_store.transformation import BaseVectorStoreConfig from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM from litellm.types.router import GenericLiteLLMParams @@ -18,6 +18,7 @@ from litellm.types.vector_stores import ( if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.router import Router else: LiteLLMLoggingObj = Any @@ -58,13 +59,18 @@ class S3VectorsVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM): return headers def get_complete_url(self, api_base: Optional[str], litellm_params: dict) -> str: - aws_region_name = litellm_params.get("aws_region_name") - if not aws_region_name: - raise ValueError("aws_region_name is required for S3 Vectors") - if not re.match(r"^[a-z][a-z0-9-]*$", aws_region_name): - raise ValueError("Invalid aws_region_name format") + # Resolve region the same way the ingestion path does: + # dynamic param -> AWS_REGION_NAME -> AWS_REGION -> default (us-west-2) + aws_region_name = self.get_aws_region_name_for_non_llm_api_calls(litellm_params.get("aws_region_name")) return f"https://s3vectors.{aws_region_name}.api.aws" + def _resolve_query_embedding_router(self, embedding_model: str, router: Optional["Router"]) -> Optional["Router"]: + """Return the router iff it serves ``embedding_model`` as a deployment.""" + if router is None: + return None + model_list = [dict(m) for m in (router.get_model_list() or [])] + return resolve_embedding_router(embedding_model=embedding_model, llm_router=router, llm_model_list=model_list) + def transform_search_vector_store_request( self, vector_store_id: str, @@ -74,6 +80,7 @@ class S3VectorsVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM): litellm_logging_obj: LiteLLMLoggingObj, litellm_params: dict, extra_body: Optional[Dict[str, Any]] = None, + router: Optional["Router"] = None, ) -> Tuple[str, Dict]: """Sync version - generates embedding synchronously.""" # For S3 Vectors, vector_store_id should be in format: bucket_name:index_name @@ -99,10 +106,14 @@ class S3VectorsVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM): # Generate embedding for the query embedding_model = litellm_params.get("embedding_model", "text-embedding-3-small") + embedding_router = self._resolve_query_embedding_router(embedding_model=embedding_model, router=router) import litellm as litellm_module - embedding_response = litellm_module.embedding(model=embedding_model, input=[query]) + if embedding_router is not None: + embedding_response = embedding_router.embedding(model=embedding_model, input=[query]) + else: + embedding_response = litellm_module.embedding(model=embedding_model, input=[query]) query_embedding = embedding_response.data[0]["embedding"] url = f"{api_base}/QueryVectors" @@ -128,6 +139,7 @@ class S3VectorsVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM): litellm_logging_obj: LiteLLMLoggingObj, litellm_params: dict, extra_body: Optional[Dict[str, Any]] = None, + router: Optional["Router"] = None, ) -> Tuple[str, Dict]: """Async version - generates embedding asynchronously.""" # For S3 Vectors, vector_store_id should be in format: bucket_name:index_name @@ -153,10 +165,14 @@ class S3VectorsVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM): # Generate embedding for the query asynchronously embedding_model = litellm_params.get("embedding_model", "text-embedding-3-small") + embedding_router = self._resolve_query_embedding_router(embedding_model=embedding_model, router=router) import litellm as litellm_module - embedding_response = await litellm_module.aembedding(model=embedding_model, input=[query]) + if embedding_router is not None: + embedding_response = await embedding_router.aembedding(model=embedding_model, input=[query]) + else: + embedding_response = await litellm_module.aembedding(model=embedding_model, input=[query]) query_embedding = embedding_response.data[0]["embedding"] url = f"{api_base}/QueryVectors" diff --git a/litellm/llms/vertex_ai/vector_stores/rag_api/transformation.py b/litellm/llms/vertex_ai/vector_stores/rag_api/transformation.py index 47a81fc07bf..93ad40616b5 100644 --- a/litellm/llms/vertex_ai/vector_stores/rag_api/transformation.py +++ b/litellm/llms/vertex_ai/vector_stores/rag_api/transformation.py @@ -19,6 +19,7 @@ from litellm.types.vector_stores import ( if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj + from litellm.router import Router LiteLLMLoggingObj = _LiteLLMLoggingObj else: @@ -97,6 +98,7 @@ class VertexVectorStoreConfig(BaseVectorStoreConfig, VertexBase): litellm_logging_obj: LiteLLMLoggingObj, litellm_params: dict, extra_body: Optional[Dict[str, Any]] = None, + router: Optional["Router"] = None, ) -> Tuple[str, Dict[str, Any]]: """ Transform search request for Vertex AI RAG API diff --git a/litellm/llms/vertex_ai/vector_stores/search_api/transformation.py b/litellm/llms/vertex_ai/vector_stores/search_api/transformation.py index 958839d4a48..f6f9e34dc75 100644 --- a/litellm/llms/vertex_ai/vector_stores/search_api/transformation.py +++ b/litellm/llms/vertex_ai/vector_stores/search_api/transformation.py @@ -23,6 +23,7 @@ from litellm.types.vector_stores import ( if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj + from litellm.router import Router LiteLLMLoggingObj = _LiteLLMLoggingObj else: @@ -197,6 +198,7 @@ class VertexSearchAPIVectorStoreConfig(BaseVectorStoreConfig, VertexBase): litellm_logging_obj: LiteLLMLoggingObj, litellm_params: dict, extra_body: Optional[Dict[str, Any]] = None, + router: Optional["Router"] = None, ) -> Tuple[str, Dict[str, Any]]: """ Transform a search request for the Vertex AI Search (Discovery Engine) API. diff --git a/litellm/proxy/rag_endpoints/endpoints.py b/litellm/proxy/rag_endpoints/endpoints.py index 27ffc49901b..0d93f20373c 100644 --- a/litellm/proxy/rag_endpoints/endpoints.py +++ b/litellm/proxy/rag_endpoints/endpoints.py @@ -26,6 +26,9 @@ from litellm.proxy.common_utils.http_parsing_utils import ( _safe_get_request_headers, get_form_data, ) +from litellm.proxy.vector_store_endpoints.endpoints import ( + _update_request_data_with_litellm_managed_vector_store_registry, +) from litellm.proxy.vector_store_endpoints.utils import ( assert_user_can_access_vector_store_id, ) @@ -652,6 +655,17 @@ async def rag_query( user_api_key_dict=user_api_key_dict, ) + # Merge litellm-managed vector store params (provider, region, embedding + # model, credentials, ...) from the registry — same source the direct + # /vector_stores/{id}/search endpoint uses. User-supplied + # retrieval_config keys win on conflict. + store_data = await _update_request_data_with_litellm_managed_vector_store_registry( + data={}, + vector_store_id=retrieval_config["vector_store_id"], + user_api_key_dict=user_api_key_dict, + ) + retrieval_config = {**store_data, **retrieval_config} + # Add litellm data request_data: Dict[str, Any] = {} request_data = await add_litellm_data_to_request( diff --git a/litellm/rag/main.py b/litellm/rag/main.py index 29891ccfd24..2329a820f1f 100644 --- a/litellm/rag/main.py +++ b/litellm/rag/main.py @@ -59,6 +59,14 @@ INGESTION_REGISTRY: Dict[str, Type[BaseRAGIngestion]] = { "vertex_ai": VertexAIRAGIngestion, } +# retrieval_config keys consumed by the query pipeline itself; everything else is +# forwarded to vector_stores.asearch as provider-specific params (e.g. +# aws_region_name, embedding_model, vector_bucket_name for S3 Vectors). +# `filters`/`retrieval_filter` are reserved for the explicit filter param. +_CONSUMED_RETRIEVAL_CONFIG_KEYS = frozenset( + {"vector_store_id", "custom_llm_provider", "top_k", "filters", "retrieval_filter"} +) + def get_ingestion_class(provider: str) -> Type[BaseRAGIngestion]: """ @@ -233,13 +241,17 @@ async def _execute_query_pipeline( raise ValueError("No query found in messages for RAG query") # 2. Search vector store + # Forward provider-specific retrieval_config extras (region, embedding model, + # bucket, credentials refs, ...) to the search call; kwargs win on conflict. + provider_search_params = {k: v for k, v in retrieval_config.items() if k not in _CONSUMED_RETRIEVAL_CONFIG_KEYS} with _suppressed_sub_call_billing(): search_response = await litellm.vector_stores.asearch( vector_store_id=retrieval_config["vector_store_id"], query=query_text, max_num_results=retrieval_config.get("top_k", 10), custom_llm_provider=retrieval_config.get("custom_llm_provider", "openai"), - **kwargs, + router=router, + **{**provider_search_params, **kwargs}, ) search_provider = retrieval_config.get("custom_llm_provider", "openai") diff --git a/litellm/router.py b/litellm/router.py index 78fe3ff025e..bffd1df3814 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -5820,6 +5820,7 @@ class Router: return await self._init_vector_store_api_endpoints( original_function=original_function, custom_llm_provider=custom_llm_provider, + call_type=call_type, **kwargs, ) elif call_type in ("afile_delete", "afile_content"): @@ -5860,6 +5861,7 @@ class Router: self, original_function: Callable, custom_llm_provider: Optional[str] = None, + call_type: Optional[str] = None, **kwargs, ): """ @@ -5878,6 +5880,12 @@ class Router: **kwargs, ) + # For search, pass the router so provider transforms can resolve + # router-managed embedding models (e.g. S3 Vectors query embeddings). + # Assigning into kwargs also overrides any client-supplied `router` key. + if call_type == "avector_store_search": + kwargs["router"] = self + # Otherwise, call the original function directly return await original_function(**kwargs) diff --git a/litellm/vector_stores/main.py b/litellm/vector_stores/main.py index f768ee75545..4035125120e 100644 --- a/litellm/vector_stores/main.py +++ b/litellm/vector_stores/main.py @@ -6,7 +6,7 @@ import asyncio import builtins import contextvars from functools import partial -from typing import Any, Coroutine, Dict, List, Optional, Union +from typing import TYPE_CHECKING, Any, Coroutine, Dict, List, Optional, Union import httpx @@ -28,6 +28,9 @@ from litellm.types.vector_stores import ( from litellm.utils import ProviderConfigManager, client from litellm.vector_stores.utils import VectorStoreRequestUtils +if TYPE_CHECKING: + from litellm.router import Router + ####### ENVIRONMENT VARIABLES ################### # Initialize any necessary instances or variables here base_llm_http_handler = BaseLLMHTTPHandler() @@ -279,6 +282,7 @@ async def asearch( timeout: Optional[Union[float, httpx.Timeout]] = None, # LiteLLM specific params, custom_llm_provider: Optional[str] = None, + router: Optional["Router"] = None, **kwargs, ) -> VectorStoreSearchResponse: """ @@ -307,6 +311,7 @@ async def asearch( extra_body=extra_body, timeout=timeout, custom_llm_provider=custom_llm_provider, + router=router, **kwargs, ) @@ -346,6 +351,7 @@ def search( timeout: Optional[Union[float, httpx.Timeout]] = None, # LiteLLM specific params, custom_llm_provider: Optional[str] = None, + router: Optional["Router"] = None, **kwargs, ) -> Union[VectorStoreSearchResponse, Coroutine[Any, Any, VectorStoreSearchResponse]]: """ @@ -449,6 +455,7 @@ def search( timeout=timeout or request_timeout, _is_async=_is_async, client=kwargs.get("client"), + router=router, ) return response diff --git a/tests/test_litellm/llms/s3_vectors/vector_stores/test_s3_vectors_transformation.py b/tests/test_litellm/llms/s3_vectors/vector_stores/test_s3_vectors_transformation.py index 7085e45cdc3..9389476bef4 100644 --- a/tests/test_litellm/llms/s3_vectors/vector_stores/test_s3_vectors_transformation.py +++ b/tests/test_litellm/llms/s3_vectors/vector_stores/test_s3_vectors_transformation.py @@ -1,4 +1,4 @@ -from unittest.mock import MagicMock, Mock +from unittest.mock import AsyncMock, MagicMock, Mock, patch import httpx import pytest @@ -9,6 +9,18 @@ from litellm.llms.s3_vectors.vector_stores.transformation import ( from litellm.types.vector_stores import VectorStoreSearchResponse +def _mock_router(model_names, sync=False): + """Router mock serving the given embedding model names.""" + router = MagicMock() + router.get_model_list.return_value = [{"model_name": name} for name in model_names] + embedding_response = Mock(data=[{"embedding": [0.1, 0.2, 0.3]}]) + if sync: + router.embedding = MagicMock(return_value=embedding_response) + else: + router.aembedding = AsyncMock(return_value=embedding_response) + return router + + class TestS3VectorsVectorStoreConfig: def test_init(self): """Test that S3VectorsVectorStoreConfig initializes correctly""" @@ -28,19 +40,174 @@ class TestS3VectorsVectorStoreConfig: url = config.get_complete_url(None, litellm_params) assert url == "https://s3vectors.us-west-2.api.aws" - def test_get_complete_url_missing_region(self): - """Test that missing region raises error""" + def test_get_complete_url_missing_region(self, monkeypatch): + """Missing region falls back to the default region (parity with ingestion)""" + monkeypatch.delenv("AWS_REGION_NAME", raising=False) + monkeypatch.delenv("AWS_REGION", raising=False) config = S3VectorsVectorStoreConfig() - litellm_params = {} - with pytest.raises(ValueError, match="aws_region_name is required"): - config.get_complete_url(None, litellm_params) + url = config.get_complete_url(None, {}) + assert url == "https://s3vectors.us-west-2.api.aws" + + def test_get_complete_url_uses_env_region(self, monkeypatch): + """Missing region param resolves from AWS_REGION_NAME env var""" + monkeypatch.setenv("AWS_REGION_NAME", "eu-west-1") + monkeypatch.delenv("AWS_REGION", raising=False) + config = S3VectorsVectorStoreConfig() + url = config.get_complete_url(None, {}) + assert url == "https://s3vectors.eu-west-1.api.aws" + + def test_get_complete_url_invalid_region_format(self): + """Invalid region format raises""" + config = S3VectorsVectorStoreConfig() + with pytest.raises(ValueError, match="Invalid AWS region format"): + config.get_complete_url(None, {"aws_region_name": "Bad_Region!"}) - @pytest.mark.skip(reason="Requires embedding API call, tested in integration tests") def test_transform_search_request(self): - """Test search request transformation""" - # This test requires making an actual embedding API call - # It's better tested in integration tests - pass + """Full request-body transformation with a router-injected embedding""" + config = S3VectorsVectorStoreConfig() + mock_logging_obj = Mock() + mock_logging_obj.model_call_details = {} + router = _mock_router(["text-embedding-3-small"], sync=True) + + url, request_body = config.transform_search_vector_store_request( + vector_store_id="test-bucket:test-index", + query="test query", + vector_store_search_optional_params={"max_num_results": 7}, + api_base="https://s3vectors.us-west-2.api.aws", + litellm_logging_obj=mock_logging_obj, + litellm_params={}, + extra_body=None, + router=router, + ) + + assert url == "https://s3vectors.us-west-2.api.aws/QueryVectors" + assert request_body == { + "vectorBucketName": "test-bucket", + "indexName": "test-index", + "queryVector": {"float32": [0.1, 0.2, 0.3]}, + "topK": 7, + "returnDistance": True, + "returnMetadata": True, + } + assert mock_logging_obj.model_call_details["query"] == "test query" + + @pytest.mark.asyncio + async def test_atransform_search_uses_router_for_virtual_model(self): + """Regression: router-served embedding models must resolve via the router, + not a bare litellm.aembedding call (which has no deployment credentials).""" + config = S3VectorsVectorStoreConfig() + mock_logging_obj = Mock() + mock_logging_obj.model_call_details = {} + router = _mock_router(["my-embedding-model"]) + + with patch("litellm.aembedding", new=AsyncMock()) as mock_bare_aembedding: + url, request_body = await config.atransform_search_vector_store_request( + vector_store_id="test-bucket:test-index", + query="test query", + vector_store_search_optional_params={}, + api_base="https://s3vectors.us-west-2.api.aws", + litellm_logging_obj=mock_logging_obj, + litellm_params={"embedding_model": "my-embedding-model"}, + extra_body=None, + router=router, + ) + + router.aembedding.assert_awaited_once_with(model="my-embedding-model", input=["test query"]) + mock_bare_aembedding.assert_not_awaited() + assert request_body["queryVector"]["float32"] == [0.1, 0.2, 0.3] + assert request_body["topK"] == 5 # default + + @pytest.mark.asyncio + async def test_atransform_search_falls_back_when_router_does_not_serve_model(self): + """Router present but embedding_model is not a router deployment -> + bare litellm.aembedding keeps working (provider-prefixed + env creds stores).""" + config = S3VectorsVectorStoreConfig() + mock_logging_obj = Mock() + mock_logging_obj.model_call_details = {} + router = _mock_router(["some-other-model"]) + + mock_bare = AsyncMock(return_value=Mock(data=[{"embedding": [0.4, 0.5]}])) + with patch("litellm.aembedding", new=mock_bare): + _, request_body = await config.atransform_search_vector_store_request( + vector_store_id="test-bucket:test-index", + query="test query", + vector_store_search_optional_params={}, + api_base="https://s3vectors.us-west-2.api.aws", + litellm_logging_obj=mock_logging_obj, + litellm_params={"embedding_model": "azure/text-embedding-3-small"}, + extra_body=None, + router=router, + ) + + mock_bare.assert_awaited_once_with(model="azure/text-embedding-3-small", input=["test query"]) + router.aembedding.assert_not_awaited() + assert request_body["queryVector"]["float32"] == [0.4, 0.5] + + @pytest.mark.asyncio + async def test_atransform_search_without_router_uses_bare_embedding(self): + """Backward compat: no router -> bare litellm.aembedding as before""" + config = S3VectorsVectorStoreConfig() + mock_logging_obj = Mock() + mock_logging_obj.model_call_details = {} + + mock_bare = AsyncMock(return_value=Mock(data=[{"embedding": [0.6, 0.7]}])) + with patch("litellm.aembedding", new=mock_bare): + _, request_body = await config.atransform_search_vector_store_request( + vector_store_id="test-bucket:test-index", + query="test query", + vector_store_search_optional_params={}, + api_base="https://s3vectors.us-west-2.api.aws", + litellm_logging_obj=mock_logging_obj, + litellm_params={}, + extra_body=None, + ) + + mock_bare.assert_awaited_once_with(model="text-embedding-3-small", input=["test query"]) + assert request_body["queryVector"]["float32"] == [0.6, 0.7] + + def test_transform_search_uses_router_for_virtual_model_sync(self): + """Sync twin: router-served embedding model resolves via router.embedding""" + config = S3VectorsVectorStoreConfig() + mock_logging_obj = Mock() + mock_logging_obj.model_call_details = {} + router = _mock_router(["my-embedding-model"], sync=True) + + with patch("litellm.embedding", new=MagicMock()) as mock_bare_embedding: + _, request_body = config.transform_search_vector_store_request( + vector_store_id="test-bucket:test-index", + query="test query", + vector_store_search_optional_params={}, + api_base="https://s3vectors.us-west-2.api.aws", + litellm_logging_obj=mock_logging_obj, + litellm_params={"embedding_model": "my-embedding-model"}, + extra_body=None, + router=router, + ) + + router.embedding.assert_called_once_with(model="my-embedding-model", input=["test query"]) + mock_bare_embedding.assert_not_called() + assert request_body["queryVector"]["float32"] == [0.1, 0.2, 0.3] + + def test_transform_search_without_router_uses_bare_embedding_sync(self): + """Sync twin: no router -> bare litellm.embedding as before""" + config = S3VectorsVectorStoreConfig() + mock_logging_obj = Mock() + mock_logging_obj.model_call_details = {} + + mock_bare = MagicMock(return_value=Mock(data=[{"embedding": [0.8, 0.9]}])) + with patch("litellm.embedding", new=mock_bare): + _, request_body = config.transform_search_vector_store_request( + vector_store_id="test-bucket:test-index", + query="test query", + vector_store_search_optional_params={}, + api_base="https://s3vectors.us-west-2.api.aws", + litellm_logging_obj=mock_logging_obj, + litellm_params={}, + extra_body=None, + ) + + mock_bare.assert_called_once_with(model="text-embedding-3-small", input=["test query"]) + assert request_body["queryVector"]["float32"] == [0.8, 0.9] def test_transform_search_request_invalid_vector_store_id(self): """Test that invalid vector_store_id format raises error""" diff --git a/tests/test_litellm/proxy/rag_endpoints/test_rag_endpoints.py b/tests/test_litellm/proxy/rag_endpoints/test_rag_endpoints.py index 15a117bd6fc..8bd67754952 100644 --- a/tests/test_litellm/proxy/rag_endpoints/test_rag_endpoints.py +++ b/tests/test_litellm/proxy/rag_endpoints/test_rag_endpoints.py @@ -327,3 +327,106 @@ def test_rag_query_stream_returns_event_stream(client_internal_user): assert response.headers.get("content-type", "").startswith("text/event-stream") assert '"object":"chat.completion.chunk"' in response.text assert "data: [DONE]" in response.text + + +def test_rag_query_merges_managed_store_params(client_internal_user): + """ + Regression: /v1/rag/query must consult the managed vector store registry + (like the direct /v1/vector_stores/{id}/search endpoint does) so that + provider, region, embedding model, etc. don't have to be repeated in + retrieval_config. Pre-fix the registry was never read, so managed S3 + Vectors stores failed with "aws_region_name is required". + """ + import litellm + from litellm.types.utils import ModelResponse + + mock_vector_store = { + "vector_store_id": "s3-store", + "custom_llm_provider": "s3_vectors", + "litellm_params": { + "aws_region_name": "eu-west-1", + "embedding_model": "my-embed", + "vector_bucket_name": "bkt", + }, + } + mock_registry = MagicMock() + mock_registry.get_litellm_managed_vector_store_from_registry.return_value = mock_vector_store + + mock_response = ModelResponse( + id="chatcmpl-test", + choices=[{"index": 0, "message": {"role": "assistant", "content": "hi"}, "finish_reason": "stop"}], + model="gpt-4o-mini", + ) + + with patch( + "litellm.proxy.rag_endpoints.endpoints.litellm.aquery", + new_callable=AsyncMock, + return_value=mock_response, + ) as mock_aquery, patch.object(litellm, "vector_store_registry", mock_registry), patch( + "litellm.proxy.rag_endpoints.endpoints.assert_user_can_access_vector_store_id", + new=AsyncMock(), + ), patch( + "litellm.proxy.vector_store_endpoints.endpoints.assert_user_can_access_vector_store", + new=AsyncMock(), + ): + response = client_internal_user.post( + "/v1/rag/query", + json={ + "model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "hello"}], + "retrieval_config": {"vector_store_id": "s3-store"}, + }, + ) + + assert response.status_code == 200, response.json() + mock_aquery.assert_awaited_once() + forwarded_config = mock_aquery.await_args.kwargs["retrieval_config"] + assert forwarded_config["vector_store_id"] == "s3-store" + assert forwarded_config["custom_llm_provider"] == "s3_vectors" + assert forwarded_config["aws_region_name"] == "eu-west-1" + assert forwarded_config["embedding_model"] == "my-embed" + assert forwarded_config["vector_bucket_name"] == "bkt" + + +def test_rag_query_user_retrieval_config_wins_over_store(client_internal_user): + """User-supplied retrieval_config keys must win over registry values.""" + import litellm + from litellm.types.utils import ModelResponse + + mock_vector_store = { + "vector_store_id": "s3-store", + "custom_llm_provider": "s3_vectors", + "litellm_params": {"aws_region_name": "eu-west-1"}, + } + mock_registry = MagicMock() + mock_registry.get_litellm_managed_vector_store_from_registry.return_value = mock_vector_store + + mock_response = ModelResponse( + id="chatcmpl-test", + choices=[{"index": 0, "message": {"role": "assistant", "content": "hi"}, "finish_reason": "stop"}], + model="gpt-4o-mini", + ) + + with patch( + "litellm.proxy.rag_endpoints.endpoints.litellm.aquery", + new_callable=AsyncMock, + return_value=mock_response, + ) as mock_aquery, patch.object(litellm, "vector_store_registry", mock_registry), patch( + "litellm.proxy.rag_endpoints.endpoints.assert_user_can_access_vector_store_id", + new=AsyncMock(), + ), patch( + "litellm.proxy.vector_store_endpoints.endpoints.assert_user_can_access_vector_store", + new=AsyncMock(), + ): + response = client_internal_user.post( + "/v1/rag/query", + json={ + "model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "hello"}], + "retrieval_config": {"vector_store_id": "s3-store", "aws_region_name": "us-east-1"}, + }, + ) + + assert response.status_code == 200, response.json() + forwarded_config = mock_aquery.await_args.kwargs["retrieval_config"] + assert forwarded_config["aws_region_name"] == "us-east-1" diff --git a/tests/test_litellm/rag/test_main.py b/tests/test_litellm/rag/test_main.py index 584124ba06a..d8ffae667b1 100644 --- a/tests/test_litellm/rag/test_main.py +++ b/tests/test_litellm/rag/test_main.py @@ -254,6 +254,96 @@ async def test_aquery_streaming_bills_sub_call_costs_into_final_event(): assert standard_logging_object["response_cost"] >= 0.003 +@pytest.mark.asyncio +async def test_aquery_forwards_provider_retrieval_config_and_router_to_search(): + """ + Regression: provider-specific retrieval_config keys (aws_region_name, + embedding_model, vector_bucket_name, ...) and the router must be forwarded + to the vector store search call. Pre-fix they were silently dropped, so + /v1/rag/query failed with provider config errors (e.g. S3 Vectors + "aws_region_name is required") even when the caller supplied them. + """ + from unittest.mock import AsyncMock + + from litellm.types.vector_stores import VectorStoreSearchResponse + + router = litellm.Router( + model_list=[ + { + "model_name": "gpt-4o-mini", + "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "test-key"}, + } + ] + ) + + fake_search = AsyncMock( + return_value=VectorStoreSearchResponse( + object="vector_store.search_results.page", search_query="q", data=[] + ) + ) + with patch("litellm.vector_stores.asearch", new=fake_search): + response = await litellm.aquery( + model="gpt-4o-mini", + messages=[{"role": "user", "content": "hello"}], + retrieval_config={ + "vector_store_id": "bkt:idx", + "custom_llm_provider": "s3_vectors", + "top_k": 5, + "aws_region_name": "eu-west-1", + "embedding_model": "my-embed", + "vector_bucket_name": "bkt", + }, + router=router, + mock_response="hi", + ) + + assert isinstance(response, ModelResponse) + fake_search.assert_awaited_once() + search_kwargs = fake_search.await_args.kwargs + assert search_kwargs["vector_store_id"] == "bkt:idx" + assert search_kwargs["custom_llm_provider"] == "s3_vectors" + assert search_kwargs["max_num_results"] == 5 + assert search_kwargs["router"] is router + # provider-specific extras forwarded + assert search_kwargs["aws_region_name"] == "eu-west-1" + assert search_kwargs["embedding_model"] == "my-embed" + assert search_kwargs["vector_bucket_name"] == "bkt" + # consumed keys are not duplicated into the spread + assert "top_k" not in search_kwargs + + +@pytest.mark.asyncio +async def test_aquery_minimal_retrieval_config_forwards_no_extras(): + """ + A minimal retrieval_config must not leak consumed keys (or invent extras) + into the vector store search call. + """ + from unittest.mock import AsyncMock + + from litellm.types.vector_stores import VectorStoreSearchResponse + + fake_search = AsyncMock( + return_value=VectorStoreSearchResponse( + object="vector_store.search_results.page", search_query="q", data=[] + ) + ) + with patch("litellm.vector_stores.asearch", new=fake_search): + await litellm.aquery( + model="gpt-4o-mini", + messages=[{"role": "user", "content": "hello"}], + retrieval_config={"vector_store_id": "vs_test_123", "custom_llm_provider": "openai"}, + mock_response="hi", + ) + + fake_search.assert_awaited_once() + search_kwargs = fake_search.await_args.kwargs + assert search_kwargs["vector_store_id"] == "vs_test_123" + assert search_kwargs["custom_llm_provider"] == "openai" + assert search_kwargs["router"] is None + leaked = {"top_k", "filters", "retrieval_filter", "aws_region_name", "embedding_model", "vector_bucket_name"} + assert not (leaked & set(search_kwargs.keys())) + + def test_rag_call_types_are_registered(): """ query/aquery/ingest/aingest are @client-decorated entry points, so their diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index a9e5b3316e0..91d2973af76 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -5936,3 +5936,58 @@ async def test_acreate_batch_request_bedrock_tags_override_deployment_tags(): bedrock_tags=request_tags, ) assert mock_sign.call_args.kwargs["data"]["tags"] == request_tags + + +@pytest.mark.asyncio +async def test_avector_store_search_injects_router(): + """ + Regression: router.avector_store_search must pass the router down to the + SDK search call so provider transforms can resolve router-managed + embedding models (e.g. S3 Vectors query embeddings). + """ + from litellm.types.vector_stores import VectorStoreSearchResponse + + mock_asearch = AsyncMock( + return_value=VectorStoreSearchResponse( + object="vector_store.search_results.page", search_query="q", data=[] + ) + ) + # Router.__init__ binds asearch via a local import, so patch the module + # attribute before constructing the Router. + with patch("litellm.vector_stores.main.asearch", new=mock_asearch): + router = litellm.Router( + model_list=[ + { + "model_name": "gpt-3.5-turbo", + "litellm_params": {"model": "openai/gpt-3.5-turbo", "api_key": "test-key"}, + } + ] + ) + await router.avector_store_search( + vector_store_id="v", query="q", custom_llm_provider="s3_vectors" + ) + + mock_asearch.assert_awaited_once() + assert mock_asearch.await_args.kwargs["router"] is router + + +@pytest.mark.asyncio +async def test_avector_store_create_does_not_inject_router(): + """The router injection is gated on the search call type: the create path + must keep calling the SDK without a router kwarg.""" + mock_acreate = AsyncMock(return_value={"id": "vs_1", "object": "vector_store"}) + # avector_store_create(model=None) resolves acreate via a local import at + # call time, so patching after Router construction works here. + router = litellm.Router( + model_list=[ + { + "model_name": "gpt-3.5-turbo", + "litellm_params": {"model": "openai/gpt-3.5-turbo", "api_key": "test-key"}, + } + ] + ) + with patch("litellm.vector_stores.main.acreate", new=mock_acreate): + await router.avector_store_create(model=None, custom_llm_provider="openai") + + mock_acreate.assert_awaited_once() + assert "router" not in mock_acreate.await_args.kwargs diff --git a/tests/test_litellm/vector_stores/test_main.py b/tests/test_litellm/vector_stores/test_main.py new file mode 100644 index 00000000000..3fdf4d9daa5 --- /dev/null +++ b/tests/test_litellm/vector_stores/test_main.py @@ -0,0 +1,77 @@ +""" +Tests for litellm/vector_stores/main.py. + +Pins the router threading contract for vector store search: the router is an +explicit named parameter that reaches the HTTP handler, and it must never leak +into litellm_params/kwargs where logging would model_dump() it (the #19550 +serialization trap). +""" + +from unittest.mock import MagicMock, patch + +import litellm.vector_stores.main as vector_stores_main +from litellm.vector_stores.main import search + +MOCK_SEARCH_RESPONSE = { + "object": "vector_store.search_results.page", + "search_query": "q", + "data": [], +} + + +def test_search_threads_router_to_handler(): + """search() must pass its router param through to the HTTP handler""" + mock_router = MagicMock() + logger = MagicMock() + + with ( + patch( + "litellm.vector_stores.main.ProviderConfigManager.get_provider_vector_stores_config", + return_value=MagicMock(), + ), + patch.object( + vector_stores_main.base_llm_http_handler, + "vector_store_search_handler", + return_value=MOCK_SEARCH_RESPONSE, + ) as mock_handler, + ): + search( + vector_store_id="bkt:idx", + query="q", + custom_llm_provider="s3_vectors", + router=mock_router, + litellm_logging_obj=logger, + ) + + mock_handler.assert_called_once() + assert mock_handler.call_args.kwargs["router"] is mock_router + + +def test_search_router_not_in_litellm_params(): + """Regression (#19550 class): the router must stay out of GenericLiteLLMParams, + otherwise pre-call logging model_dump()s it and breaks serialization.""" + mock_router = MagicMock() + logger = MagicMock() + + with ( + patch( + "litellm.vector_stores.main.ProviderConfigManager.get_provider_vector_stores_config", + return_value=MagicMock(), + ), + patch.object( + vector_stores_main.base_llm_http_handler, + "vector_store_search_handler", + return_value=MOCK_SEARCH_RESPONSE, + ) as mock_handler, + ): + search( + vector_store_id="bkt:idx", + query="q", + custom_llm_provider="s3_vectors", + router=mock_router, + litellm_logging_obj=logger, + ) + + litellm_params = mock_handler.call_args.kwargs["litellm_params"] + assert "router" not in litellm_params.model_dump(exclude_none=True) + assert getattr(litellm_params, "router", None) is None diff --git a/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/VectorStoreTester.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/VectorStoreTester.test.tsx index cbabcc6dca5..f375bfdd351 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/VectorStoreTester.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/VectorStoreTester.test.tsx @@ -128,16 +128,33 @@ describe("VectorStoreTester", () => { await waitFor(() => expect(mockSearch).toHaveBeenCalledTimes(1)); }); - it("reports a failed search and keeps the history empty", async () => { + it("shows the backend error in the history when a search fails", async () => { const user = userEvent.setup(); - mockSearch.mockRejectedValue(new Error("boom")); + const errorBody = '{"error":{"message":"OpenAIException - api_key is required"}}'; + mockSearch.mockRejectedValue(new Error(errorBody)); renderTester(); await user.type(queryInput(), "hello"); await user.click(searchButton()); - await waitFor(() => expect(mockFromBackend).toHaveBeenCalledWith("Failed to search vector store")); - expect(screen.getByText(EMPTY_STATE)).toBeInTheDocument(); + await waitFor(() => expect(mockFromBackend).toHaveBeenCalledWith(errorBody)); + expect(screen.getByText(`Search failed: ${errorBody}`)).toBeInTheDocument(); + expect(screen.queryByText("No results found")).not.toBeInTheDocument(); + expect(screen.queryByText(EMPTY_STATE)).not.toBeInTheDocument(); + // the failed query stays in the input for retry + expect(queryInput()).toHaveValue("hello"); + }); + + it('renders "No results found" for an empty result set, not an error', async () => { + const user = userEvent.setup(); + mockSearch.mockResolvedValue({ object: "vector_store.search_results.page", search_query: "hello", data: [] }); + renderTester(); + + await user.type(queryInput(), "hello"); + await user.click(searchButton()); + + expect(await screen.findByText("No results found")).toBeInTheDocument(); + expect(screen.queryByText(/search failed/i)).not.toBeInTheDocument(); }); it("clears the search history", async () => { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/VectorStoreTester.tsx b/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/VectorStoreTester.tsx index 015d58e8649..65b31bb911a 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/VectorStoreTester.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/VectorStoreTester.tsx @@ -41,6 +41,7 @@ export const VectorStoreTester: React.FC = ({ vectorStor { query: string; response: VectorStoreSearchResponse | null; + error: string | null; timestamp: number; }[] >([]); @@ -60,6 +61,7 @@ export const VectorStoreTester: React.FC = ({ vectorStor const historyEntry = { query, response, + error: null, timestamp: Date.now(), }; @@ -67,7 +69,9 @@ export const VectorStoreTester: React.FC = ({ vectorStor setQuery(""); } catch (error) { console.error("Error searching vector store:", error); - NotificationsManager.fromBackend("Failed to search vector store"); + const errorMessage = error instanceof Error ? error.message : String(error); + NotificationsManager.fromBackend(errorMessage); + setSearchHistory((prev) => [{ query, response: null, error: errorMessage, timestamp: Date.now() }, ...prev]); } finally { setIsLoading(false); } @@ -228,6 +232,8 @@ export const VectorStoreTester: React.FC = ({ vectorStor ); })} + ) : entry.error ? ( +
Search failed: {entry.error}
) : (
No results found
)} diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index 576e16cbb37..a8408fecaad 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -6851,7 +6851,7 @@ export const vectorStoreSearchCall = async ( if (!response.ok) { const errorData = await response.text(); await handleError(errorData); - return null; + throw new Error(errorData); } const data = await response.json(); From 9e25dd708fa7a8b2af354e6901683fd604c4a19a Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 31 Aug 2026 21:47:13 -0700 Subject: [PATCH 02/48] feat(streaming): carry final response cost on streamed usage by default Streamed responses through the proxy previously exposed no usable cost: the x-litellm-response-cost header is unreadable mid-stream and the final usage chunk carried only tokens, priced against an alias model name the client cannot resolve. The include_cost_in_streaming_usage flag existed but was off by default and only fixed the wire, not SDK clients. Stamp usage.cost into the joined streaming response by default wherever a final usage object is built: the chat-completions stream_chunk_builder, the native /v1/responses RESPONSE_COMPLETED event, and synthetic response events. Provider-reported cost always wins over the computed value, and only positive computed costs are stamped so unpriceable alias responses keep deferring to the logging object's own calculation. Per-chunk SSE cost injection (/v1/messages, generateContent, passthrough) stays behind the flag. Also normalize non-litellm usage objects in stream_chunk_builder: openai CompletionUsage lacks Usage.__contains__, so membership probes silently returned False and client-side rebuilds dropped the wire cost and recounted token usage locally. Wire token counts and cost now survive. Resolves LIT-6427 --- .../streaming_chunk_builder_utils.py | 8 ++- litellm/main.py | 22 ++++--- .../streaming_iterator.py | 10 --- litellm/responses/streaming_iterator.py | 46 ++++++------- .../test_streaming_chunk_builder_utils.py | 53 +++++++++++++++ .../responses/test_streaming_iterator.py | 52 +++++++++++++++ tests/test_litellm/test_main.py | 66 +++++++++++++++++-- 7 files changed, 204 insertions(+), 53 deletions(-) diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index 0e2139d688b..276616f9eee 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -36,6 +36,8 @@ from litellm.types.utils import ( from litellm.utils import print_verbose, token_counter if TYPE_CHECKING: + from openai.types.completion_usage import CompletionUsage + from litellm.litellm_core_utils.litellm_logging import Logging from litellm.types.litellm_core_utils.streaming_chunk_builder_utils import ( UsagePerChunk, @@ -782,7 +784,7 @@ class ChunkProcessor: @staticmethod def _extract_usage_chunk(chunk: "_UsageBearingChunk | ModelResponse | ModelResponseStream") -> Usage | None: - usage_chunk: Usage | None = None + usage_chunk: Usage | CompletionUsage | None = None if hasattr(chunk, "usage") and chunk.usage is not None: usage_chunk = chunk.usage elif "usage" in chunk: @@ -794,7 +796,9 @@ class ChunkProcessor: if isinstance(usage_chunk, dict): return Usage(**usage_chunk) - return usage_chunk + if usage_chunk is None or isinstance(usage_chunk, Usage): + return usage_chunk + return Usage(**usage_chunk.model_dump()) def _calculate_usage_per_chunk( self, diff --git a/litellm/main.py b/litellm/main.py index 0c8bff16f81..7ca84226b09 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -8634,6 +8634,16 @@ def _set_stream_builder_response_cost(response: ModelResponse, logging_obj: Opti hidden_params["response_cost"] = response_cost +def _stamp_streaming_usage_cost(usage: Usage, response: ModelResponse, logging_obj: Optional["Logging"]) -> None: + if logging_obj is None: + return + if isinstance(getattr(usage, "cost", None), (int, float)): + return + computed_cost: Final = logging_obj._response_cost_calculator(result=response) + if isinstance(computed_cost, (int, float)) and computed_cost > 0: + setattr(usage, "cost", computed_cost) + + def stream_chunk_builder( chunks: list, messages: list | None = None, @@ -8728,12 +8738,7 @@ def stream_chunk_builder( ) break - if litellm.include_cost_in_streaming_usage and logging_obj is not None: - setattr( - usage, - "cost", - logging_obj._response_cost_calculator(result=response), - ) + _stamp_streaming_usage_cost(usage, response, logging_obj) _set_stream_builder_response_cost(response, logging_obj) processor.apply_provider_assembled_streaming_metadata(response, chunks, logging_obj) @@ -8912,10 +8917,7 @@ def stream_chunk_builder( ) break - # Add cost to usage object if include_cost_in_streaming_usage is True - if litellm.include_cost_in_streaming_usage and logging_obj is not None: - setattr(usage, "cost", logging_obj._response_cost_calculator(result=response)) - + _stamp_streaming_usage_cost(usage, response, logging_obj) _set_stream_builder_response_cost(response, logging_obj) processor.apply_provider_assembled_streaming_metadata(response, chunks, logging_obj) diff --git a/litellm/responses/litellm_completion_transformation/streaming_iterator.py b/litellm/responses/litellm_completion_transformation/streaming_iterator.py index 8b1eeb30306..27afff39c0f 100644 --- a/litellm/responses/litellm_completion_transformation/streaming_iterator.py +++ b/litellm/responses/litellm_completion_transformation/streaming_iterator.py @@ -1164,16 +1164,6 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): def _emit_response_completed_event(self, litellm_model_response: ModelResponse) -> ResponseCompletedEvent | None: if litellm_model_response: - # Add cost to usage object if include_cost_in_streaming_usage is True - if litellm.include_cost_in_streaming_usage and self.litellm_logging_obj is not None: - usage: Final[object] = getattr(litellm_model_response, "usage", None) - if usage is not None: - setattr( - usage, - "cost", - self.litellm_logging_obj._response_cost_calculator(result=litellm_model_response), - ) - # Transform the response responses_api_response: Final = ( LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index d070f7758fd..2b4252aa1c5 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -405,23 +405,7 @@ class BaseResponsesAPIStreamingIterator: openai_types.ResponsesAPIStreamEvents.RESPONSE_FAILED, ): self.completed_response = openai_responses_api_chunk - # Add cost to usage object if include_cost_in_streaming_usage is True - if litellm.include_cost_in_streaming_usage and self.logging_obj is not None: - response_obj: Final[ResponsesAPIResponse | None] = getattr( - openai_responses_api_chunk, "response", None - ) - if response_obj: - usage_obj: Final[ResponseAPIUsage | None] = getattr(response_obj, "usage", None) - if usage_obj is not None: - try: - cost: Final[float | None] = self.logging_obj._response_cost_calculator( - result=response_obj - ) - if cost is not None: - setattr(usage_obj, "cost", cost) - except Exception: - # Best-effort usage cost annotation should not break stream replay. - pass + _stamp_responses_usage_cost(getattr(openai_responses_api_chunk, "response", None), self.logging_obj) if _chunk_type == openai_types.ResponsesAPIStreamEvents.RESPONSE_FAILED: self._handle_logging_failed_response() @@ -1272,6 +1256,24 @@ def _add_text_like_part_events( ) +def _stamp_responses_usage_cost( + response_obj: ResponsesAPIResponse | None, logging_obj: LiteLLMLoggingObj | None +) -> None: + if response_obj is None or logging_obj is None: + return + usage_obj: Final[ResponseAPIUsage | None] = getattr(response_obj, "usage", None) + if usage_obj is None: + return + if isinstance(getattr(usage_obj, "cost", None), (int, float)): + return + try: + cost: Final[float | None] = logging_obj._response_cost_calculator(result=response_obj) + except Exception: + return + if isinstance(cost, (int, float)) and cost > 0: + setattr(usage_obj, "cost", cost) + + def _build_synthetic_response_events( *, transformed: ResponsesAPIResponse, @@ -1279,15 +1281,7 @@ def _build_synthetic_response_events( chunk_size: int, ) -> list[ResponsesAPIStreamingResponse]: openai_types: Final = _get_openai_response_types() - if litellm.include_cost_in_streaming_usage and logging_obj is not None: - usage_obj: Final = transformed.usage if hasattr(transformed, "usage") else None - if usage_obj is not None: - try: - cost: Final[float | None] = logging_obj._response_cost_calculator(result=transformed) - if cost is not None: - setattr(usage_obj, "cost", cost) - except Exception: - pass + _stamp_responses_usage_cost(transformed, logging_obj) events: Final[list[ResponsesAPIStreamingResponse]] = [ _build_response_status_event(openai_types.ResponsesAPIStreamEvents.RESPONSE_CREATED, transformed), diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py b/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py index 8ac050a04f9..bacbcbf132b 100644 --- a/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py +++ b/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py @@ -592,6 +592,59 @@ def test_stream_chunk_builder_litellm_usage_chunks(): assert usage.total_tokens == 77 +def test_calculate_usage_honors_openai_sdk_completion_usage_chunks(): + from openai.types.completion_usage import CompletionUsage + + content_chunk = ModelResponseStream( + id="chatcmpl-sdk-usage-1", + created=1745513206, + model="mantle-claude", + object="chat.completion.chunk", + system_fingerprint=None, + choices=[ + StreamingChoices( + finish_reason="stop", + index=0, + delta=Delta( + provider_specific_fields=None, + content="ok", + role=None, + function_call=None, + tool_calls=None, + audio=None, + ), + logprobs=None, + ) + ], + provider_specific_fields=None, + stream_options={"include_usage": True}, + ) + usage_chunk = ModelResponseStream( + id="chatcmpl-sdk-usage-1", + created=1745513207, + model="mantle-claude", + object="chat.completion.chunk", + system_fingerprint=None, + choices=[], + provider_specific_fields=None, + stream_options={"include_usage": True}, + ) + usage_chunk.usage = CompletionUsage( + prompt_tokens=20, completion_tokens=60, total_tokens=80, cost=0.000704 + ) + assert type(usage_chunk.usage) is CompletionUsage + + chunks = [content_chunk, usage_chunk] + usage = ChunkProcessor(chunks=chunks).calculate_usage( + chunks=chunks, model="mantle-claude", completion_output="" + ) + + assert usage.prompt_tokens == 20 + assert usage.completion_tokens == 60 + assert usage.total_tokens == 80 + assert getattr(usage, "cost", None) == pytest.approx(0.000704) + + def test_get_model_from_chunks_azure_model_router(): """ Test that _get_model_from_chunks finds the actual model from Azure Model Router chunks. diff --git a/tests/test_litellm/responses/test_streaming_iterator.py b/tests/test_litellm/responses/test_streaming_iterator.py index 677faf7f655..9edcaaef034 100644 --- a/tests/test_litellm/responses/test_streaming_iterator.py +++ b/tests/test_litellm/responses/test_streaming_iterator.py @@ -326,3 +326,55 @@ def test_run_post_success_hooks_does_not_report_generation_time_as_overhead(): assert iterator.completed_response._hidden_params["_response_ms"] == 10000.0 assert "litellm_overhead_time_ms" not in iterator.completed_response._hidden_params + + +def _responses_api_response_with_usage() -> ResponsesAPIResponse: + from litellm.types.llms.openai import ResponseAPIUsage + + return ResponsesAPIResponse( + id="resp_lit6427", + created_at=int(datetime(2025, 1, 1).timestamp()), + status="completed", + model="mantle-claude", + object="response", + output=[], + usage=ResponseAPIUsage(input_tokens=20, output_tokens=60, total_tokens=80), + ) + + +def test_stamp_responses_usage_cost_stamps_computed_cost(): + from litellm.responses.streaming_iterator import _stamp_responses_usage_cost + + response = _responses_api_response_with_usage() + logging_obj = Mock(spec=LiteLLMLoggingObj) + logging_obj._response_cost_calculator.return_value = 0.000704 + + _stamp_responses_usage_cost(response, logging_obj) + + assert getattr(response.usage, "cost", None) == pytest.approx(0.000704) + logging_obj._response_cost_calculator.assert_called_once_with(result=response) + + +def test_stamp_responses_usage_cost_keeps_provider_reported_cost(): + from litellm.responses.streaming_iterator import _stamp_responses_usage_cost + + response = _responses_api_response_with_usage() + setattr(response.usage, "cost", 0.5) + logging_obj = Mock(spec=LiteLLMLoggingObj) + + _stamp_responses_usage_cost(response, logging_obj) + + assert getattr(response.usage, "cost", None) == pytest.approx(0.5) + logging_obj._response_cost_calculator.assert_not_called() + + +def test_stamp_responses_usage_cost_survives_calculator_failure(): + from litellm.responses.streaming_iterator import _stamp_responses_usage_cost + + response = _responses_api_response_with_usage() + logging_obj = Mock(spec=LiteLLMLoggingObj) + logging_obj._response_cost_calculator.side_effect = RuntimeError("cost map unavailable") + + _stamp_responses_usage_cost(response, logging_obj) + + assert getattr(response.usage, "cost", None) is None diff --git a/tests/test_litellm/test_main.py b/tests/test_litellm/test_main.py index 8cf878d05d9..7c2b9d0be05 100644 --- a/tests/test_litellm/test_main.py +++ b/tests/test_litellm/test_main.py @@ -3150,8 +3150,8 @@ def _stream_builder_logging_obj() -> LiteLLMLogging: return logging_obj -def test_stream_chunk_builder_reports_streaming_usage_cost_when_enabled(monkeypatch: pytest.MonkeyPatch): - monkeypatch.setattr(litellm, "include_cost_in_streaming_usage", True) +def test_stream_chunk_builder_stamps_streaming_usage_cost_by_default(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(litellm, "include_cost_in_streaming_usage", False) chunks: Final = [ _stream_builder_text_chunk("gpt-4o", "Hello "), _stream_builder_text_chunk("gpt-4o", "world.", finish_reason="stop"), @@ -3168,11 +3168,45 @@ def test_stream_chunk_builder_reports_streaming_usage_cost_when_enabled(monkeypa assert response._hidden_params["response_cost"] == pytest.approx(usage_cost) -def test_stream_chunk_builder_defers_cost_to_logging_obj_when_usage_cost_absent(monkeypatch: pytest.MonkeyPatch): - monkeypatch.setattr(litellm, "include_cost_in_streaming_usage", False) +def test_stream_chunk_builder_skips_stamp_when_cost_is_unpriceable(): + import time as time_module + + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging + + logging_obj: Final = LiteLLMLogging( + model="us.anthropic.claude-opus-5", + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="completion", + start_time=time_module.time(), + litellm_call_id="stream-builder-alias-unpriceable", + function_id="1", + ) + logging_obj.model_call_details["custom_llm_provider"] = "bedrock" + logging_obj.optional_params = {} + usage_chunk: Final = _stream_builder_text_chunk("bedrock-claude-opus-5", "") + usage_chunk.usage = Usage(prompt_tokens=40, completion_tokens=5, total_tokens=45) + chunks: Final = [ + _stream_builder_text_chunk("bedrock-claude-opus-5", "Hello ", finish_reason="stop"), + usage_chunk, + ] + + response: Final = litellm.stream_chunk_builder( + chunks=chunks, messages=[{"role": "user", "content": "hi"}], logging_obj=logging_obj + ) + + assert response is not None + assert getattr(response.usage, "cost", None) is None + assert response._hidden_params.get("response_cost") is None + + +def test_stream_chunk_builder_keeps_provider_reported_usage_cost(): + usage_chunk: Final = _stream_builder_text_chunk("gpt-4o", "") + usage_chunk.usage = Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15, cost=0.5) chunks: Final = [ _stream_builder_text_chunk("gpt-4o", "Hello "), _stream_builder_text_chunk("gpt-4o", "world.", finish_reason="stop"), + usage_chunk, ] response: Final = litellm.stream_chunk_builder( @@ -3180,4 +3214,26 @@ def test_stream_chunk_builder_defers_cost_to_logging_obj_when_usage_cost_absent( ) assert response is not None - assert response._hidden_params.get("response_cost") is None + assert getattr(response.usage, "cost", None) == pytest.approx(0.5) + assert response._hidden_params["response_cost"] == pytest.approx(0.5) + + +def test_stream_chunk_builder_prices_alias_from_openai_sdk_usage_chunk(): + from openai.types.completion_usage import CompletionUsage + + usage_chunk: Final = _stream_builder_text_chunk("mantle-claude", "") + usage_chunk.usage = CompletionUsage(prompt_tokens=20, completion_tokens=60, total_tokens=80, cost=0.000704) + assert type(usage_chunk.usage) is CompletionUsage + chunks: Final = [ + _stream_builder_text_chunk("mantle-claude", "Hello "), + _stream_builder_text_chunk("mantle-claude", "world.", finish_reason="stop"), + usage_chunk, + ] + + response: Final = litellm.stream_chunk_builder(chunks=chunks, messages=[{"role": "user", "content": "hi"}]) + + assert response is not None + assert response.usage.prompt_tokens == 20 + assert response.usage.completion_tokens == 60 + assert getattr(response.usage, "cost", None) == pytest.approx(0.000704) + assert response._hidden_params["response_cost"] == pytest.approx(0.000704) From 7ca035f310f891eea216f3162191739f3720bcb1 Mon Sep 17 00:00:00 2001 From: Kris Xia Date: Tue, 1 Sep 2026 11:39:33 +0800 Subject: [PATCH 03/48] fix(gemini): return enabled thinking content by default --- .../gemini/vertex_and_google_ai_studio_gemini.py | 2 +- .../test_vertex_and_google_ai_studio_gemini.py | 12 ++++++++++++ 2 files changed, 13 insertions(+), 1 deletion(-) diff --git a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py index d8b1e7ba17c..69fe5678de9 100644 --- a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py +++ b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py @@ -949,7 +949,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): # For Gemini 3+ models, use thinkingLevel instead of thinkingBudget if model and VertexGeminiConfig._is_gemini_3_or_newer(model): if thinking_enabled: - if thinking_budget is None or thinking_budget == 0: + if thinking_budget == 0: params["includeThoughts"] = False else: params["includeThoughts"] = True 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 bd07bec900f..d2788408e09 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 @@ -1185,6 +1185,18 @@ def test_vertex_ai_map_thinking_param_with_budget_tokens_0(): } +def test_vertex_ai_map_thinking_param_without_budget_tokens_for_gemini_3(): + v = VertexGeminiConfig() + result = v.map_openai_params( + non_default_params={"thinking": {"type": "enabled"}}, + optional_params={}, + model="gemini-3.5-flash", + drop_params=False, + ) + + assert result["thinkingConfig"] == {"includeThoughts": True} + + def test_vertex_ai_map_tools(): v = VertexGeminiConfig() optional_params = {} From f8298fa35c310eaf9a2cd8133cc4811c99a42c99 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 1 Sep 2026 12:22:15 -0700 Subject: [PATCH 04/48] fix(bedrock): stop Converse crashing on bearer-token auth without SigV4 credentials Since the Rust core handoff in #37241, BedrockConverseLLM.completion read access_key, secret_key and token off the boto3 credentials before asking the Rust gate whether it wanted the call. On a deployment that only sets AWS_BEARER_TOKEN_BEDROCK boto3 resolves no credentials, so every Converse call through /v1/chat/completions and /v1/responses failed with "'NoneType' object has no attribute 'access_key'", with or without the Rust opt-in Bearer auth resolves no SigV4 principal at all, and both the Python and the Rust path read the bearer token themselves, so only hand the principal keys down when boto3 actually resolved one get_request_headers now accepts credentials=None and raises botocore's NoCredentialsError when neither a bearer token nor a principal exists instead of handing SigV4Auth a None Fixes #38579 Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01BFDpYC45u9p8eKd4aBATnS --- litellm/llms/bedrock/base_aws_llm.py | 6 +- litellm/llms/bedrock/chat/converse_handler.py | 40 ++++++++----- .../chat/test_bedrock_converse_handler.py | 58 +++++++++++++++---- .../llms/bedrock/test_base_aws_llm.py | 18 ++++++ 4 files changed, 97 insertions(+), 25 deletions(-) diff --git a/litellm/llms/bedrock/base_aws_llm.py b/litellm/llms/bedrock/base_aws_llm.py index 852cfaa24f2..1e634ced29b 100644 --- a/litellm/llms/bedrock/base_aws_llm.py +++ b/litellm/llms/bedrock/base_aws_llm.py @@ -1442,7 +1442,7 @@ class BaseAWSLLM: @tracer.wrap() def get_request_headers( self, - credentials: Credentials, + credentials: Credentials | None, aws_region_name: str, extra_headers: dict | None, endpoint_url: str, @@ -1469,9 +1469,13 @@ class BaseAWSLLM: try: from botocore.auth import SigV4Auth from botocore.awsrequest import AWSRequest + from botocore.exceptions import NoCredentialsError except ImportError: raise ImportError("Missing boto3 to call bedrock. Run 'pip install boto3'.") + if credentials is None: + raise NoCredentialsError() + # Filter headers for AWS signature calculation # AWS SigV4 only includes specific headers in signature calculation aws_signature_headers: Final = self._filter_headers_for_aws_signature(headers) diff --git a/litellm/llms/bedrock/chat/converse_handler.py b/litellm/llms/bedrock/chat/converse_handler.py index ca5f1298360..7d5f99ca893 100644 --- a/litellm/llms/bedrock/chat/converse_handler.py +++ b/litellm/llms/bedrock/chat/converse_handler.py @@ -1,4 +1,6 @@ import json +from collections.abc import Mapping +from types import MappingProxyType from typing import Any, Final import httpx @@ -24,6 +26,22 @@ from ..common_utils import BedrockError, _get_all_bedrock_regions from .invoke_handler import AWSEventStreamDecoder, MockResponseIterator, make_call +def _sigv4_principal(credentials: Credentials | None) -> Mapping[str, str]: + if credentials is None: + return MappingProxyType({}) + return MappingProxyType( + { + key: value + for key, value in ( + ("aws_access_key_id", credentials.access_key), + ("aws_secret_access_key", credentials.secret_key), + ("aws_session_token", credentials.token), + ) + if value is not None + } + ) + + def make_sync_call( client: HTTPHandler | None, api_base: str, @@ -95,7 +113,7 @@ class BedrockConverseLLM(BaseAWSLLM): stream, optional_params: dict, litellm_params: dict, - credentials: Credentials, + credentials: Credentials | None, logger_fn=None, headers={}, client: AsyncHTTPHandler | None = None, @@ -167,7 +185,7 @@ class BedrockConverseLLM(BaseAWSLLM): stream, optional_params: dict, litellm_params: dict, - credentials: Credentials, + credentials: Credentials | None, logger_fn=None, headers: dict = {}, client: AsyncHTTPHandler | None = None, @@ -331,7 +349,7 @@ class BedrockConverseLLM(BaseAWSLLM): litellm_params["aws_region_name"] = aws_region_name # [DO NOT DELETE] important for async calls - credentials: Final[Credentials] = self.get_credentials( + credentials: Final[Credentials | None] = self.get_credentials( aws_access_key_id=aws_access_key_id, aws_secret_access_key=aws_secret_access_key, aws_session_token=aws_session_token, @@ -368,19 +386,13 @@ class BedrockConverseLLM(BaseAWSLLM): # The Rust core owns the whole call for the subset it accepts. Ask # before transforming so whichever path runs emits pre_call once, and # hand down the credentials, region and endpoint this handler already - # resolved so both paths sign as the same principal. + # resolved so both paths sign as the same principal. Bearer-token auth + # resolves no SigV4 principal at all, and each path reads that token + # itself. rust_optional_params: Final = { # mutable-ok: json.dumps in the bridge rejects a mappingproxy **optional_params, - **{ # mutable-ok: merged into its mutable parent above - key: value - for key, value in ( - ("aws_access_key_id", credentials.access_key), - ("aws_secret_access_key", credentials.secret_key), - ("aws_session_token", credentials.token), - ("aws_region_name", aws_region_name), - ) - if value is not None - }, + **_sigv4_principal(credentials), + "aws_region_name": aws_region_name, } serves_via_rust: Final = rust_chat_completions_accepts( model=model, diff --git a/tests/test_litellm/llms/bedrock/chat/test_bedrock_converse_handler.py b/tests/test_litellm/llms/bedrock/chat/test_bedrock_converse_handler.py index 8e67a7e3438..4a296d9296b 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_bedrock_converse_handler.py +++ b/tests/test_litellm/llms/bedrock/chat/test_bedrock_converse_handler.py @@ -96,10 +96,8 @@ def _completion_kwargs(**overrides): return kwargs -def _run(**overrides): - with patch.object( - BedrockConverseLLM, "get_credentials", return_value=RESOLVED_CREDENTIALS - ): +def _run(*, credentials: Credentials | None = RESOLVED_CREDENTIALS, **overrides): + with patch.object(BedrockConverseLLM, "get_credentials", return_value=credentials): return BedrockConverseLLM().completion(**_completion_kwargs(**overrides)) @@ -358,14 +356,27 @@ async def test_async_completion_logs_pre_call_by_default(): assert logging_obj.pre_call.call_count == 1 -def _sync_client_returning_converse_response(): +def _recording_sync_client(): + """A sync transport that answers with a Converse response and keeps every + `post` payload, so a test can assert which headers were sent.""" + posted: list[dict] = [] client = MagicMock() - client.post = lambda **_kwargs: httpx.Response( - 200, - json=CONVERSE_RESPONSE, - request=httpx.Request("POST", "https://bedrock-runtime.us-west-2.amazonaws.com"), - ) + + def post(**kwargs): + posted.append(kwargs) + return httpx.Response( + 200, + json=CONVERSE_RESPONSE, + request=httpx.Request("POST", "https://bedrock-runtime.us-west-2.amazonaws.com"), + ) + + client.post = post client.__class__ = HTTPHandler + return client, posted + + +def _sync_client_returning_converse_response(): + client, _ = _recording_sync_client() return client @@ -487,3 +498,30 @@ def test_post_call_is_not_logged_twice_when_the_sync_rust_call_declines(): assert response.choices[0].message.content == "hi" assert len(calls["post_call"]) == 1 assert "hi" in calls["post_call"][0]["original_response"] + + +def test_bearer_token_auth_serves_when_boto3_resolves_no_sigv4_credentials(monkeypatch): + """With only `AWS_BEARER_TOKEN_BEDROCK` configured boto3 resolves no + credentials at all. Preparing the Rust handoff must not dereference that + None: the bearer token signs the request on its own.""" + monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", "bedrock-bearer-token") + client, posted = _recording_sync_client() + + response = _run(credentials=None, litellm_params={}, client=client) + + assert response.choices[0].message.content == "hi" + assert posted[0]["headers"]["Authorization"] == "Bearer bedrock-bearer-token" + + +def test_the_rust_opt_in_needs_no_sigv4_principal(): + """The core resolves the bearer token itself, so a bearer-only deployment + keeps its opt-in and the gate sees no aws_* credential keys to sign with.""" + seen = _inject() + + response = _run(credentials=None, api_key="bedrock-bearer-token") + + assert response.choices[0].message.content == "hello from rust" + params = seen["call"][0]["optional_params"] + assert not {"aws_access_key_id", "aws_secret_access_key", "aws_session_token"} & params.keys() + assert params["aws_region_name"] == "us-east-1" + assert seen["call"][0]["api_key"] == "bedrock-bearer-token" 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 7d07ac947b1..f854d806bdc 100644 --- a/tests/test_litellm/llms/bedrock/test_base_aws_llm.py +++ b/tests/test_litellm/llms/bedrock/test_base_aws_llm.py @@ -15,6 +15,7 @@ from unittest.mock import MagicMock, patch from botocore.awsrequest import AWSPreparedRequest, AWSRequest from botocore.auth import SigV4Auth from botocore.credentials import Credentials +from botocore.exceptions import NoCredentialsError import litellm from litellm.llms.bedrock.base_aws_llm import ( @@ -801,6 +802,23 @@ def test_get_request_headers_with_sigv4(): assert result == mock_request.prepare.return_value +def test_get_request_headers_without_credentials_or_bearer_token_raises_no_credentials(): + """Bearer-token auth needs no SigV4 principal, so `credentials` may be None. + Reaching the SigV4 branch with neither must fail the way botocore always + has instead of signing with a missing principal.""" + llm = BaseAWSLLM() + + with patch.dict(os.environ, {}, clear=True), pytest.raises(NoCredentialsError): + llm.get_request_headers( + credentials=None, + aws_region_name="us-west-2", + extra_headers=None, + endpoint_url="https://api.example.com", + data='{"prompt": "test"}', + headers={"Content-Type": "application/json"}, + ) + + def test_sigv4_matches_rust_golden_vector(): request = AWSRequest( method="POST", From 24ee419c85f6758369e6582250a50490ce1d6819 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 1 Sep 2026 19:28:04 +0000 Subject: [PATCH 05/48] fix(models): registry audit 2026-09-01 for openai realtime, mistral aliases, voyage, xai, fireworks Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...odel_prices_and_context_window_backup.json | 133 ++++++++++++++---- model_prices_and_context_window.json | 133 ++++++++++++++---- 2 files changed, 212 insertions(+), 54 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 27ff525c15e..08add22c998 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -30606,17 +30606,18 @@ }, "gpt-realtime-2": { "cache_creation_input_audio_token_cost": 4e-07, + "cache_read_input_audio_token_cost": 4e-07, "cache_read_input_token_cost": 4e-07, "input_cost_per_audio_token": 3.2e-05, "input_cost_per_image": 5e-06, "input_cost_per_token": 4e-06, "litellm_provider": "openai", - "max_input_tokens": 32000, - "max_output_tokens": 4096, - "max_tokens": 4096, + "max_input_tokens": 128000, + "max_output_tokens": 32000, + "max_tokens": 32000, "mode": "realtime", "output_cost_per_audio_token": 6.4e-05, - "output_cost_per_token": 1.6e-05, + "output_cost_per_token": 2.4e-05, "supported_endpoints": [ "/v1/realtime" ], @@ -30680,8 +30681,8 @@ "input_cost_per_token": 6e-07, "litellm_provider": "openai", "max_input_tokens": 128000, - "max_output_tokens": 4096, - "max_tokens": 4096, + "max_output_tokens": 32000, + "max_tokens": 32000, "mode": "realtime", "output_cost_per_audio_token": 2e-05, "output_cost_per_token": 2.4e-06, @@ -30713,7 +30714,7 @@ "input_cost_per_audio_token": 1e-05, "input_cost_per_token": 6e-07, "litellm_provider": "openai", - "max_input_tokens": 128000, + "max_input_tokens": 32000, "max_output_tokens": 4096, "max_tokens": 4096, "mode": "realtime", @@ -33477,19 +33478,21 @@ "source": "https://mistral.ai/pricing#api-pricing" }, "mistral/magistral-medium-latest": { - "input_cost_per_token": 2e-06, + "cache_read_input_token_cost": 1.5e-07, + "input_cost_per_token": 1.5e-06, "litellm_provider": "mistral", - "max_input_tokens": 40000, - "max_output_tokens": 40000, - "max_tokens": 40000, + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, "mode": "chat", - "output_cost_per_token": 5e-06, - "source": "https://mistral.ai/news/magistral", + "output_cost_per_token": 7.5e-06, + "source": "https://docs.mistral.ai/models/model-cards/mistral-medium-3-5-26-04", "supports_assistant_prefill": true, "supports_function_calling": true, "supports_reasoning": true, "supports_response_schema": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_vision": true }, "mistral/magistral-small-2506": { "deprecation_date": "2025-11-30", @@ -33508,19 +33511,21 @@ "supports_tool_choice": true }, "mistral/magistral-small-latest": { - "input_cost_per_token": 5e-07, + "cache_read_input_token_cost": 1.5e-08, + "input_cost_per_token": 1.5e-07, "litellm_provider": "mistral", - "max_input_tokens": 40000, - "max_output_tokens": 40000, - "max_tokens": 40000, + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, "mode": "chat", - "output_cost_per_token": 1.5e-06, - "source": "https://mistral.ai/pricing#api-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_reasoning": true, "supports_response_schema": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_vision": true }, "mistral/magistral-small-1-2-2509": { "deprecation_date": "2026-07-31", @@ -33652,16 +33657,21 @@ "supports_vision": true }, "mistral/mistral-medium": { - "input_cost_per_token": 2.7e-06, + "cache_read_input_token_cost": 1.5e-07, + "input_cost_per_token": 1.5e-06, "litellm_provider": "mistral", - "max_input_tokens": 32000, - "max_output_tokens": 8191, - "max_tokens": 8191, + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, "mode": "chat", - "output_cost_per_token": 8.1e-06, + "output_cost_per_token": 7.5e-06, + "source": "https://docs.mistral.ai/models/model-cards/mistral-medium-3-5-26-04", "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_reasoning": true, "supports_response_schema": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_vision": true }, "mistral/mistral-medium-2312": { "deprecation_date": "2025-06-16", @@ -45539,6 +45549,26 @@ "mode": "rerank", "output_cost_per_token": 0.0 }, + "voyage/rerank-3": { + "input_cost_per_token": 5e-08, + "litellm_provider": "voyage", + "max_input_tokens": 32000, + "max_output_tokens": 32000, + "max_tokens": 32000, + "mode": "rerank", + "output_cost_per_token": 0.0, + "source": "https://docs.voyageai.com/docs/pricing" + }, + "voyage/rerank-3-lite": { + "input_cost_per_token": 2e-08, + "litellm_provider": "voyage", + "max_input_tokens": 32000, + "max_output_tokens": 32000, + "max_tokens": 32000, + "mode": "rerank", + "output_cost_per_token": 0.0, + "source": "https://docs.voyageai.com/docs/pricing" + }, "voyage/voyage-2": { "input_cost_per_token": 1e-07, "litellm_provider": "voyage", @@ -46790,6 +46820,27 @@ "supports_vision": true, "supports_web_search": true }, + "xai/grok-build-latest": { + "cache_read_input_token_cost": 3e-07, + "cache_read_input_token_cost_above_200k_tokens": 6e-07, + "input_cost_per_token": 2e-06, + "input_cost_per_token_above_200k_tokens": 4e-06, + "litellm_provider": "xai", + "max_input_tokens": 500000, + "max_output_tokens": 500000, + "max_tokens": 500000, + "mode": "chat", + "output_cost_per_token": 6e-06, + "output_cost_per_token_above_200k_tokens": 1.2e-05, + "source": "https://docs.x.ai/developers/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true + }, "xai/grok-4.6": { "cache_read_input_token_cost": 5e-07, "cache_read_input_token_cost_above_200k_tokens": 1e-06, @@ -57110,6 +57161,34 @@ "supports_tool_choice": true, "supports_vision": false }, + "fireworks_ai/accounts/fireworks/models/glm-5p3-flash": { + "cache_read_input_token_cost": 3e-08, + "input_cost_per_token": 1.5e-07, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 1048576, + "max_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 5e-07, + "source": "https://docs.fireworks.ai/serverless/pricing", + "supports_function_calling": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "fireworks_ai/accounts/fireworks/models/inkling": { + "cache_read_input_token_cost": 1.7e-07, + "input_cost_per_token": 1e-06, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 1048576, + "max_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 4.05e-06, + "source": "https://fireworks.ai/models/fireworks/inkling", + "supports_function_calling": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, "fireworks_ai/accounts/fireworks/models/qwen3-embedding-8b": { "input_cost_per_token": 1e-07, "output_cost_per_token": 0.0, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 27ff525c15e..08add22c998 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -30606,17 +30606,18 @@ }, "gpt-realtime-2": { "cache_creation_input_audio_token_cost": 4e-07, + "cache_read_input_audio_token_cost": 4e-07, "cache_read_input_token_cost": 4e-07, "input_cost_per_audio_token": 3.2e-05, "input_cost_per_image": 5e-06, "input_cost_per_token": 4e-06, "litellm_provider": "openai", - "max_input_tokens": 32000, - "max_output_tokens": 4096, - "max_tokens": 4096, + "max_input_tokens": 128000, + "max_output_tokens": 32000, + "max_tokens": 32000, "mode": "realtime", "output_cost_per_audio_token": 6.4e-05, - "output_cost_per_token": 1.6e-05, + "output_cost_per_token": 2.4e-05, "supported_endpoints": [ "/v1/realtime" ], @@ -30680,8 +30681,8 @@ "input_cost_per_token": 6e-07, "litellm_provider": "openai", "max_input_tokens": 128000, - "max_output_tokens": 4096, - "max_tokens": 4096, + "max_output_tokens": 32000, + "max_tokens": 32000, "mode": "realtime", "output_cost_per_audio_token": 2e-05, "output_cost_per_token": 2.4e-06, @@ -30713,7 +30714,7 @@ "input_cost_per_audio_token": 1e-05, "input_cost_per_token": 6e-07, "litellm_provider": "openai", - "max_input_tokens": 128000, + "max_input_tokens": 32000, "max_output_tokens": 4096, "max_tokens": 4096, "mode": "realtime", @@ -33477,19 +33478,21 @@ "source": "https://mistral.ai/pricing#api-pricing" }, "mistral/magistral-medium-latest": { - "input_cost_per_token": 2e-06, + "cache_read_input_token_cost": 1.5e-07, + "input_cost_per_token": 1.5e-06, "litellm_provider": "mistral", - "max_input_tokens": 40000, - "max_output_tokens": 40000, - "max_tokens": 40000, + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, "mode": "chat", - "output_cost_per_token": 5e-06, - "source": "https://mistral.ai/news/magistral", + "output_cost_per_token": 7.5e-06, + "source": "https://docs.mistral.ai/models/model-cards/mistral-medium-3-5-26-04", "supports_assistant_prefill": true, "supports_function_calling": true, "supports_reasoning": true, "supports_response_schema": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_vision": true }, "mistral/magistral-small-2506": { "deprecation_date": "2025-11-30", @@ -33508,19 +33511,21 @@ "supports_tool_choice": true }, "mistral/magistral-small-latest": { - "input_cost_per_token": 5e-07, + "cache_read_input_token_cost": 1.5e-08, + "input_cost_per_token": 1.5e-07, "litellm_provider": "mistral", - "max_input_tokens": 40000, - "max_output_tokens": 40000, - "max_tokens": 40000, + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, "mode": "chat", - "output_cost_per_token": 1.5e-06, - "source": "https://mistral.ai/pricing#api-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_reasoning": true, "supports_response_schema": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_vision": true }, "mistral/magistral-small-1-2-2509": { "deprecation_date": "2026-07-31", @@ -33652,16 +33657,21 @@ "supports_vision": true }, "mistral/mistral-medium": { - "input_cost_per_token": 2.7e-06, + "cache_read_input_token_cost": 1.5e-07, + "input_cost_per_token": 1.5e-06, "litellm_provider": "mistral", - "max_input_tokens": 32000, - "max_output_tokens": 8191, - "max_tokens": 8191, + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, "mode": "chat", - "output_cost_per_token": 8.1e-06, + "output_cost_per_token": 7.5e-06, + "source": "https://docs.mistral.ai/models/model-cards/mistral-medium-3-5-26-04", "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_reasoning": true, "supports_response_schema": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_vision": true }, "mistral/mistral-medium-2312": { "deprecation_date": "2025-06-16", @@ -45539,6 +45549,26 @@ "mode": "rerank", "output_cost_per_token": 0.0 }, + "voyage/rerank-3": { + "input_cost_per_token": 5e-08, + "litellm_provider": "voyage", + "max_input_tokens": 32000, + "max_output_tokens": 32000, + "max_tokens": 32000, + "mode": "rerank", + "output_cost_per_token": 0.0, + "source": "https://docs.voyageai.com/docs/pricing" + }, + "voyage/rerank-3-lite": { + "input_cost_per_token": 2e-08, + "litellm_provider": "voyage", + "max_input_tokens": 32000, + "max_output_tokens": 32000, + "max_tokens": 32000, + "mode": "rerank", + "output_cost_per_token": 0.0, + "source": "https://docs.voyageai.com/docs/pricing" + }, "voyage/voyage-2": { "input_cost_per_token": 1e-07, "litellm_provider": "voyage", @@ -46790,6 +46820,27 @@ "supports_vision": true, "supports_web_search": true }, + "xai/grok-build-latest": { + "cache_read_input_token_cost": 3e-07, + "cache_read_input_token_cost_above_200k_tokens": 6e-07, + "input_cost_per_token": 2e-06, + "input_cost_per_token_above_200k_tokens": 4e-06, + "litellm_provider": "xai", + "max_input_tokens": 500000, + "max_output_tokens": 500000, + "max_tokens": 500000, + "mode": "chat", + "output_cost_per_token": 6e-06, + "output_cost_per_token_above_200k_tokens": 1.2e-05, + "source": "https://docs.x.ai/developers/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true + }, "xai/grok-4.6": { "cache_read_input_token_cost": 5e-07, "cache_read_input_token_cost_above_200k_tokens": 1e-06, @@ -57110,6 +57161,34 @@ "supports_tool_choice": true, "supports_vision": false }, + "fireworks_ai/accounts/fireworks/models/glm-5p3-flash": { + "cache_read_input_token_cost": 3e-08, + "input_cost_per_token": 1.5e-07, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 1048576, + "max_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 5e-07, + "source": "https://docs.fireworks.ai/serverless/pricing", + "supports_function_calling": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "fireworks_ai/accounts/fireworks/models/inkling": { + "cache_read_input_token_cost": 1.7e-07, + "input_cost_per_token": 1e-06, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 1048576, + "max_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 4.05e-06, + "source": "https://fireworks.ai/models/fireworks/inkling", + "supports_function_calling": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, "fireworks_ai/accounts/fireworks/models/qwen3-embedding-8b": { "input_cost_per_token": 1e-07, "output_cost_per_token": 0.0, From cf4738c3b736edb62f111ab26d899da6a7f9f284 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 1 Sep 2026 12:32:04 -0700 Subject: [PATCH 06/48] style: format s3 vectors transformation and rag endpoints --- litellm/llms/s3_vectors/vector_stores/transformation.py | 4 +++- litellm/proxy/rag_endpoints/endpoints.py | 5 ++++- 2 files changed, 7 insertions(+), 2 deletions(-) diff --git a/litellm/llms/s3_vectors/vector_stores/transformation.py b/litellm/llms/s3_vectors/vector_stores/transformation.py index 3c1b0025d08..733358381fe 100644 --- a/litellm/llms/s3_vectors/vector_stores/transformation.py +++ b/litellm/llms/s3_vectors/vector_stores/transformation.py @@ -68,7 +68,9 @@ class S3VectorsVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM): """Return the router iff it serves ``embedding_model`` as a deployment.""" if router is None: return None - model_list: Final = [dict(m) for m in (router.get_model_list() or ())] # mutable-ok: resolve_embedding_router requires list[dict] + model_list: Final = [ + dict(m) for m in (router.get_model_list() or ()) + ] # mutable-ok: resolve_embedding_router requires list[dict] return resolve_embedding_router(embedding_model=embedding_model, llm_router=router, llm_model_list=model_list) def transform_search_vector_store_request( diff --git a/litellm/proxy/rag_endpoints/endpoints.py b/litellm/proxy/rag_endpoints/endpoints.py index 31998d2040e..5c392c30018 100644 --- a/litellm/proxy/rag_endpoints/endpoints.py +++ b/litellm/proxy/rag_endpoints/endpoints.py @@ -717,7 +717,10 @@ async def rag_query( vector_store_id=retrieval_config["vector_store_id"], user_api_key_dict=user_api_key_dict, ) - merged_retrieval_config: Final = {**store_data, **retrieval_config} # mutable-ok: litellm.aquery requires a plain dict payload + merged_retrieval_config: Final = { + **store_data, + **retrieval_config, + } # mutable-ok: litellm.aquery requires a plain dict payload # Add litellm data request_data: dict[str, object] = {} From d3dab8e294b06badb2d34f31ce6aade9380a49ea Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 1 Sep 2026 12:46:04 -0700 Subject: [PATCH 07/48] fix(rerank): map provider errors with the resolved provider on sync and async paths --- litellm/rerank_api/main.py | 22 ++++++-- tests/test_litellm/rerank_api/test_main.py | 61 ++++++++++++++++++++++ 2 files changed, 80 insertions(+), 3 deletions(-) diff --git a/litellm/rerank_api/main.py b/litellm/rerank_api/main.py index c8f7842aebf..597d1cfb863 100644 --- a/litellm/rerank_api/main.py +++ b/litellm/rerank_api/main.py @@ -43,10 +43,17 @@ async def arerank( """ Async: Reranks a list of documents based on their relevance to the query """ + _custom_llm_provider: str | None = None # rebind-ok: set by the get_llm_provider unpack; read in the except try: loop: Final = asyncio.get_event_loop() kwargs["arerank"] = True + _, _custom_llm_provider, _, _ = litellm.get_llm_provider( # rebind-ok: see pre-declaration above + model=model, + custom_llm_provider=custom_llm_provider, + api_base=kwargs.get("api_base", None), + ) + func: Final = partial( rerank, model, @@ -70,7 +77,11 @@ async def arerank( response = init_response return response except Exception as e: - raise e + raise exception_type( + model=model, + custom_llm_provider=_custom_llm_provider or custom_llm_provider, + original_exception=e, + ) @client @@ -115,6 +126,7 @@ def rerank( model_info: Final = kwargs.get("model_info", None) user: Final = kwargs.get("user", None) client: Final = kwargs.get("client", None) + _custom_llm_provider: str | None = None # rebind-ok: set by the get_llm_provider unpack; read in the except try: _is_async: Final = kwargs.pop("arerank", False) is True optional_params: Final = GenericLiteLLMParams(**kwargs) @@ -127,7 +139,7 @@ def rerank( ( model, - _custom_llm_provider, + _custom_llm_provider, # rebind-ok: see pre-declaration above dynamic_api_key, dynamic_api_base, ) = litellm.get_llm_provider( @@ -538,4 +550,8 @@ def rerank( return response except Exception as e: verbose_logger.error("Error in rerank: %s", e) - raise exception_type(model=model, custom_llm_provider=custom_llm_provider, original_exception=e) + raise exception_type( + model=model, + custom_llm_provider=_custom_llm_provider or custom_llm_provider, + original_exception=e, + ) diff --git a/tests/test_litellm/rerank_api/test_main.py b/tests/test_litellm/rerank_api/test_main.py index 587be59c550..62149c742d6 100644 --- a/tests/test_litellm/rerank_api/test_main.py +++ b/tests/test_litellm/rerank_api/test_main.py @@ -111,6 +111,67 @@ def test_together_rerank_honors_api_base(respx_mock: respx.MockRouter): assert mock_route.calls[0].request.headers["authorization"] == "Bearer fake-together-key" +DASHSCOPE_404_BODY = { + "error": { + "message": "The model `does-not-exist` does not exist or you do not have access to it.", + "type": "invalid_request_error", + "param": None, + "code": "model_not_found", + }, + "request_id": "mock-request-id", +} + + +def test_rerank_error_names_provider_and_keeps_body(respx_mock: respx.MockRouter, monkeypatch): + """Regression for the rerank error path mapping with the unresolved provider param: + a provider 404 surfaced as 'None - ' instead of naming the provider and its error body.""" + monkeypatch.delenv("DASHSCOPE_API_BASE", raising=False) + monkeypatch.delenv("DASHSCOPE_API_BASE_RERANK", raising=False) + + mock_route = respx_mock.post("https://dashscope.example/v1/reranks") + mock_route.return_value = httpx.Response(404, json=DASHSCOPE_404_BODY) + + with pytest.raises(litellm.NotFoundError) as exc_info: + litellm.rerank( + model="dashscope/does-not-exist", + query=MARKER_QUERY, + documents=[MARKER_DOC], + api_key="fake-dashscope-key", + api_base="https://dashscope.example/v1", + ) + + assert mock_route.called + assert "DashscopeException" in str(exc_info.value) + assert "does not exist or you do not have access to it" in str(exc_info.value) + assert "None - " not in str(exc_info.value) + + +@pytest.mark.asyncio +async def test_arerank_error_is_mapped_to_litellm_exception(respx_mock: respx.MockRouter, monkeypatch): + """Regression for arerank's bare re-raise: provider errors escaped as raw + provider exception classes instead of the mapped litellm exception contract.""" + monkeypatch.delenv("DASHSCOPE_API_BASE", raising=False) + monkeypatch.delenv("DASHSCOPE_API_BASE_RERANK", raising=False) + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + + mock_route = respx_mock.post("https://dashscope.example/v1/reranks") + mock_route.return_value = httpx.Response(404, json=DASHSCOPE_404_BODY) + + with pytest.raises(litellm.NotFoundError) as exc_info: + await litellm.arerank( + model="dashscope/does-not-exist", + query=MARKER_QUERY, + documents=[MARKER_DOC], + api_key="fake-dashscope-key", + api_base="https://dashscope.example/v1", + ) + + assert mock_route.called + assert "DashscopeException" in str(exc_info.value) + assert "does not exist or you do not have access to it" in str(exc_info.value) + assert "None - " not in str(exc_info.value) + + @pytest.mark.asyncio async def test_together_rerank_async_honors_env_api_base(respx_mock: respx.MockRouter, monkeypatch): """Regression: TOGETHER_AI_API_BASE was honored by chat but ignored by rerank.""" From 3914de24eff6c0f3deda46ec2f2217cb0305d916 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 1 Sep 2026 12:52:45 -0700 Subject: [PATCH 08/48] fix(router): route model-less sync vector store calls to the SDK _generic_api_call_with_fallbacks requires a model, so sync vector_store_search and vector_store_create raised a TypeError whenever the call carried no model. Model-less calls now go directly to the SDK function, with the router injected for search, matching the async wrapper's behavior --- litellm/router.py | 19 +++++++---- tests/test_litellm/test_router.py | 55 +++++++++++++++++++++++++++++++ 2 files changed, 67 insertions(+), 7 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index d8769525815..96d5b1e1488 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -6299,8 +6299,6 @@ class Router: "responses", "generate_content", "generate_content_stream", - "vector_store_search", - "vector_store_create", "ocr", "search", "video_generation", @@ -6324,6 +6322,8 @@ class Router: return sync_wrapper if call_type in ( + "vector_store_search", + "vector_store_create", "vector_store_retrieve", "vector_store_list", "vector_store_update", @@ -6335,11 +6335,16 @@ class Router: client: object | None = None, **kwargs, ): - if custom_llm_provider and "custom_llm_provider" not in kwargs: - kwargs["custom_llm_provider"] = custom_llm_provider - if kwargs.get("model"): - return self._generic_api_call_with_fallbacks(original_function=original_function, **kwargs) - return original_function(**kwargs) + provider_kwargs: Final = ( + MappingProxyType({**kwargs, "custom_llm_provider": custom_llm_provider}) + if custom_llm_provider and "custom_llm_provider" not in kwargs + else MappingProxyType(kwargs) + ) + if provider_kwargs.get("model"): + return self._generic_api_call_with_fallbacks(original_function=original_function, **provider_kwargs) + if call_type == "vector_store_search": + return original_function(**MappingProxyType({**provider_kwargs, "router": self})) + return original_function(**provider_kwargs) return vector_store_sync_wrapper diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index c222fab79f0..3abfe8ce522 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -7567,6 +7567,61 @@ async def test_avector_store_create_does_not_inject_router(): assert "router" not in mock_acreate.await_args.kwargs +def test_vector_store_search_injects_router(): + """ + Sync parity for the router injection: router.vector_store_search must pass + the router down to the SDK search call so provider transforms can resolve + router-managed embedding models, same as avector_store_search. + """ + from litellm.types.vector_stores import VectorStoreSearchResponse + + expected_response = VectorStoreSearchResponse( + object="vector_store.search_results.page", search_query="q", data=[] + ) + mock_search = MagicMock(return_value=expected_response) + # Router.__init__ binds search via a local import, so patch the module + # attribute before constructing the Router. + with patch("litellm.vector_stores.main.search", new=mock_search): # test-quality-ok: the SDK call is the only place the injected router kwarg is observable + router = litellm.Router( + model_list=[ + { + "model_name": "gpt-3.5-turbo", + "litellm_params": {"model": "openai/gpt-3.5-turbo", "api_key": "test-key"}, + } + ] + ) + search_response = router.vector_store_search( + vector_store_id="v", query="q", custom_llm_provider="s3_vectors" + ) + + assert search_response is expected_response + mock_search.assert_called_once() + assert mock_search.call_args.kwargs["router"] is router + assert mock_search.call_args.kwargs["custom_llm_provider"] == "s3_vectors" + + +def test_vector_store_create_does_not_inject_router(): + """The sync create path must keep calling the SDK without a router kwarg.""" + expected_response = {"id": "vs_1", "object": "vector_store"} + mock_create = MagicMock(return_value=expected_response) + # Router.__init__ binds create via a local import, so patch the module + # attribute before constructing the Router. + with patch("litellm.vector_stores.main.create", new=mock_create): # test-quality-ok: the SDK call is the only place a leaked router kwarg would surface + router = litellm.Router( + model_list=[ + { + "model_name": "gpt-3.5-turbo", + "litellm_params": {"model": "openai/gpt-3.5-turbo", "api_key": "test-key"}, + } + ] + ) + create_response = router.vector_store_create(custom_llm_provider="openai") + + assert create_response is expected_response + mock_create.assert_called_once() + assert "router" not in mock_create.call_args.kwargs + + class TestPreRoutingStrategyRegistryLifecycle: """ Regression tests: a deployment leaving the model_list must release the From babe7816ada8d622be993b542fc2512037d2466f Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 1 Sep 2026 12:52:45 -0700 Subject: [PATCH 09/48] fix(rag): store-wins merge, single lookup, allowlisted search params rag_query reuses the store resolved during authorization instead of a second registry lookup, merges registry data store-wins so callers cannot override a managed store's provider or credentials, and logs ids instead of the merged config, which can carry resolved credentials. aquery forwards only allowlisted retrieval_config keys to vector store search, keeping caller-supplied connection overrides like api_base and api_key away from the search call --- litellm/proxy/rag_endpoints/endpoints.py | 53 ++++++++---- .../proxy/vector_store_endpoints/endpoints.py | 86 ++++++++++--------- .../management_endpoints.py | 4 +- litellm/rag/main.py | 25 ++++-- .../proxy/rag_endpoints/test_rag_endpoints.py | 26 +++--- tests/test_litellm/rag/test_main.py | 39 +++++++++ 6 files changed, 149 insertions(+), 84 deletions(-) diff --git a/litellm/proxy/rag_endpoints/endpoints.py b/litellm/proxy/rag_endpoints/endpoints.py index 5c392c30018..d2c7d6f93ee 100644 --- a/litellm/proxy/rag_endpoints/endpoints.py +++ b/litellm/proxy/rag_endpoints/endpoints.py @@ -9,6 +9,7 @@ Provides: import base64 import json from collections.abc import Mapping +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final import orjson @@ -19,6 +20,9 @@ from starlette.datastructures import UploadFile import litellm from litellm._logging import verbose_proxy_logger from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH +from litellm.integrations.vector_store_integrations.vector_store_pre_call_hook import ( + LiteLLM_ManagedVectorStore, +) from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper from litellm.proxy._types import * from litellm.proxy.auth.auth_utils import is_request_body_safe @@ -37,7 +41,7 @@ from litellm.proxy.rag_endpoints.upload_security import ( validate_upload, ) from litellm.proxy.vector_store_endpoints.endpoints import ( - _update_request_data_with_litellm_managed_vector_store_registry, # pyright: ignore[reportPrivateUsage] # shared registry-merge helper used by the direct search endpoint + build_request_data_from_managed_vector_store, ) from litellm.proxy.vector_store_endpoints.utils import ( assert_user_can_access_vector_store_id, @@ -123,12 +127,21 @@ def _collect_vector_store_ids_from_payload(payload: object) -> set[str]: async def _authorize_nested_vector_store_ids( payload: object, user_api_key_dict: UserAPIKeyAuth, -) -> None: - for vector_store_id in sorted(_collect_vector_store_ids_from_payload(payload)): - await assert_user_can_access_vector_store_id( - vector_store_id=vector_store_id, - user_api_key_dict=user_api_key_dict, - ) +) -> Mapping[str, LiteLLM_ManagedVectorStore]: + """Authorize every nested vector store id and return the managed stores it resolved.""" + return MappingProxyType( + { + vector_store_id: store + for vector_store_id in sorted(_collect_vector_store_ids_from_payload(payload)) + if ( + store := await assert_user_can_access_vector_store_id( + vector_store_id=vector_store_id, + user_api_key_dict=user_api_key_dict, + ) + ) + is not None + } + ) def _build_file_metadata_entry( @@ -703,23 +716,24 @@ async def rag_query( status_code=400, detail={"error": "retrieval_config must contain 'vector_store_id'"}, ) - await _authorize_nested_vector_store_ids( + resolved_stores: Final = await _authorize_nested_vector_store_ids( payload=retrieval_config, user_api_key_dict=user_api_key_dict, ) # Merge litellm-managed vector store params (provider, region, embedding - # model, credentials, ...) from the registry — same source the direct - # /vector_stores/{id}/search endpoint uses. User-supplied - # retrieval_config keys win on conflict. - store_data: Final = await _update_request_data_with_litellm_managed_vector_store_registry( - data={}, # mutable-ok: the helper mutates and returns the seed dict - vector_store_id=retrieval_config["vector_store_id"], - user_api_key_dict=user_api_key_dict, + # model, credentials, ...) from the registry: the same source the direct + # /vector_stores/{id}/search endpoint uses. Store-managed keys win on + # conflict so callers cannot override the store's provider or credentials. + managed_store: Final = resolved_stores.get(retrieval_config["vector_store_id"]) + store_data: Final = ( + await build_request_data_from_managed_vector_store(managed_store) + if managed_store is not None + else MappingProxyType({}) ) merged_retrieval_config: Final = { - **store_data, **retrieval_config, + **store_data, } # mutable-ok: litellm.aquery requires a plain dict payload # Add litellm data @@ -733,7 +747,12 @@ async def rag_query( proxy_config=proxy_config, ) - verbose_proxy_logger.debug("RAG Query - model: %s, retrieval_config: %s", model, merged_retrieval_config) + verbose_proxy_logger.debug( + "RAG Query - model: %s, vector_store_id: %s, custom_llm_provider: %s", + model, + retrieval_config["vector_store_id"], + merged_retrieval_config.get("custom_llm_provider"), + ) # Call query response: Final = await litellm.aquery( diff --git a/litellm/proxy/vector_store_endpoints/endpoints.py b/litellm/proxy/vector_store_endpoints/endpoints.py index a59d7a277cc..3fc6749f18a 100644 --- a/litellm/proxy/vector_store_endpoints/endpoints.py +++ b/litellm/proxy/vector_store_endpoints/endpoints.py @@ -1,3 +1,5 @@ +from collections.abc import Mapping +from types import MappingProxyType from typing import ( Annotated, Any, # noqa: TID251 # jsonify_object in proxy/utils.py is annotated with a bare dict @@ -32,6 +34,41 @@ router: Final = APIRouter() ######################################################## +async def build_request_data_from_managed_vector_store( + vector_store: LiteLLM_ManagedVectorStore, +) -> Mapping[str, object]: + """ + Build request params (provider, credential ref, litellm_params) from an + already-resolved managed vector store. + + ``litellm_embedding_config`` is resolved here, at request-handling time, + instead of at row-creation time: the resolved api_key/api_base/api_version + lives only in the returned per-request mapping and is never persisted back + to the registry cache. Legacy rows that already carry a resolved + (cleartext) config skip the lookup and pass through unchanged. + """ + top_level: Final = MappingProxyType( + { + key: vector_store.get(key) + for key in ("custom_llm_provider", "litellm_credential_name") + if key in vector_store + } + ) + litellm_params: Final = vector_store.get("litellm_params") or MappingProxyType({}) + embedding_model: Final = litellm_params.get("litellm_embedding_model") + if not embedding_model or litellm_params.get("litellm_embedding_config"): + return MappingProxyType({**top_level, **litellm_params}) + + from litellm.proxy.proxy_server import prisma_client + + resolved_config: Final = await _resolve_embedding_config( + embedding_model=embedding_model, prisma_client=prisma_client + ) + if not resolved_config: + return MappingProxyType({**top_level, **litellm_params}) + return MappingProxyType({**top_level, **litellm_params, "litellm_embedding_config": resolved_config}) + + async def _update_request_data_with_litellm_managed_vector_store_registry( data: dict, vector_store_id: str, @@ -51,47 +88,14 @@ async def _update_request_data_with_litellm_managed_vector_store_registry( vector_store_to_run: Final[LiteLLM_ManagedVectorStore | None] = await get_litellm_managed_vector_store( vector_store_id=vector_store_id ) - if vector_store_to_run is not None: - if user_api_key_dict is not None: - await assert_user_can_access_vector_store( - vector_store=vector_store_to_run, - user_api_key_dict=user_api_key_dict, - ) - - if "custom_llm_provider" in vector_store_to_run: - data["custom_llm_provider"] = vector_store_to_run.get("custom_llm_provider") - - if "litellm_credential_name" in vector_store_to_run: - data["litellm_credential_name"] = vector_store_to_run.get("litellm_credential_name") - - if "litellm_params" in vector_store_to_run: - litellm_params = vector_store_to_run.get("litellm_params", {}) or {} - # Resolve ``litellm_embedding_config`` here, at request-handling - # time, instead of at row-creation time. The resolved - # ``api_key`` / ``api_base`` / ``api_version`` lives only in - # this per-request ``data`` dict and is never persisted. - # Legacy rows that already carry a resolved (cleartext) - # ``litellm_embedding_config`` skip the lookup and pass through - # unchanged so the embed call keeps working. - embedding_model: Final = litellm_params.get("litellm_embedding_model") - if embedding_model and not litellm_params.get("litellm_embedding_config"): - from litellm.proxy.proxy_server import prisma_client - - resolved_config: Final = await _resolve_embedding_config( - embedding_model=embedding_model, prisma_client=prisma_client - ) - if resolved_config: - # Build a fresh dict via spread instead of mutating - # ``litellm_params`` in place — the registry hands back - # a reference to its cached object, so an in-place - # update would persist the resolved cleartext into the - # in-memory cache for the lifetime of the process. - litellm_params = { - **litellm_params, - "litellm_embedding_config": resolved_config, - } - data.update(litellm_params) - return data + if vector_store_to_run is None: + return data + if user_api_key_dict is not None: + await assert_user_can_access_vector_store( + vector_store=vector_store_to_run, + user_api_key_dict=user_api_key_dict, + ) + return {**data, **(await build_request_data_from_managed_vector_store(vector_store_to_run))} @router.post( diff --git a/litellm/proxy/vector_store_endpoints/management_endpoints.py b/litellm/proxy/vector_store_endpoints/management_endpoints.py index 183a03cc13c..244798ba05e 100644 --- a/litellm/proxy/vector_store_endpoints/management_endpoints.py +++ b/litellm/proxy/vector_store_endpoints/management_endpoints.py @@ -470,7 +470,7 @@ async def create_vector_store_in_db( # exposed every env-stored embedding-model credential on the # ``/vector_store/{new,info,update,list}`` responses. Keep the user's # raw ``litellm_embedding_model`` reference; resolution now happens in - # ``_update_request_data_with_litellm_managed_vector_store_registry`` + # ``build_request_data_from_managed_vector_store`` # at request-handling time so the cleartext config exists only in # per-request memory and never reaches the database. if litellm_params: @@ -864,7 +864,7 @@ async def update_vector_store( # embedding-config auto-resolve previously persisted cleartext # credentials into the row; resolution now happens at request- # handling time in - # ``_update_request_data_with_litellm_managed_vector_store_registry`` + # ``build_request_data_from_managed_vector_store`` # so this row only ever stores the user-supplied # ``litellm_embedding_model`` reference. if "litellm_params" in update_data: diff --git a/litellm/rag/main.py b/litellm/rag/main.py index bd6788b3a1b..94bfc305a6a 100644 --- a/litellm/rag/main.py +++ b/litellm/rag/main.py @@ -51,12 +51,19 @@ INGESTION_REGISTRY: Final[dict[str, type[BaseRAGIngestion]]] = { "vertex_ai": VertexAIRAGIngestion, } -# retrieval_config keys consumed by the query pipeline itself; everything else is -# forwarded to vector_stores.asearch as provider-specific params (e.g. -# aws_region_name, embedding_model, vector_bucket_name for S3 Vectors). -# `filters`/`retrieval_filter` are reserved for the explicit filter param. -_CONSUMED_RETRIEVAL_CONFIG_KEYS: Final = frozenset( - {"vector_store_id", "custom_llm_provider", "top_k", "filters", "retrieval_filter"} +# Only these retrieval_config keys are forwarded to vector_stores.asearch as +# provider-specific params. The explicit allowlist keeps caller-controlled +# connection overrides (api_base, api_key, ...) away from the search call, +# where they could redirect store credentials to an attacker-chosen host. +_FORWARDABLE_RETRIEVAL_CONFIG_KEYS: Final = frozenset( + { + "aws_region_name", + "vector_bucket_name", + "embedding_model", + "litellm_embedding_model", + "litellm_embedding_config", + "litellm_credential_name", + } ) @@ -233,10 +240,10 @@ async def _execute_query_pipeline( raise ValueError("No query found in messages for RAG query") # 2. Search vector store - # Forward provider-specific retrieval_config extras (region, embedding model, - # bucket, credentials refs, ...) to the search call; kwargs win on conflict. + # Forward allowlisted provider retrieval_config extras (region, embedding + # model, bucket, credential refs) to the search call; kwargs win on conflict. provider_search_params: Final = MappingProxyType( - {k: v for k, v in retrieval_config.items() if k not in _CONSUMED_RETRIEVAL_CONFIG_KEYS} + {k: v for k, v in retrieval_config.items() if k in _FORWARDABLE_RETRIEVAL_CONFIG_KEYS} ) forwarded_search_params: Final = MappingProxyType({**provider_search_params, **kwargs}) with _suppressed_sub_call_billing(): diff --git a/tests/test_litellm/proxy/rag_endpoints/test_rag_endpoints.py b/tests/test_litellm/proxy/rag_endpoints/test_rag_endpoints.py index 5561ee1e6ae..342a4535b21 100644 --- a/tests/test_litellm/proxy/rag_endpoints/test_rag_endpoints.py +++ b/tests/test_litellm/proxy/rag_endpoints/test_rag_endpoints.py @@ -357,12 +357,9 @@ def test_rag_query_merges_managed_store_params(client_internal_user): "litellm.proxy.rag_endpoints.endpoints.litellm.aquery", new_callable=AsyncMock, return_value=mock_response, - ) as mock_aquery, patch.object(litellm, "vector_store_registry", mock_registry), patch( # test-quality-ok: seeds the managed-store registry the merge under test reads and stubs the access assert covered by auth tests - "litellm.proxy.rag_endpoints.endpoints.assert_user_can_access_vector_store_id", - new=AsyncMock(), - ), patch( # test-quality-ok: stubs the direct-endpoint access assert covered by auth tests - "litellm.proxy.vector_store_endpoints.endpoints.assert_user_can_access_vector_store", - new=AsyncMock(), + ) as mock_aquery, patch.object(litellm, "vector_store_registry", mock_registry), patch( # test-quality-ok: seeds the managed-store registry the merge under test reads and grants access so real store resolution runs + "litellm.proxy.vector_store_endpoints.utils.can_user_access_vector_store", + new=AsyncMock(return_value=True), ): response = client_internal_user.post( "/v1/rag/query", @@ -383,8 +380,8 @@ def test_rag_query_merges_managed_store_params(client_internal_user): assert forwarded_config["vector_bucket_name"] == "bkt" -def test_rag_query_user_retrieval_config_wins_over_store(client_internal_user): - """User-supplied retrieval_config keys must win over registry values.""" +def test_rag_query_store_params_win_over_user_retrieval_config(client_internal_user): + """Registry values must win over user-supplied retrieval_config keys so callers cannot override store credentials.""" import litellm from litellm.types.utils import ModelResponse @@ -406,12 +403,9 @@ def test_rag_query_user_retrieval_config_wins_over_store(client_internal_user): "litellm.proxy.rag_endpoints.endpoints.litellm.aquery", new_callable=AsyncMock, return_value=mock_response, - ) as mock_aquery, patch.object(litellm, "vector_store_registry", mock_registry), patch( # test-quality-ok: seeds the managed-store registry the merge under test reads and stubs the access assert covered by auth tests - "litellm.proxy.rag_endpoints.endpoints.assert_user_can_access_vector_store_id", - new=AsyncMock(), - ), patch( # test-quality-ok: stubs the direct-endpoint access assert covered by auth tests - "litellm.proxy.vector_store_endpoints.endpoints.assert_user_can_access_vector_store", - new=AsyncMock(), + ) as mock_aquery, patch.object(litellm, "vector_store_registry", mock_registry), patch( # test-quality-ok: seeds the managed-store registry the merge under test reads and grants access so real store resolution runs + "litellm.proxy.vector_store_endpoints.utils.can_user_access_vector_store", + new=AsyncMock(return_value=True), ): response = client_internal_user.post( "/v1/rag/query", @@ -424,7 +418,9 @@ def test_rag_query_user_retrieval_config_wins_over_store(client_internal_user): assert response.status_code == 200, response.json() forwarded_config = mock_aquery.await_args.kwargs["retrieval_config"] - assert forwarded_config["aws_region_name"] == "us-east-1" + assert forwarded_config["aws_region_name"] == "eu-west-1" + + EICAR = r"X5O!P%@AP[4\PZX54(P^)7CC)7}$EICAR-STANDARD-ANTIVIRUS-TEST-FILE!$H+H*" INGEST_REQUEST = '{"ingest_options":{"vector_store":{"custom_llm_provider":"openai"}}}' diff --git a/tests/test_litellm/rag/test_main.py b/tests/test_litellm/rag/test_main.py index fdcdf342eea..51d03544910 100644 --- a/tests/test_litellm/rag/test_main.py +++ b/tests/test_litellm/rag/test_main.py @@ -349,6 +349,45 @@ async def test_aquery_minimal_retrieval_config_forwards_no_extras(): assert not (leaked & set(search_kwargs.keys())) +@pytest.mark.asyncio +async def test_aquery_does_not_forward_connection_override_keys_to_search(): + """ + Only allowlisted retrieval_config keys may reach the vector store search + call. Caller-controlled connection overrides (api_base, api_key, arbitrary + extras) must be dropped, otherwise a caller could redirect store + credentials to an attacker-chosen host. + """ + from unittest.mock import AsyncMock + + from litellm.types.vector_stores import VectorStoreSearchResponse + + fake_search = AsyncMock( + return_value=VectorStoreSearchResponse( + object="vector_store.search_results.page", search_query="q", data=[] + ) + ) + with patch("litellm.vector_stores.asearch", new=fake_search): # test-quality-ok: asearch is the boundary the forwarding contract under test targets + await litellm.aquery( + model="gpt-4o-mini", + messages=[{"role": "user", "content": "hello"}], + retrieval_config={ + "vector_store_id": "bkt:idx", + "custom_llm_provider": "s3_vectors", + "aws_region_name": "eu-west-1", + "api_base": "https://attacker.example.com", + "api_key": "attacker-key", + "arbitrary_extra": "nope", + }, + mock_response="hi", + ) + + fake_search.assert_awaited_once() + search_kwargs = fake_search.await_args.kwargs + assert search_kwargs["aws_region_name"] == "eu-west-1" + blocked = {"api_base", "api_key", "arbitrary_extra"} + assert not (blocked & set(search_kwargs.keys())) + + def test_rag_call_types_are_registered(): """ query/aquery/ingest/aingest are @client-decorated entry points, so their From 8b5ae3da9d49e08c4f6c8ca22d6f59c64c6bfae5 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 1 Sep 2026 13:15:51 -0700 Subject: [PATCH 10/48] test(vector_stores): package the suite dir to avoid test_main basename collision --- tests/test_litellm/vector_stores/__init__.py | 0 1 file changed, 0 insertions(+), 0 deletions(-) create mode 100644 tests/test_litellm/vector_stores/__init__.py diff --git a/tests/test_litellm/vector_stores/__init__.py b/tests/test_litellm/vector_stores/__init__.py new file mode 100644 index 00000000000..e69de29bb2d From 8b0441a628c01f0cd6caa10176ae887c06d75fa7 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 1 Sep 2026 13:29:25 -0700 Subject: [PATCH 11/48] fix(vector_stores): block caller-supplied embedding selection params on query surfaces --- litellm/proxy/rag_endpoints/endpoints.py | 2 ++ .../proxy/vector_store_endpoints/endpoints.py | 24 ++++++++++++++ .../proxy/rag_endpoints/test_rag_endpoints.py | 24 ++++++++++++++ .../test_vector_store_endpoints.py | 32 +++++++++++++++++++ 4 files changed, 82 insertions(+) diff --git a/litellm/proxy/rag_endpoints/endpoints.py b/litellm/proxy/rag_endpoints/endpoints.py index d2c7d6f93ee..0ab7d99e4e4 100644 --- a/litellm/proxy/rag_endpoints/endpoints.py +++ b/litellm/proxy/rag_endpoints/endpoints.py @@ -42,6 +42,7 @@ from litellm.proxy.rag_endpoints.upload_security import ( ) from litellm.proxy.vector_store_endpoints.endpoints import ( build_request_data_from_managed_vector_store, + reject_caller_embedding_selection_params, ) from litellm.proxy.vector_store_endpoints.utils import ( assert_user_can_access_vector_store_id, @@ -716,6 +717,7 @@ async def rag_query( status_code=400, detail={"error": "retrieval_config must contain 'vector_store_id'"}, ) + reject_caller_embedding_selection_params(payload=retrieval_config, source="retrieval_config") resolved_stores: Final = await _authorize_nested_vector_store_ids( payload=retrieval_config, user_api_key_dict=user_api_key_dict, diff --git a/litellm/proxy/vector_store_endpoints/endpoints.py b/litellm/proxy/vector_store_endpoints/endpoints.py index 3fc6749f18a..7d64e648e08 100644 --- a/litellm/proxy/vector_store_endpoints/endpoints.py +++ b/litellm/proxy/vector_store_endpoints/endpoints.py @@ -29,6 +29,29 @@ from litellm.types.vector_stores import IndexCreateRequest, IndexListResponse from litellm.vector_stores.vector_store_registry import VectorStoreIndexRegistry router: Final = APIRouter() + +BLOCKED_QUERY_EMBEDDING_SELECTION_PARAMS: Final = frozenset( + { + "embedding_model", + "litellm_embedding_model", + "litellm_embedding_config", + "litellm_credential_name", + } +) + + +def reject_caller_embedding_selection_params(payload: Mapping[str, object], source: str) -> None: + blocked: Final = sorted(BLOCKED_QUERY_EMBEDDING_SELECTION_PARAMS & payload.keys()) + if blocked: + raise HTTPException( + status_code=400, + detail={ + "error": f"'{blocked[0]}' cannot be set in {source}. " + "Embedding configuration comes from the vector store's server-side registration." + }, + ) + + ######################################################## # OpenAI Compatible Endpoints ######################################################## @@ -134,6 +157,7 @@ async def vector_store_search( ) data = await _read_request_body(request=request) + reject_caller_embedding_selection_params(payload=data, source="the search request body") data["vector_store_id"] = vector_store_id # Check for legacy vector store registry (non-managed vector stores) diff --git a/tests/test_litellm/proxy/rag_endpoints/test_rag_endpoints.py b/tests/test_litellm/proxy/rag_endpoints/test_rag_endpoints.py index 342a4535b21..0085b6ebd36 100644 --- a/tests/test_litellm/proxy/rag_endpoints/test_rag_endpoints.py +++ b/tests/test_litellm/proxy/rag_endpoints/test_rag_endpoints.py @@ -421,6 +421,30 @@ def test_rag_query_store_params_win_over_user_retrieval_config(client_internal_u assert forwarded_config["aws_region_name"] == "eu-west-1" +@pytest.mark.parametrize( + "blocked_key", + ["embedding_model", "litellm_embedding_model", "litellm_embedding_config", "litellm_credential_name"], +) +def test_rag_query_rejects_caller_embedding_selection_params(client_internal_user, blocked_key): + """ + Regression: a caller must not pick the embedding model or credential used at + search time. Those resolve through the Router with the proxy's credentials, + bypassing the key's model permissions, so they may only come from the + managed store's server-side registration. + """ + response = client_internal_user.post( + "/v1/rag/query", + json={ + "model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "hello"}], + "retrieval_config": {"vector_store_id": "s3-store", blocked_key: "attacker-choice"}, + }, + ) + + assert response.status_code == 400, response.json() + assert blocked_key in str(response.json()) + + EICAR = r"X5O!P%@AP[4\PZX54(P^)7CC)7}$EICAR-STANDARD-ANTIVIRUS-TEST-FILE!$H+H*" INGEST_REQUEST = '{"ingest_options":{"vector_store":{"custom_llm_provider":"openai"}}}' diff --git a/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_endpoints.py b/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_endpoints.py index eae6f90863a..45a0221c8a6 100644 --- a/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_endpoints.py +++ b/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_endpoints.py @@ -3158,3 +3158,35 @@ class TestAzureAIAnalyzeNamedIndexClassification: user_api_key_dict=self._team_member("analyze", ["read"]), ) assert result is True + + +@pytest.mark.parametrize( + "blocked_key", + ["embedding_model", "litellm_embedding_model", "litellm_embedding_config", "litellm_credential_name"], +) +def test_vector_store_search_rejects_caller_embedding_selection_params(blocked_key): + """ + Regression: the search request body must not pick the embedding model or + credential used to embed the query. Those resolve through the Router with + the proxy's credentials, bypassing the key's model permissions, so they may + only come from the managed store's server-side registration. + """ + from fastapi.testclient import TestClient + + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + from litellm.proxy.proxy_server import app + + mock_auth = UserAPIKeyAuth(user_id="test_internal_user", user_role=LitellmUserRoles.INTERNAL_USER.value) + original_overrides = app.dependency_overrides.copy() + app.dependency_overrides[user_api_key_auth] = lambda: mock_auth + try: + client = TestClient(app) + response = client.post( + "/v1/vector_stores/s3-store/search", + json={"query": "hello", blocked_key: "attacker-choice"}, + ) + finally: + app.dependency_overrides = original_overrides + + assert response.status_code == 400, response.json() + assert blocked_key in str(response.json()) From 1eea8e283157d3c92be36749b2713607aebc9786 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 1 Sep 2026 13:54:34 -0700 Subject: [PATCH 12/48] fix(deps): raise the tornado floor to 6.5.8 for GHSA-8423-8fgw-73vq and GHSA-wwv5-g3v4-889x --- pyproject.toml | 2 +- uv.lock | 26 +++++++++++++------------- 2 files changed, 14 insertions(+), 14 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 2866e27e84c..d0f14722acd 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -292,7 +292,7 @@ exclude = [ [tool.uv] constraint-dependencies = [ - "tornado>=6.5.6", + "tornado>=6.5.8", "aiohttp>=3.14.2,<4.0", "packaging>=24.0", "soupsieve>=2.8.4", diff --git a/uv.lock b/uv.lock index 27be919eea1..8bac024d49e 100644 --- a/uv.lock +++ b/uv.lock @@ -10,7 +10,7 @@ resolution-markers = [ ] [options] -exclude-newer = "2026-08-29T17:58:57.633306Z" +exclude-newer = "2026-08-29T20:52:40.322465Z" exclude-newer-span = "P3D" [manifest] @@ -25,7 +25,7 @@ constraints = [ { name = "packaging", specifier = ">=24.0" }, { name = "setuptools", specifier = ">=83.0.0" }, { name = "soupsieve", specifier = ">=2.8.4" }, - { name = "tornado", specifier = ">=6.5.6" }, + { name = "tornado", specifier = ">=6.5.8" }, ] overrides = [ { name = "cryptography", specifier = ">=50.0.0,<51.0" }, @@ -9441,19 +9441,19 @@ wheels = [ [[package]] name = "tornado" -version = "6.5.7" +version = "6.5.8" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/64/24/95ec527ad67b76d59299e5465b3935d05e4294b7e0290a3924b7487df30b/tornado-6.5.7.tar.gz", hash = "sha256:66c513a76cda70d53907bc27cf1447557699c2e95aa48ba27a442ff61c3ddfc2", size = 519252, upload-time = "2026-06-08T17:34:51.232Z" } +sdist = { url = "https://files.pythonhosted.org/packages/10/d3/343e5bb989d6515b1646cf3d40135d73f3d5e45339bded401b56cdac24dd/tornado-6.5.8.tar.gz", hash = "sha256:9452e1b208a8bd771e2cb1f2ff564985b9b214bdebbe622793e1799e0a6bd23f", size = 520493, upload-time = "2026-08-07T02:12:42.971Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/02/dc/c7043cab6fed8ae159fc1923ce829ada35c4dbd797d408a43858ffaf9639/tornado-6.5.7-cp39-abi3-macosx_10_9_universal2.whl", hash = "sha256:148b2eb15c2c765a50796172c1e499649b35f30d2e3c3d3e15913cfa56bfb163", size = 448543, upload-time = "2026-06-08T17:34:38.052Z" }, - { url = "https://files.pythonhosted.org/packages/92/4f/090b1431e5a43df696feceffc268c5383cc079ecb5f08ce58f917109aafe/tornado-6.5.7-cp39-abi3-macosx_10_9_x86_64.whl", hash = "sha256:9da38de27f1da3b78a966f0dae12b5a1ea9afe72ca805d84ff06508272ddf100", size = 446707, upload-time = "2026-06-08T17:34:39.594Z" }, - { url = "https://files.pythonhosted.org/packages/37/d8/ef374952fd5da67d4463122c2b8e5a96536ec10b4b339254c6dcde81d01c/tornado-6.5.7-cp39-abi3-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:8d759e71906ee783f8867b93bf26a265743da4c1e2f4a018464c1ba019862972", size = 449774, upload-time = "2026-06-08T17:34:41.204Z" }, - { url = "https://files.pythonhosted.org/packages/35/37/d434c73f4c6e014b745b9b37085f34f40c022f007efff3d7fe65991899f3/tornado-6.5.7-cp39-abi3-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:8a46347a18f23fb92b396beebe0fb78f61dda0cc302445202c16203d8a18848b", size = 450745, upload-time = "2026-06-08T17:34:42.531Z" }, - { url = "https://files.pythonhosted.org/packages/b6/2b/56b9aff361d7f1ab728a805ec7d7ea835f8807afa9f5cc690ea0e630efb9/tornado-6.5.7-cp39-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:7778b30bef919231265e91c69963ce0f49a1e9c07ac900bbe75b19ce2575ba92", size = 450578, upload-time = "2026-06-08T17:34:43.787Z" }, - { url = "https://files.pythonhosted.org/packages/02/30/a7444fb23aa76860a14198fab96ac79f1866b0a6e19e26c4381b0938e50f/tornado-6.5.7-cp39-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:e726f0c75da7726eec023aa62751ff8878bd2737e34fbdd33b1ae5897d2200f5", size = 449985, upload-time = "2026-06-08T17:34:45.326Z" }, - { url = "https://files.pythonhosted.org/packages/5c/42/5f0e56c01e8d9d36f4e23f367b85ae6cae0c1ecddd5e6977d8388ad27488/tornado-6.5.7-cp39-abi3-win32.whl", hash = "sha256:f8de3bf12d3efdd0cbe7c8887868198f8a91415e3f29fcf258d9b8eb7b1d9ae4", size = 451047, upload-time = "2026-06-08T17:34:46.784Z" }, - { url = "https://files.pythonhosted.org/packages/c9/a4/b393076ffb21b469eec5b328a0534cf03a3b90bfc6b1f09507cdd075d938/tornado-6.5.7-cp39-abi3-win_amd64.whl", hash = "sha256:de942f843533a039ef9fa3d9c88c7cd8a7c94553fb5ad0154270989b3d99a2c4", size = 451485, upload-time = "2026-06-08T17:34:48.248Z" }, - { url = "https://files.pythonhosted.org/packages/71/2e/7b1c769803121b809112cf9a00681c472eae1d80e32d7ec0e0bd61d0d0e1/tornado-6.5.7-cp39-abi3-win_arm64.whl", hash = "sha256:ff934fce95643af5f11efdae618eaa73d469dc588641e5c8d19295a0c65c4796", size = 450506, upload-time = "2026-06-08T17:34:49.702Z" }, + { url = "https://files.pythonhosted.org/packages/f2/d5/007086fd8df5489338e204f65adce33fd4f21a4999dbb2b9cff2f897b5f4/tornado-6.5.8-cp39-abi3-macosx_10_9_universal2.whl", hash = "sha256:cc6aa787d7cfab7c3d35189dc7a56fbd2399a569624c730c6b55b3d6531d0403", size = 449487, upload-time = "2026-08-07T02:12:28.682Z" }, + { url = "https://files.pythonhosted.org/packages/70/c8/5a24a99495903f594f6a199dd7beead1cbc0a13e2cb9102727bcaaf2a997/tornado-6.5.8-cp39-abi3-macosx_10_9_x86_64.whl", hash = "sha256:9715b5eb79735b2bcd454ce216a9275b7c0470e64ea1bf5742f78b2f72b26eeb", size = 447649, upload-time = "2026-08-07T02:12:30.306Z" }, + { url = "https://files.pythonhosted.org/packages/6e/de/f2e733f386b85962d1b1dc82cd63d169b5b4580062b35397eac9244a41fe/tornado-6.5.8-cp39-abi3-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:547d63f450d570c14fe0e8db2cfb14c9bbd1c2503b4a6612586267955aa47b58", size = 450707, upload-time = "2026-08-07T02:12:31.95Z" }, + { url = "https://files.pythonhosted.org/packages/0b/94/20efeee9a01c141e9ac47c397f81679dfda24b32768fc4fff24e76d36c2c/tornado-6.5.8-cp39-abi3-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:7e2360a0ffbe145eca8af0b19cb7203d79b1a98dd4cccdd6b368f6f49c2e3808", size = 451677, upload-time = "2026-08-07T02:12:33.512Z" }, + { url = "https://files.pythonhosted.org/packages/42/ec/a96ccb8ccf0de2b7bc2c5fa1608a4803735018242e90c4882365a9fd418f/tornado-6.5.8-cp39-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:5d242290bdf7ab3151bc1065fdd75c0dcc21cbc7b49f22a4c56329c2d6566d22", size = 451510, upload-time = "2026-08-07T02:12:35.346Z" }, + { url = "https://files.pythonhosted.org/packages/29/b5/93185859245ad3f00e62175f29607346788b696369347f0146e0421286bb/tornado-6.5.8-cp39-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:7b94ff0e128fe0542f3bd331fb44d06260fc4ac16881545159f34ef08aad4195", size = 450917, upload-time = "2026-08-07T02:12:36.963Z" }, + { url = "https://files.pythonhosted.org/packages/97/cf/fe33cf062834487d34d1559746a4a12521033c22645b6d74d4bca702e018/tornado-6.5.8-cp39-abi3-win32.whl", hash = "sha256:67832909c4779c64942380cb5f044a5c6163d00831472d80e25e115de9917836", size = 451952, upload-time = "2026-08-07T02:12:38.512Z" }, + { url = "https://files.pythonhosted.org/packages/cb/e1/468ad54333e92ccb62627e62cb88e5fc14a2171daa67ed47b1b8542d5b86/tornado-6.5.8-cp39-abi3-win_amd64.whl", hash = "sha256:11881db6b7c168494be2c2d12e65931451bdf7ee718535418ae1d8855dd5a0ee", size = 452391, upload-time = "2026-08-07T02:12:39.971Z" }, + { url = "https://files.pythonhosted.org/packages/ad/3e/cd5e4f06e34cde33b8ef66cf36aa2b5ad46354cc1af7d2136bbe365fee1d/tornado-6.5.8-cp39-abi3-win_arm64.whl", hash = "sha256:68a7468c7e289f8514d7d664101753903217eff1bb6822c6b5994a0b5f5bcb26", size = 451411, upload-time = "2026-08-07T02:12:41.469Z" }, ] [[package]] From 386946353ac425c6cbcac28368846dc2ab413bc2 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Tue, 1 Sep 2026 14:39:15 -0700 Subject: [PATCH 13/48] fix(vertex): avoid duplicate DeepSeek OCR model namespace --- .../vertex_ai/ocr/deepseek_transformation.py | 3 ++- tests/ocr_tests/test_ocr_vertex_ai.py | 20 ++++++++++++++++++- 2 files changed, 21 insertions(+), 2 deletions(-) diff --git a/litellm/llms/vertex_ai/ocr/deepseek_transformation.py b/litellm/llms/vertex_ai/ocr/deepseek_transformation.py index 2603552152d..b57a87c3325 100644 --- a/litellm/llms/vertex_ai/ocr/deepseek_transformation.py +++ b/litellm/llms/vertex_ai/ocr/deepseek_transformation.py @@ -177,8 +177,9 @@ class VertexAIDeepSeekOCRConfig(BaseOCRConfig): content_item = {"type": "image_url", "image_url": document_url} # Build DeepSeek OCR request + provider_model: Final = model if model.startswith("deepseek-ai/") else f"deepseek-ai/{model}" data: Final = { - "model": "deepseek-ai/" + model, + "model": provider_model, "messages": [{"role": "user", "content": [content_item]}], } diff --git a/tests/ocr_tests/test_ocr_vertex_ai.py b/tests/ocr_tests/test_ocr_vertex_ai.py index 1ba5b9d0883..1842eb063a5 100644 --- a/tests/ocr_tests/test_ocr_vertex_ai.py +++ b/tests/ocr_tests/test_ocr_vertex_ai.py @@ -5,9 +5,11 @@ Note: Vertex AI OCR automatically converts URLs to base64 data URIs since the Vertex AI endpoint doesn't have internet access. """ -import os import json +import os import tempfile +from typing import Final + import pytest from base_ocr_unit_tests import BaseOCRTest @@ -139,3 +141,19 @@ def test_vertex_ai_ocr_routing(): assert isinstance( deepseek_variant, VertexAIDeepSeekOCRConfig ), "DeepSeek variant should route to VertexAIDeepSeekOCRConfig" + + +@pytest.mark.parametrize("model", ("deepseek-ocr-maas", "deepseek-ai/deepseek-ocr-maas")) +def test_deepseek_request_uses_single_provider_namespace(model: str) -> None: + from litellm.llms.vertex_ai.ocr.deepseek_transformation import ( + VertexAIDeepSeekOCRConfig, + ) + + request: Final = VertexAIDeepSeekOCRConfig().transform_ocr_request( + model=model, + document={"type": "image_url", "image_url": "data:image/png;base64,AA=="}, + optional_params={}, + headers={}, + ) + + assert request.data["model"] == "deepseek-ai/deepseek-ocr-maas" From d59fcda8af69f5545a8e7c29b29d12362e2bffad Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 1 Sep 2026 14:47:59 -0700 Subject: [PATCH 14/48] fix(rerank): adopt declared authenticating providers in arerank instead of resolving them get_llm_provider runs the OAuth device flow for github_copilot and chatgpt, so calling it on the event loop before the executor dispatch let an authenticated caller block the loop for the length of the polling window. Adopt the declared provider via declared_authenticating_provider, matching the metadata callers in utils.py, and only resolve for everything else. --- litellm/rerank_api/main.py | 19 +++++++++---- tests/test_litellm/rerank_api/test_main.py | 32 ++++++++++++++++++++++ 2 files changed, 45 insertions(+), 6 deletions(-) diff --git a/litellm/rerank_api/main.py b/litellm/rerank_api/main.py index 597d1cfb863..37ca989b8d3 100644 --- a/litellm/rerank_api/main.py +++ b/litellm/rerank_api/main.py @@ -6,6 +6,7 @@ from typing import Any, Final, Literal import litellm from litellm._logging import verbose_logger +from litellm.litellm_core_utils.get_llm_provider_logic import declared_authenticating_provider from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.base_llm.rerank.transformation import BaseRerankConfig from litellm.llms.bedrock.rerank.handler import BedrockRerankHandler @@ -43,16 +44,22 @@ async def arerank( """ Async: Reranks a list of documents based on their relevance to the query """ - _custom_llm_provider: str | None = None # rebind-ok: set by the get_llm_provider unpack; read in the except + _custom_llm_provider: str | None = ( + None # rebind-ok: set by the declared-provider guard or the get_llm_provider unpack; read in the except + ) try: loop: Final = asyncio.get_event_loop() kwargs["arerank"] = True - _, _custom_llm_provider, _, _ = litellm.get_llm_provider( # rebind-ok: see pre-declaration above - model=model, - custom_llm_provider=custom_llm_provider, - api_base=kwargs.get("api_base", None), - ) + declared_provider: Final = declared_authenticating_provider(model, custom_llm_provider) + if declared_provider is not None: + _custom_llm_provider = declared_provider # rebind-ok: see pre-declaration above + else: + _, _custom_llm_provider, _, _ = litellm.get_llm_provider( # rebind-ok: see pre-declaration above + model=model, + custom_llm_provider=custom_llm_provider, + api_base=kwargs.get("api_base", None), + ) func: Final = partial( rerank, diff --git a/tests/test_litellm/rerank_api/test_main.py b/tests/test_litellm/rerank_api/test_main.py index 62149c742d6..2b6cfeda2c2 100644 --- a/tests/test_litellm/rerank_api/test_main.py +++ b/tests/test_litellm/rerank_api/test_main.py @@ -172,6 +172,38 @@ async def test_arerank_error_is_mapped_to_litellm_exception(respx_mock: respx.Mo assert "None - " not in str(exc_info.value) +@pytest.mark.asyncio +async def test_arerank_declared_authenticating_provider_skips_resolution(monkeypatch): + """Regression for the event-loop hazard in arerank's provider pre-resolution: + get_llm_provider runs the blocking OAuth device flow for github_copilot/chatgpt, + so arerank must adopt the declared provider instead of resolving it, while the + except path still maps with that declared provider.""" + from litellm.llms.base_llm.chat.transformation import BaseLLMException + + resolution_calls = [] + + def record_resolution(*args, **kwargs): + resolution_calls.append((args, kwargs)) + return "gpt-4o", "github_copilot", None, None + + def rerank_raises_provider_error(*args, **kwargs): + raise BaseLLMException(status_code=401, message='{"error":"bad key"}') + + monkeypatch.setattr(litellm, "get_llm_provider", record_resolution) + monkeypatch.setattr("litellm.rerank_api.main.rerank", rerank_raises_provider_error) + + with pytest.raises(litellm.AuthenticationError) as exc_info: + await litellm.arerank( + model="github_copilot/gpt-4o", + query=MARKER_QUERY, + documents=[MARKER_DOC], + ) + + assert resolution_calls == [] + assert "Github_copilotException" in str(exc_info.value) + assert "None - " not in str(exc_info.value) + + @pytest.mark.asyncio async def test_together_rerank_async_honors_env_api_base(respx_mock: respx.MockRouter, monkeypatch): """Regression: TOGETHER_AI_API_BASE was honored by chat but ignored by rerank.""" From fcd9052179f039f075b3e25a8c1aec8657fc98a0 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 1 Sep 2026 17:10:52 -0700 Subject: [PATCH 15/48] feat(proxy): honor model_info.display_name in the Anthropic-shaped /v1/models listing --- litellm/llms/anthropic/common_utils.py | 17 ++-- .../proxy/common_utils/model_listing_utils.py | 25 +++++- litellm/proxy/proxy_server.py | 21 +++-- litellm/router.py | 20 +++++ .../proxy/proxy_server/test_routes_models.py | 78 +++++++++++++++++++ .../test_team_model_name_translation.py | 26 ++++++- tests/test_litellm/test_router.py | 65 ++++++++++++++++ 7 files changed, 240 insertions(+), 12 deletions(-) diff --git a/litellm/llms/anthropic/common_utils.py b/litellm/llms/anthropic/common_utils.py index c60ebd844ba..d23690976ad 100644 --- a/litellm/llms/anthropic/common_utils.py +++ b/litellm/llms/anthropic/common_utils.py @@ -1378,31 +1378,38 @@ def process_anthropic_headers(headers: httpx.Headers | dict) -> dict: return additional_headers -def _anthropic_model_entry(model: ModelInfoResponse, created_at: str) -> Mapping[str, object]: +def _anthropic_model_entry( + model: ModelInfoResponse, created_at: str, display_names: Mapping[str, str] +) -> Mapping[str, object]: return { # mutable-ok: JSON response body, serialized by the route and never mutated "type": "model", "id": model["id"], - "display_name": model["id"], + "display_name": display_names.get(model["id"], model["id"]), "created_at": created_at, "max_input_tokens": model.get("max_input_tokens"), "max_tokens": model.get("max_output_tokens"), } -def create_anthropic_model_list_response(models: Sequence[ModelInfoResponse]) -> Mapping[str, object]: +def create_anthropic_model_list_response( + models: Sequence[ModelInfoResponse], + display_names: Mapping[str, str] = MappingProxyType({}), +) -> Mapping[str, object]: """Build the Anthropic-native /v1/models envelope. Clients that send an anthropic-version header parse the Anthropic Models API shape (type/display_name/created_at plus has_more/first_id/last_id) and filter the list themselves, so every model is returned here. The token limits carry over from the OpenAI-shaped listing, named as the Messages API names them, and - are always present because the vendor shape declares them nullable, not optional + are always present because the vendor shape declares them nullable, not optional. + display_names maps a listed model id to a configured human-readable name; ids + without an entry fall back to the id itself, matching the vendor behavior """ created_at: Final = ( datetime.fromtimestamp(DEFAULT_MODEL_CREATED_AT_TIME, tz=timezone.utc).isoformat().replace("+00:00", "Z") ) data: Final = [ # mutable-ok: JSON response body, serialized by the route and never mutated - _anthropic_model_entry(model, created_at) for model in models + _anthropic_model_entry(model, created_at, display_names) for model in models ] return { # mutable-ok: JSON response body, serialized by the route and never mutated "data": data, diff --git a/litellm/proxy/common_utils/model_listing_utils.py b/litellm/proxy/common_utils/model_listing_utils.py index 9fd24162f7e..213a697b3dd 100644 --- a/litellm/proxy/common_utils/model_listing_utils.py +++ b/litellm/proxy/common_utils/model_listing_utils.py @@ -10,13 +10,36 @@ legacy internal names with `general_settings.use_team_public_model_name: false`. from __future__ import annotations -from collections.abc import Mapping +from collections.abc import Mapping, Sequence +from types import MappingProxyType from typing import TYPE_CHECKING, Final, cast if TYPE_CHECKING: from litellm.router import Router +def configured_display_names( + entries: Sequence[tuple[str, str]], + llm_router: Router | None, +) -> Mapping[str, str]: + """response_id -> configured `model_info.display_name` for the listing entries + that have one. + + Metadata is looked up by each entry's internal lookup id (so team-scoped rows + resolve), while the returned map is keyed by the public response id the + Anthropic-shaped listing is built from. Entries without a configured name are + omitted so the listing falls back to the id itself. + """ + if llm_router is None: + return MappingProxyType({}) + resolved: Final = ( + (response_id, llm_router.get_configured_display_name(lookup_id)) for response_id, lookup_id in entries + ) + return MappingProxyType( + {response_id: display_name for response_id, display_name in resolved if display_name is not None} + ) + + class TeamModelNameTranslator: """Translates internal team routing keys to their public names for the model listing/retrieve responses. Stateless; the live router and general_settings diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 2c600667283..4f172ca29b9 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -351,7 +351,10 @@ from litellm.proxy.common_utils.load_config_utils import ( get_file_contents_from_s3, ) from litellm.proxy.common_utils.model_deprecation import collect_model_deprecations -from litellm.proxy.common_utils.model_listing_utils import TeamModelNameTranslator +from litellm.proxy.common_utils.model_listing_utils import ( + TeamModelNameTranslator, + configured_display_names, +) from litellm.proxy.common_utils.openai_endpoint_utils import ( remove_sensitive_info_from_deployment, ) @@ -10193,7 +10196,8 @@ async def model_list( # The internal routing key drives the metadata/fallback lookup, while the # public name is what the client sees as the model id. model_data = [] - for response_id, lookup_id in TeamModelNameTranslator.listing_entries(all_models, llm_router, settings): + admin_entries: Final = TeamModelNameTranslator.listing_entries(all_models, llm_router, settings) + for response_id, lookup_id in admin_entries: model_info = create_model_info_response( model_id=lookup_id, provider="openai", @@ -10206,7 +10210,10 @@ async def model_list( if wants_anthropic_format: admin_listing: Final = cast(Sequence[ModelInfoResponse], model_data) # cast-ok: rows built above - return create_anthropic_model_list_response(admin_listing) + return create_anthropic_model_list_response( + admin_listing, + display_names=configured_display_names(admin_entries, llm_router), + ) return dict( data=model_data, @@ -10237,7 +10244,8 @@ async def model_list( # The internal routing key drives the metadata/fallback lookup, while the # public name is what the client sees as the model id. model_data = [] - for response_id, lookup_id in TeamModelNameTranslator.listing_entries(all_models, llm_router, settings): + entries: Final = TeamModelNameTranslator.listing_entries(all_models, llm_router, settings) + for response_id, lookup_id in entries: model_info = create_model_info_response( model_id=lookup_id, provider="openai", @@ -10250,7 +10258,10 @@ async def model_list( if wants_anthropic_format: listing: Final = cast(Sequence[ModelInfoResponse], model_data) # cast-ok: rows built above - return create_anthropic_model_list_response(listing) + return create_anthropic_model_list_response( + listing, + display_names=configured_display_names(entries, llm_router), + ) return dict( data=model_data, diff --git a/litellm/router.py b/litellm/router.py index 471a1116f44..6245f8a6a03 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -9607,6 +9607,26 @@ class Router: coerce_token_limit(model_info.get("max_output_tokens")), ) + def get_configured_display_name(self, model_name: str) -> "str | None": + """ + Return the display_name explicitly configured in a concrete deployment's + model_info for model_name, via O(1) index lookup. + + Returns None for wildcard-expanded or unknown names, and treats a + non-string or empty configured value as absent rather than failing the + listing. Like get_configured_token_limits, this never triggers pattern + matching or deep copies, so it is safe to call per listed model on the + /v1/models hot path. + """ + deployment: Final = self.get_deployment_by_model_group_name(model_group_name=model_name) + if deployment is None: + return None + + display_name: Final = deployment.model_info.get("display_name") + if isinstance(display_name, str) and display_name.strip(): + return display_name + return None + def get_deployment_credentials_with_provider( self, model_id: str, team_id: str | None = None ) -> dict[str, Any] | None: diff --git a/tests/test_litellm/proxy/proxy_server/test_routes_models.py b/tests/test_litellm/proxy/proxy_server/test_routes_models.py index 2b126b1ea95..bc6106a06f8 100644 --- a/tests/test_litellm/proxy/proxy_server/test_routes_models.py +++ b/tests/test_litellm/proxy/proxy_server/test_routes_models.py @@ -45,6 +45,7 @@ def patched_models(monkeypatch): deployment = MagicMock() deployment.litellm_params.model = "gpt-4" router.get_deployment_by_model_group_name = MagicMock(return_value=deployment) + router.get_configured_display_name = MagicMock(return_value=None) monkeypatch.setattr(proxy_server, "llm_router", router) monkeypatch.setattr(proxy_server, "prisma_client", MagicMock()) @@ -187,6 +188,83 @@ def test_anthropic_format_carries_router_configured_token_limits(client, auth_as assert (claude["max_input_tokens"], claude["max_tokens"]) == (500000, 4096) +@pytest.mark.parametrize("path", ["/v1/models", "/models"]) +def test_anthropic_format_uses_configured_display_name(client, auth_as, patched_models, path): + """A deployment's ``model_info.display_name`` becomes the Anthropic-native + ``display_name`` so Claude Code's picker shows a clean name while the id keeps + routing; models without one keep the id fallback, and the OpenAI-shaped + listing carries no display_name either way.""" + + def _configured(model_name): + return "Kimi K3" if model_name == "gpt-4" else None + + patched_models.get_configured_display_name = MagicMock(side_effect=_configured) + + with auth_as(): + anthropic_response = client.get(path, headers={"anthropic-version": "2023-06-01"}) + openai_response = client.get(path) + + assert anthropic_response.status_code == 200 + gpt_4, claude = anthropic_response.json()["data"] + assert (gpt_4["id"], gpt_4["display_name"]) == ("gpt-4", "Kimi K3") + assert (claude["id"], claude["display_name"]) == ("claude-sonnet", "claude-sonnet") + + assert openai_response.status_code == 200 + openai_models = openai_response.json()["data"] + assert [m["id"] for m in openai_models] == ["gpt-4", "claude-sonnet"] + assert all("display_name" not in m for m in openai_models) + + +@pytest.mark.parametrize("params", [{}, {"scope": "expand"}]) +def test_anthropic_display_name_resolved_via_internal_team_key( + client, auth_as, patched_models, monkeypatch, params +): + """For a team-scoped row the configured display name must be looked up by the + internal routing key while the entry itself is keyed by the public name, so + the clean name lands on the id the client actually sees.""" + from litellm.proxy import utils as proxy_utils + from litellm.proxy.auth import model_checks + + internal_name = "model_name_team-1_c0ffee" + + patched_models.get_model_list = MagicMock( + return_value=[ + { + "model_name": internal_name, + "model_info": { + "team_id": "team-1", + "team_public_model_name": "gpt-4-team", + }, + } + ] + ) + patched_models.get_model_names = MagicMock(return_value=[internal_name]) + patched_models.get_configured_display_name = MagicMock( + side_effect=lambda model_name: "Team GPT" if model_name == internal_name else None + ) + + async def _fake_get_available_models_for_user(**kwargs): + return [internal_name] + + monkeypatch.setattr( + proxy_utils, + "get_available_models_for_user", + _fake_get_available_models_for_user, + ) + monkeypatch.setattr( + model_checks, "get_complete_model_list", lambda **kwargs: [internal_name] + ) + + with auth_as(): + response = client.get( + "/v1/models", params=params, headers={"anthropic-version": "2023-06-01"} + ) + + assert response.status_code == 200 + (entry,) = response.json()["data"] + assert (entry["id"], entry["display_name"]) == ("gpt-4-team", "Team GPT") + + @pytest.mark.parametrize("path", ["/v1/models", "/models"]) def test_get_models_invalid_scope_returns_400(client, auth_as, patched_models, path): """Pins: ``GET /v1/models``, ``GET /models`` (error path: invalid scope).""" diff --git a/tests/test_litellm/proxy/proxy_server/test_team_model_name_translation.py b/tests/test_litellm/proxy/proxy_server/test_team_model_name_translation.py index aa35fd64f18..0fb9b1a6d88 100644 --- a/tests/test_litellm/proxy/proxy_server/test_team_model_name_translation.py +++ b/tests/test_litellm/proxy/proxy_server/test_team_model_name_translation.py @@ -19,7 +19,10 @@ from litellm.proxy._types import ( LitellmUserRoles, UserAPIKeyAuth, ) -from litellm.proxy.common_utils.model_listing_utils import TeamModelNameTranslator +from litellm.proxy.common_utils.model_listing_utils import ( + TeamModelNameTranslator, + configured_display_names, +) from litellm.proxy.proxy_server import ( _get_proxy_model_info, _translate_model_name_for_response, @@ -1391,6 +1394,27 @@ def test_resolve_public_name_respects_legacy_flag(): ) +def test_configured_display_names_keyed_by_response_id(): + """The map is keyed by the public response id while the router lookup uses + the internal routing key, and entries without a configured name are omitted.""" + router = MagicMock() + router.get_configured_display_name = MagicMock( + side_effect=lambda model_name: "Team Sonnet" if model_name == "model_name_team-abc-123_4a6b8" else None + ) + + assert configured_display_names( + entries=[ + ("team-claude-sonnet", "model_name_team-abc-123_4a6b8"), + ("gpt-4o", "gpt-4o"), + ], + llm_router=router, + ) == {"team-claude-sonnet": "Team Sonnet"} + + +def test_configured_display_names_empty_without_router(): + assert configured_display_names(entries=[("gpt-4o", "gpt-4o")], llm_router=None) == {} + + @pytest.mark.asyncio async def test_retrieve_model_by_public_name_returns_200(monkeypatch): """Regression: `GET /v1/models/{public_name}` must NOT 404. The listing diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 44c1cdbff06..f4ea9b03a80 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -7271,6 +7271,71 @@ def test_get_configured_token_limits_coerces_numeric_strings(): assert router.get_configured_token_limits("quoted-limits-model") == (32000, 8000) +def test_get_configured_display_name_reads_deployment_model_info(): + router = litellm.Router( + model_list=[ + { + "model_name": "Kimi K3-claude-compatible", + "litellm_params": {"model": "openai/some-unmapped-model"}, + "model_info": {"display_name": "Kimi K3"}, + } + ] + ) + + assert router.get_configured_display_name("Kimi K3-claude-compatible") == "Kimi K3" + + +def test_get_configured_display_name_returns_none_for_unset_or_unknown(): + router = litellm.Router( + model_list=[ + { + "model_name": "no-display-model", + "litellm_params": {"model": "openai/some-unmapped-model"}, + } + ] + ) + + assert router.get_configured_display_name("no-display-model") is None + assert router.get_configured_display_name("not-a-real-model") is None + + +def test_get_configured_display_name_skips_wildcard_pattern_matching(): + router = litellm.Router( + model_list=[ + { + "model_name": "bedrock/*", + "litellm_params": {"model": "bedrock/*"}, + "model_info": {"display_name": "Bedrock"}, + } + ] + ) + + with patch.object( + router.pattern_router, "route", side_effect=AssertionError("pattern route called") + ): + assert ( + router.get_configured_display_name("bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0") + is None + ) + + +def test_get_configured_display_name_treats_malformed_values_as_absent(): + malformed = ["", " ", 12345, ["Kimi K3"], {"name": "Kimi K3"}, True] + router = litellm.Router( + model_list=[ + { + "model_name": f"bad-display-{i}", + "litellm_params": {"model": "openai/some-unmapped-model"}, + "model_info": {"display_name": bad}, + } + for i, bad in enumerate(malformed) + ] + ) + + for i in range(len(malformed)): + assert router.get_configured_display_name(f"bad-display-{i}") is None + + @pytest.mark.asyncio async def test_acreate_batch_disable_fallbacks_surfaces_owning_provider_error(): router = litellm.Router( From 0e7a05d8781963fbbc60301df350c74db46fa82c Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 1 Sep 2026 18:46:41 -0700 Subject: [PATCH 16/48] test(bedrock): read the sent bearer header off the mock instead of a hand-rolled recorder --- .../chat/test_bedrock_converse_handler.py | 32 ++++++------------- 1 file changed, 10 insertions(+), 22 deletions(-) diff --git a/tests/test_litellm/llms/bedrock/chat/test_bedrock_converse_handler.py b/tests/test_litellm/llms/bedrock/chat/test_bedrock_converse_handler.py index 4a296d9296b..21e3239f623 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_bedrock_converse_handler.py +++ b/tests/test_litellm/llms/bedrock/chat/test_bedrock_converse_handler.py @@ -356,27 +356,14 @@ async def test_async_completion_logs_pre_call_by_default(): assert logging_obj.pre_call.call_count == 1 -def _recording_sync_client(): - """A sync transport that answers with a Converse response and keeps every - `post` payload, so a test can assert which headers were sent.""" - posted: list[dict] = [] - client = MagicMock() - - def post(**kwargs): - posted.append(kwargs) - return httpx.Response( - 200, - json=CONVERSE_RESPONSE, - request=httpx.Request("POST", "https://bedrock-runtime.us-west-2.amazonaws.com"), - ) - - client.post = post - client.__class__ = HTTPHandler - return client, posted - - def _sync_client_returning_converse_response(): - client, _ = _recording_sync_client() + client = MagicMock() + client.post.side_effect = lambda **_kwargs: httpx.Response( + 200, + json=CONVERSE_RESPONSE, + request=httpx.Request("POST", "https://bedrock-runtime.us-west-2.amazonaws.com"), + ) + client.__class__ = HTTPHandler return client @@ -505,12 +492,13 @@ def test_bearer_token_auth_serves_when_boto3_resolves_no_sigv4_credentials(monke credentials at all. Preparing the Rust handoff must not dereference that None: the bearer token signs the request on its own.""" monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", "bedrock-bearer-token") - client, posted = _recording_sync_client() + client = _sync_client_returning_converse_response() response = _run(credentials=None, litellm_params={}, client=client) assert response.choices[0].message.content == "hi" - assert posted[0]["headers"]["Authorization"] == "Bearer bedrock-bearer-token" + sent_headers = client.post.call_args.kwargs["headers"] + assert sent_headers["Authorization"] == "Bearer bedrock-bearer-token" def test_the_rust_opt_in_needs_no_sigv4_principal(): From f49a3e15a8b936c94e5d050017f6887e8d34e997 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 1 Sep 2026 19:12:07 -0700 Subject: [PATCH 17/48] test(e2e): read JUnit properties off the real collected pytest Item tests/e2e/test_junit_properties.py fed a hand-rolled FakeItem to result_properties and attach_result_properties, both typed pytest.Item, so uv run basedpyright tests/e2e reported 3 reportArgumentType errors on litellm_internal_staging and every make check that scopes a litellm/ or tests/e2e/ Python file failed. Each test now looks up its own collected Item in request.session.items and applies the covers marker at run time through request.applymarker, so the coverage registry's collect-only pass never sees the test ids and the production functions keep their pytest.Item signatures. No casts, no ignores. Resolves LIT-6669 --- tests/e2e/test_junit_properties.py | 45 ++++++++++-------------------- 1 file changed, 15 insertions(+), 30 deletions(-) diff --git a/tests/e2e/test_junit_properties.py b/tests/e2e/test_junit_properties.py index c0596177cc1..02c1413c840 100644 --- a/tests/e2e/test_junit_properties.py +++ b/tests/e2e/test_junit_properties.py @@ -24,25 +24,10 @@ from junit_properties import ( ) -class FakeMarker: - def __init__(self, name: str, *args: object) -> None: - self.name = name - self.args = args - - -class FakeItem: - """The three attributes junit_properties reads off a pytest Item.""" - - def __init__( - self, nodeid: str, location: tuple[str, int | None, str], markers: tuple[FakeMarker, ...] = () - ) -> None: - self.nodeid = nodeid - self.location = location - self.user_properties: list[tuple[str, str]] = [] - self._markers = markers - - def iter_markers(self, name: str): - return (marker for marker in self._markers if marker.name == name) +def collected_item(request: pytest.FixtureRequest, name: str) -> pytest.Item: + """The Item pytest collected for test ``name`` in this file: the real nodeid, + location and marker machinery the collection hook reads, as pytest built it.""" + return next(item for item in request.session.items if item.path == request.path and item.name == name) def repo_root() -> Path | None: @@ -109,22 +94,22 @@ class TestSourceFromLocation: class TestResultProperties: - def test_every_test_carries_package_covers_and_source(self) -> None: - item = FakeItem( - "logging/test_x.py::TestFoo::test_bar", - ("logging/test_x.py", 40, "TestFoo.test_bar"), - (FakeMarker("covers", "LOG-1", "LOG-2"),), - ) - assert result_properties(item) == ( - ("package", "logging"), + def test_every_test_carries_package_covers_and_source(self, request: pytest.FixtureRequest) -> None: + """Read off this test's own collected Item, so the nodeid and location are + whatever pytest reports for the launch shape in use, and the marker is added + at run time so the coverage registry's collect-only pass never sees it.""" + test = type(self).test_every_test_carries_package_covers_and_source + request.applymarker(pytest.mark.covers("LOG-1", "LOG-2")) + assert result_properties(collected_item(request, test.__name__)) == ( + ("package", "root"), ("covers", "LOG-1,LOG-2"), - ("source", "tests/e2e/logging/test_x.py:41"), + ("source", f"tests/e2e/test_junit_properties.py:{test.__code__.co_firstlineno}"), ) - def test_attach_is_idempotent(self) -> None: + def test_attach_is_idempotent(self, request: pytest.FixtureRequest) -> None: """Collection can run the hook more than once; a second pass must not double the entries in the report.""" - item = FakeItem("logging/test_x.py::test_bar", ("logging/test_x.py", 40, "test_bar")) + item = collected_item(request, type(self).test_attach_is_idempotent.__name__) attach_result_properties(item) attach_result_properties(item) assert [name for name, _ in item.user_properties] == ["package", "covers", "source"] From 0608f0a00f2c76b50640ed7f4559f4c8551fdc44 Mon Sep 17 00:00:00 2001 From: mateo Date: Wed, 2 Sep 2026 02:17:02 +0000 Subject: [PATCH 18/48] fix: reject unknown runtime router settings Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/constants.py | 22 +++++++++++ litellm/proxy/proxy_server.py | 21 +++++++++- litellm/router.py | 30 ++++---------- litellm/types/router.py | 29 +++++++------- .../proxy/proxy_server/test_routes_config.py | 39 +++++++++++++++++++ .../test_router_retry_policy_update.py | 21 +++++++++- ui/litellm-dashboard/src/lib/http/schema.d.ts | 2 + 7 files changed, 125 insertions(+), 39 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index 1bd977dd9a9..a7506cb6378 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -38,6 +38,28 @@ DEFAULT_MAX_TOKENS: Final = int(os.getenv("DEFAULT_MAX_TOKENS", 4096)) DEFAULT_ALLOWED_FAILS: Final = int(os.getenv("DEFAULT_ALLOWED_FAILS", 3)) DEFAULT_REDIS_SYNC_INTERVAL: Final = int(os.getenv("DEFAULT_REDIS_SYNC_INTERVAL", 1)) DEFAULT_COOLDOWN_TIME_SECONDS: Final = int(os.getenv("DEFAULT_COOLDOWN_TIME_SECONDS", 5)) +RUNTIME_UPDATABLE_ROUTER_SETTINGS: Final[frozenset[str]] = frozenset( + { + "routing_strategy_args", + "routing_strategy", + "routing_groups", + "allowed_fails", + "cooldown_time", + "num_retries", + "timeout", + "max_retries", + "retry_after", + "fallbacks", + "context_window_fallbacks", + "retry_policy", + "model_group_retry_policy", + "model_group_alias", + "enable_weighted_failover", + "enable_tag_filtering", + "tag_routing_prefix", + "optional_pre_call_checks", + } +) DEFAULT_REPLICATE_POLLING_RETRIES: Final = int(os.getenv("DEFAULT_REPLICATE_POLLING_RETRIES", 5)) DEFAULT_REPLICATE_POLLING_DELAY_SECONDS: Final = int(os.getenv("DEFAULT_REPLICATE_POLLING_DELAY_SECONDS", 1)) DEFAULT_IMAGE_TOKEN_COUNT: Final = int(os.getenv("DEFAULT_IMAGE_TOKEN_COUNT", 250)) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 77a80ea0052..9de6b38265a 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -39,7 +39,7 @@ from typing import ( import anyio import websockets import websockets.exceptions -from pydantic import BaseModel, Json, JsonValue, ValidationError +from pydantic import BaseModel, Json, JsonValue, TypeAdapter, ValidationError from typing_extensions import NotRequired, ReadOnly, assert_never from litellm._uuid import uuid @@ -60,6 +60,7 @@ from litellm.constants import ( LITELLM_SETTINGS_SAFE_DB_OVERRIDES, LITELLM_UI_ALLOW_HEADERS, LITELLM_UI_SESSION_DURATION, + RUNTIME_UPDATABLE_ROUTER_SETTINGS, ) from litellm.litellm_core_utils.litellm_logging import ( _init_custom_logger_compatible_class, @@ -16207,6 +16208,7 @@ async def invitation_delete( ) async def update_config( config_info: ConfigYAML, + request: Request, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """ @@ -16218,6 +16220,23 @@ async def update_config( a side effect of an unrelated update. """ global llm_router, llm_model_list, general_settings, proxy_config, proxy_logging_obj, master_key, prisma_client + request_body: Final[Mapping[str, JsonValue]] = TypeAdapter(Mapping[str, JsonValue]).validate_python( + await request.json() + ) + raw_router_settings: Final = request_body.get("router_settings") + if isinstance(raw_router_settings, dict): + unsupported_router_settings: Final = sorted(set(raw_router_settings) - RUNTIME_UPDATABLE_ROUTER_SETTINGS) + if unsupported_router_settings: + raise HTTPException( + status_code=400, + detail={ + "error": ( + f"Unsupported router settings: {', '.join(unsupported_router_settings)} " + "are not runtime-updatable router settings" + ) + }, + ) + try: if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: raise HTTPException(status_code=403, detail="Only proxy admins can update config") diff --git a/litellm/router.py b/litellm/router.py index 23d8907fb49..fc4f815170d 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -50,6 +50,7 @@ from litellm.constants import ( DEFAULT_HEALTH_CHECK_INTERVAL, DEFAULT_HEALTH_CHECK_STALENESS_MULTIPLIER, DEFAULT_MAX_LRU_CACHE_SIZE, + RUNTIME_UPDATABLE_ROUTER_SETTINGS, SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY, ) from litellm.integrations.custom_logger import CustomLogger @@ -2072,6 +2073,10 @@ class Router: if _callback is None: continue + if self.optional_callbacks is not None and any( + isinstance(callback, type(_callback)) for callback in self.optional_callbacks + ): + continue if self.optional_callbacks is None: self.optional_callbacks = [] self.optional_callbacks.append(_callback) @@ -11331,27 +11336,6 @@ class Router: """ Update the router settings. """ - # only the following settings are allowed to be configured - _allowed_settings: Final = [ - "routing_strategy_args", - "routing_strategy", - "routing_groups", - "allowed_fails", - "cooldown_time", - "num_retries", - "timeout", - "max_retries", - "retry_after", - "fallbacks", - "context_window_fallbacks", - "retry_policy", - "model_group_retry_policy", - "model_group_alias", - "enable_weighted_failover", - "enable_tag_filtering", - "tag_routing_prefix", - ] - _int_settings: Final = [ "timeout", "num_retries", @@ -11364,13 +11348,15 @@ class Router: rebuild_routing_groups = False relink_lar1_from_args = False for var in kwargs: - if var in _allowed_settings: + if var in RUNTIME_UPDATABLE_ROUTER_SETTINGS: if var in _int_settings: _casted_value = int(kwargs[var]) setattr(self, var, _casted_value) elif var == "routing_groups": self._routing_groups_input = kwargs[var] rebuild_routing_groups = True + elif var == "optional_pre_call_checks": + self.add_optional_pre_call_checks(kwargs[var]) elif var == "retry_policy": value = kwargs[var] if isinstance(value, dict): diff --git a/litellm/types/router.py b/litellm/types/router.py index e0957383aac..2a5f264cee3 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -106,6 +106,20 @@ class RetryPolicy(BaseModel): InternalServerErrorRetries: int | None = None +OptionalPreCallChecks = list[ + Literal[ + "prompt_caching", + "router_budget_limiting", + "responses_api_deployment_check", + "deployment_affinity", + "session_affinity", + "forward_client_headers_by_model_group", + "enforce_model_rate_limits", + "encrypted_content_affinity", + ] +] + + class UpdateRouterConfig(BaseModel): """ Set of params that you can modify via `router.update_settings()`. @@ -128,6 +142,7 @@ class UpdateRouterConfig(BaseModel): model_group_alias: dict[str, str | dict] | None = {} enable_tag_filtering: bool | None = None tag_routing_prefix: str | None = None + optional_pre_call_checks: OptionalPreCallChecks | None = None model_config = ConfigDict(protected_namespaces=()) @@ -869,20 +884,6 @@ class FallbackAccessCheck(Protocol): async def __call__(self, *, model: str, request_kwargs: Mapping[str, object], llm_router: "Router") -> bool: ... -OptionalPreCallChecks = list[ - Literal[ - "prompt_caching", - "router_budget_limiting", - "responses_api_deployment_check", - "deployment_affinity", - "session_affinity", - "forward_client_headers_by_model_group", - "enforce_model_rate_limits", - "encrypted_content_affinity", - ] -] - - class LiteLLM_RouterFileObject(TypedDict, total=False): """ Tracking the litellm params hash, used for mapping the file id to the right model diff --git a/tests/test_litellm/proxy/proxy_server/test_routes_config.py b/tests/test_litellm/proxy/proxy_server/test_routes_config.py index ad3c470acf3..0df8fb663e2 100644 --- a/tests/test_litellm/proxy/proxy_server/test_routes_config.py +++ b/tests/test_litellm/proxy/proxy_server/test_routes_config.py @@ -60,6 +60,45 @@ def test_config_update_happy_admin(client, auth_as, mock_prisma, monkeypatch): assert normalize(response.json()) == {"message": "Config updated successfully"} +def test_config_update_persists_optional_pre_call_checks(client, auth_as, mock_prisma, monkeypatch): + from litellm.proxy import proxy_server as ps + from litellm.proxy._types import LitellmUserRoles + + table = _install_litellm_config(mock_prisma) + monkeypatch.setattr(ps, "prisma_client", mock_prisma) + fake_proxy_config = MagicMock() + fake_proxy_config.add_deployment = AsyncMock() + monkeypatch.setattr(ps, "proxy_config", fake_proxy_config) + + with auth_as(LitellmUserRoles.PROXY_ADMIN): + response = client.post( + "/config/update", + json={"router_settings": {"optional_pre_call_checks": ["prompt_caching"]}}, + ) + + assert response.status_code == 200 + persisted = json.loads(table.upsert.call_args.kwargs["data"]["create"]["param_value"]) + assert persisted["optional_pre_call_checks"] == ["prompt_caching"] + + +def test_config_update_rejects_unknown_router_setting(client, auth_as, mock_prisma, monkeypatch): + from litellm.proxy import proxy_server as ps + from litellm.proxy._types import LitellmUserRoles + + table = _install_litellm_config(mock_prisma) + monkeypatch.setattr(ps, "prisma_client", mock_prisma) + + with auth_as(LitellmUserRoles.PROXY_ADMIN): + response = client.post( + "/config/update", + json={"router_settings": {"optional_precall_checks": ["prompt_caching"]}}, + ) + + assert response.status_code == 400 + assert "optional_precall_checks" in response.json()["detail"]["error"] + table.upsert.assert_not_called() + + def test_config_update_non_admin_forbidden(client, auth_as, mock_prisma, monkeypatch): """POST /config/update by a non-admin caller is rejected; the error surfaces as a ProxyException with the admin-only message.""" diff --git a/tests/test_litellm/test_router_retry_policy_update.py b/tests/test_litellm/test_router_retry_policy_update.py index 1b98b8c1ae8..1b014cd8401 100644 --- a/tests/test_litellm/test_router_retry_policy_update.py +++ b/tests/test_litellm/test_router_retry_policy_update.py @@ -26,8 +26,8 @@ from unittest.mock import AsyncMock, MagicMock import pytest from pydantic import ValidationError - import litellm +from litellm.router_utils.pre_call_checks.prompt_caching_deployment_check import PromptCachingDeploymentCheck from litellm.types.router import RetryPolicy, UpdateRouterConfig # --------------------------------------------------------------------------- @@ -100,6 +100,19 @@ def _build_router() -> litellm.Router: ) +def test_update_settings_adds_optional_pre_call_check_once(): + router = _build_router() + + router.update_settings(num_retries=7, optional_pre_call_checks=["prompt_caching"]) + router.update_settings(optional_pre_call_checks=["prompt_caching"]) + + prompt_caching_callbacks = [ + callback for callback in router.optional_callbacks if isinstance(callback, PromptCachingDeploymentCheck) + ] + assert len(prompt_caching_callbacks) == 1 + assert router.num_retries == 7 + + def test_update_settings_persists_retry_policy_dict(): """When the proxy's ``_add_router_settings_from_db_config`` calls ``llm_router.update_settings(retry_policy={...})`` after reading the @@ -228,7 +241,7 @@ async def test_config_update_persists_and_reads_back_retry_policy(monkeypatch): """The exact global retry_policy save the UI performs must survive the real ``/config/update`` -> DB -> apply -> ``/get/config/callbacks`` path, not snap back to the ``num_retries`` fallback the ticket reported.""" - import litellm.proxy.proxy_server as proxy_server + from litellm.proxy import proxy_server from litellm.proxy._types import ConfigYAML, LitellmUserRoles, UserAPIKeyAuth router = _build_router() @@ -255,8 +268,12 @@ async def test_config_update_persists_and_reads_back_retry_policy(monkeypatch): RateLimitErrorRetries=7, ) ) + request = MagicMock() + request.json = AsyncMock(return_value={"router_settings": {"retry_policy": posted.model_dump()}}) + await proxy_server.update_config( config_info=ConfigYAML(router_settings=posted), + request=request, user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-1234"), ) diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 6f044fec3f3..bde7fd611d5 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -37473,6 +37473,8 @@ export interface components { } | null; /** Num Retries */ num_retries?: number | null; + /** Optional Pre Call Checks */ + optional_pre_call_checks?: ("prompt_caching" | "router_budget_limiting" | "responses_api_deployment_check" | "deployment_affinity" | "session_affinity" | "forward_client_headers_by_model_group" | "enforce_model_rate_limits" | "encrypted_content_affinity")[] | null; /** Retry After */ retry_after?: number | null; retry_policy?: components["schemas"]["RetryPolicy"] | null; From cb511f70ccbee8cc9257b52cc6ad5d7721219ce3 Mon Sep 17 00:00:00 2001 From: mateo Date: Wed, 2 Sep 2026 02:20:51 +0000 Subject: [PATCH 19/48] fix: preserve config update authorization order Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/constants.py | 44 +++++++++---------- litellm/proxy/proxy_server.py | 34 +++++++------- .../proxy/proxy_server/test_routes_config.py | 19 +++++++- .../test_router_retry_policy_update.py | 3 +- 4 files changed, 59 insertions(+), 41 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index a7506cb6378..5b44e8b5f51 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -9,6 +9,28 @@ DEFAULT_HEALTH_CHECK_PROMPT: Final = str(os.getenv("DEFAULT_HEALTH_CHECK_PROMPT" AZURE_DEFAULT_RESPONSES_API_VERSION: Final = str(os.getenv("AZURE_DEFAULT_RESPONSES_API_VERSION", "preview")) ROUTER_MAX_FALLBACKS: Final = int(os.getenv("ROUTER_MAX_FALLBACKS", 5)) ROUTER_FALLBACK_ERROR_DETAIL_MAX_CHARS: Final = 2000 +RUNTIME_UPDATABLE_ROUTER_SETTINGS: Final[frozenset[str]] = frozenset( + { + "routing_strategy_args", + "routing_strategy", + "routing_groups", + "allowed_fails", + "cooldown_time", + "num_retries", + "timeout", + "max_retries", + "retry_after", + "fallbacks", + "context_window_fallbacks", + "retry_policy", + "model_group_retry_policy", + "model_group_alias", + "enable_weighted_failover", + "enable_tag_filtering", + "tag_routing_prefix", + "optional_pre_call_checks", + } +) DEFAULT_BATCH_SIZE: Final = int(os.getenv("DEFAULT_BATCH_SIZE", 512)) DEFAULT_FLUSH_INTERVAL_SECONDS: Final = int(os.getenv("DEFAULT_FLUSH_INTERVAL_SECONDS", 5)) DEFAULT_S3_FLUSH_INTERVAL_SECONDS: Final = int(os.getenv("DEFAULT_S3_FLUSH_INTERVAL_SECONDS", 10)) @@ -38,28 +60,6 @@ DEFAULT_MAX_TOKENS: Final = int(os.getenv("DEFAULT_MAX_TOKENS", 4096)) DEFAULT_ALLOWED_FAILS: Final = int(os.getenv("DEFAULT_ALLOWED_FAILS", 3)) DEFAULT_REDIS_SYNC_INTERVAL: Final = int(os.getenv("DEFAULT_REDIS_SYNC_INTERVAL", 1)) DEFAULT_COOLDOWN_TIME_SECONDS: Final = int(os.getenv("DEFAULT_COOLDOWN_TIME_SECONDS", 5)) -RUNTIME_UPDATABLE_ROUTER_SETTINGS: Final[frozenset[str]] = frozenset( - { - "routing_strategy_args", - "routing_strategy", - "routing_groups", - "allowed_fails", - "cooldown_time", - "num_retries", - "timeout", - "max_retries", - "retry_after", - "fallbacks", - "context_window_fallbacks", - "retry_policy", - "model_group_retry_policy", - "model_group_alias", - "enable_weighted_failover", - "enable_tag_filtering", - "tag_routing_prefix", - "optional_pre_call_checks", - } -) DEFAULT_REPLICATE_POLLING_RETRIES: Final = int(os.getenv("DEFAULT_REPLICATE_POLLING_RETRIES", 5)) DEFAULT_REPLICATE_POLLING_DELAY_SECONDS: Final = int(os.getenv("DEFAULT_REPLICATE_POLLING_DELAY_SECONDS", 1)) DEFAULT_IMAGE_TOKEN_COUNT: Final = int(os.getenv("DEFAULT_IMAGE_TOKEN_COUNT", 250)) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 9de6b38265a..dbeb8486539 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -16220,27 +16220,27 @@ async def update_config( a side effect of an unrelated update. """ global llm_router, llm_model_list, general_settings, proxy_config, proxy_logging_obj, master_key, prisma_client - request_body: Final[Mapping[str, JsonValue]] = TypeAdapter(Mapping[str, JsonValue]).validate_python( - await request.json() - ) - raw_router_settings: Final = request_body.get("router_settings") - if isinstance(raw_router_settings, dict): - unsupported_router_settings: Final = sorted(set(raw_router_settings) - RUNTIME_UPDATABLE_ROUTER_SETTINGS) - if unsupported_router_settings: - raise HTTPException( - status_code=400, - detail={ - "error": ( - f"Unsupported router settings: {', '.join(unsupported_router_settings)} " - "are not runtime-updatable router settings" - ) - }, - ) - try: if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: raise HTTPException(status_code=403, detail="Only proxy admins can update config") + request_body: Final[Mapping[str, JsonValue]] = TypeAdapter(Mapping[str, JsonValue]).validate_python( + await request.json() + ) + raw_router_settings: Final = request_body.get("router_settings") + if isinstance(raw_router_settings, dict): + unsupported_router_settings: Final = sorted(set(raw_router_settings) - RUNTIME_UPDATABLE_ROUTER_SETTINGS) + if unsupported_router_settings: + raise HTTPException( + status_code=400, + detail={ + "error": ( + f"Unsupported router settings: {', '.join(unsupported_router_settings)} " + "are not runtime-updatable router settings" + ) + }, + ) + if prisma_client is None: raise Exception("No DB Connected") diff --git a/tests/test_litellm/proxy/proxy_server/test_routes_config.py b/tests/test_litellm/proxy/proxy_server/test_routes_config.py index 0df8fb663e2..4b0954c350a 100644 --- a/tests/test_litellm/proxy/proxy_server/test_routes_config.py +++ b/tests/test_litellm/proxy/proxy_server/test_routes_config.py @@ -95,10 +95,27 @@ def test_config_update_rejects_unknown_router_setting(client, auth_as, mock_pris ) assert response.status_code == 400 - assert "optional_precall_checks" in response.json()["detail"]["error"] + assert "optional_precall_checks" in response.json()["error"]["message"] table.upsert.assert_not_called() +def test_config_update_unknown_router_setting_non_admin_forbidden(client, auth_as, mock_prisma, monkeypatch): + from litellm.proxy import proxy_server as ps + from litellm.proxy._types import LitellmUserRoles + + _install_litellm_config(mock_prisma) + monkeypatch.setattr(ps, "prisma_client", mock_prisma) + + with auth_as(LitellmUserRoles.INTERNAL_USER): + response = client.post( + "/config/update", + json={"router_settings": {"optional_precall_checks": ["prompt_caching"]}}, + ) + + assert response.status_code == 403 + assert "admin" in response.json()["error"]["message"].lower() + + def test_config_update_non_admin_forbidden(client, auth_as, mock_prisma, monkeypatch): """POST /config/update by a non-admin caller is rejected; the error surfaces as a ProxyException with the admin-only message.""" diff --git a/tests/test_litellm/test_router_retry_policy_update.py b/tests/test_litellm/test_router_retry_policy_update.py index 1b014cd8401..e386eebf3d9 100644 --- a/tests/test_litellm/test_router_retry_policy_update.py +++ b/tests/test_litellm/test_router_retry_policy_update.py @@ -26,6 +26,7 @@ from unittest.mock import AsyncMock, MagicMock import pytest from pydantic import ValidationError + import litellm from litellm.router_utils.pre_call_checks.prompt_caching_deployment_check import PromptCachingDeploymentCheck from litellm.types.router import RetryPolicy, UpdateRouterConfig @@ -241,7 +242,7 @@ async def test_config_update_persists_and_reads_back_retry_policy(monkeypatch): """The exact global retry_policy save the UI performs must survive the real ``/config/update`` -> DB -> apply -> ``/get/config/callbacks`` path, not snap back to the ``num_retries`` fallback the ticket reported.""" - from litellm.proxy import proxy_server + import litellm.proxy.proxy_server as proxy_server from litellm.proxy._types import ConfigYAML, LitellmUserRoles, UserAPIKeyAuth router = _build_router() From 29f0110fe08d5b8f4798e20eed6d2368e4cada2d Mon Sep 17 00:00:00 2001 From: mateo Date: Wed, 2 Sep 2026 02:29:41 +0000 Subject: [PATCH 20/48] test: pass request to config update test Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/proxy_unit_tests/test_proxy_server.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/proxy_unit_tests/test_proxy_server.py b/tests/proxy_unit_tests/test_proxy_server.py index 47554913419..54cce9cdd78 100644 --- a/tests/proxy_unit_tests/test_proxy_server.py +++ b/tests/proxy_unit_tests/test_proxy_server.py @@ -3076,7 +3076,9 @@ async def test_update_config_success_callback_normalization(): admin_user = UserAPIKeyAuth( user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-test" ) - await proxy_server.update_config(config_update, user_api_key_dict=admin_user) + request = MagicMock() + request.json = AsyncMock(return_value={"litellm_settings": {"success_callback": ["SQS", "sQs"]}}) + await proxy_server.update_config(config_update, request=request, user_api_key_dict=admin_user) assert ( "litellm_settings" in upserted From e67f98feb1cb1758e253beaf005eaaad44ab3abe Mon Sep 17 00:00:00 2001 From: mateo Date: Wed, 2 Sep 2026 02:44:17 +0000 Subject: [PATCH 21/48] fix: reconcile runtime pre-call checks Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/router.py | 30 +++++++++++++- .../test_router_retry_policy_update.py | 40 +++++++++++++++++++ 2 files changed, 69 insertions(+), 1 deletion(-) diff --git a/litellm/router.py b/litellm/router.py index fc4f815170d..e2865542e89 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -355,6 +355,13 @@ _PreRoutingStrategyT = TypeVar("_PreRoutingStrategyT") _ALIAS_PARAMS_NEVER_FORWARDED: Final = frozenset({"model", "api_base", "api_key", "api_version"}) _ALIAS_MARKER_FORWARDED_PARAMS_KWARG: Final = "_alias_marker_forwarded_params" +_RUNTIME_TOGGLEABLE_PRE_CALL_CHECKS: Final[Mapping[str, type[CustomLogger]]] = MappingProxyType( + { + "prompt_caching": PromptCachingDeploymentCheck, + "enforce_model_rate_limits": ModelRateLimitingCheck, + } +) + def _stream_chunks_have_generated_content(chunks: Sequence[ModelResponseStream]) -> bool: for chunk in chunks: @@ -2082,6 +2089,27 @@ class Router: self.optional_callbacks.append(_callback) litellm.logging_callback_manager.add_litellm_callback(_callback) + def set_optional_pre_call_checks(self, optional_pre_call_checks: OptionalPreCallChecks | None) -> None: + if optional_pre_call_checks is None: + return + requested: Final = frozenset(optional_pre_call_checks) + for name, callback_cls in _RUNTIME_TOGGLEABLE_PRE_CALL_CHECKS.items(): + if name not in requested: + self._remove_optional_callbacks_of_type(callback_cls) + self.add_optional_pre_call_checks(optional_pre_call_checks) + + def _remove_optional_callbacks_of_type(self, callback_cls: type[CustomLogger]) -> None: + if self.optional_callbacks is None: + return + removed: Final = [cb for cb in self.optional_callbacks if isinstance(cb, callback_cls)] + if not removed: + return + self.optional_callbacks = [cb for cb in self.optional_callbacks if not isinstance(cb, callback_cls)] + for cb in removed: + litellm.logging_callback_manager.remove_callback_from_list_by_object( + litellm.callbacks, cb, require_self=False + ) + def print_deployment(self, deployment: dict): """ returns a copy of the deployment with the api key masked @@ -11356,7 +11384,7 @@ class Router: self._routing_groups_input = kwargs[var] rebuild_routing_groups = True elif var == "optional_pre_call_checks": - self.add_optional_pre_call_checks(kwargs[var]) + self.set_optional_pre_call_checks(kwargs[var]) elif var == "retry_policy": value = kwargs[var] if isinstance(value, dict): diff --git a/tests/test_litellm/test_router_retry_policy_update.py b/tests/test_litellm/test_router_retry_policy_update.py index e386eebf3d9..2c23d0da7e7 100644 --- a/tests/test_litellm/test_router_retry_policy_update.py +++ b/tests/test_litellm/test_router_retry_policy_update.py @@ -28,6 +28,8 @@ from pydantic import ValidationError import litellm +from litellm.router_strategy.budget_limiter import RouterBudgetLimiting +from litellm.router_utils.pre_call_checks.model_rate_limit_check import ModelRateLimitingCheck from litellm.router_utils.pre_call_checks.prompt_caching_deployment_check import PromptCachingDeploymentCheck from litellm.types.router import RetryPolicy, UpdateRouterConfig @@ -114,6 +116,44 @@ def test_update_settings_adds_optional_pre_call_check_once(): assert router.num_retries == 7 +def test_update_settings_clears_omitted_toggleable_pre_call_checks(): + router = _build_router() + + router.update_settings(optional_pre_call_checks=["prompt_caching"]) + router.update_settings(optional_pre_call_checks=[]) + + assert not any(isinstance(callback, PromptCachingDeploymentCheck) for callback in (router.optional_callbacks or [])) + assert not any(isinstance(callback, PromptCachingDeploymentCheck) for callback in litellm.callbacks) + + +def test_update_settings_replaces_toggleable_pre_call_checks(): + router = _build_router() + + router.update_settings(optional_pre_call_checks=["prompt_caching"]) + router.update_settings(optional_pre_call_checks=["enforce_model_rate_limits"]) + + assert not any(isinstance(callback, PromptCachingDeploymentCheck) for callback in (router.optional_callbacks or [])) + assert not any(isinstance(callback, PromptCachingDeploymentCheck) for callback in litellm.callbacks) + assert any(isinstance(callback, ModelRateLimitingCheck) for callback in (router.optional_callbacks or [])) + + +@pytest.mark.asyncio +async def test_update_settings_preserves_router_budget_limiting_when_omitted(monkeypatch): + async def _disable_periodic_sync(*args, **kwargs): + return None + + monkeypatch.setattr( + "litellm.router_strategy.budget_limiter.RouterBudgetLimiting.periodic_sync_in_memory_spend_with_redis", + _disable_periodic_sync, + ) + router = _build_router() + + router.add_optional_pre_call_checks(["router_budget_limiting"]) + router.update_settings(optional_pre_call_checks=[]) + + assert any(isinstance(callback, RouterBudgetLimiting) for callback in (router.optional_callbacks or [])) + + def test_update_settings_persists_retry_policy_dict(): """When the proxy's ``_add_router_settings_from_db_config`` calls ``llm_router.update_settings(retry_policy={...})`` after reading the From 7b86b7f4cd7576aeba56d9721f28002a4e5c6383 Mon Sep 17 00:00:00 2001 From: mateo Date: Wed, 2 Sep 2026 02:45:45 +0000 Subject: [PATCH 22/48] test: isolate router callback state Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/test_litellm/test_router_retry_policy_update.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/tests/test_litellm/test_router_retry_policy_update.py b/tests/test_litellm/test_router_retry_policy_update.py index 2c23d0da7e7..db710f76887 100644 --- a/tests/test_litellm/test_router_retry_policy_update.py +++ b/tests/test_litellm/test_router_retry_policy_update.py @@ -21,6 +21,7 @@ This file pins both halves of the fix. import json from dataclasses import dataclass +from typing import Final from unittest.mock import AsyncMock, MagicMock import pytest @@ -33,6 +34,14 @@ from litellm.router_utils.pre_call_checks.model_rate_limit_check import ModelRat from litellm.router_utils.pre_call_checks.prompt_caching_deployment_check import PromptCachingDeploymentCheck from litellm.types.router import RetryPolicy, UpdateRouterConfig + +@pytest.fixture(autouse=True) +def isolate_litellm_callbacks(): + callbacks_before: Final = litellm.callbacks.copy() + yield + litellm.callbacks = callbacks_before + + # --------------------------------------------------------------------------- # UpdateRouterConfig schema membership (LIT-3152 part 1) # --------------------------------------------------------------------------- From 3dea586d3be34c52156fe764051dd8d750d4570b Mon Sep 17 00:00:00 2001 From: mateo Date: Wed, 2 Sep 2026 02:48:09 +0000 Subject: [PATCH 23/48] test: cover runtime callback reconciliation Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../test_router_retry_policy_update.py | 20 +++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/tests/test_litellm/test_router_retry_policy_update.py b/tests/test_litellm/test_router_retry_policy_update.py index db710f76887..9e0bb0b9bef 100644 --- a/tests/test_litellm/test_router_retry_policy_update.py +++ b/tests/test_litellm/test_router_retry_policy_update.py @@ -135,6 +135,26 @@ def test_update_settings_clears_omitted_toggleable_pre_call_checks(): assert not any(isinstance(callback, PromptCachingDeploymentCheck) for callback in litellm.callbacks) +def test_set_optional_pre_call_checks_reconciles_callback_types(): + router = _build_router() + + router.set_optional_pre_call_checks(["prompt_caching"]) + router.set_optional_pre_call_checks([]) + + assert not any(isinstance(callback, PromptCachingDeploymentCheck) for callback in (router.optional_callbacks or [])) + assert not any(isinstance(callback, PromptCachingDeploymentCheck) for callback in litellm.callbacks) + + +def test_remove_optional_pre_call_check_removes_local_and_global_callbacks(): + router = _build_router() + + router.set_optional_pre_call_checks(["prompt_caching"]) + router._remove_optional_callbacks_of_type(PromptCachingDeploymentCheck) + + assert not any(isinstance(callback, PromptCachingDeploymentCheck) for callback in (router.optional_callbacks or [])) + assert not any(isinstance(callback, PromptCachingDeploymentCheck) for callback in litellm.callbacks) + + def test_update_settings_replaces_toggleable_pre_call_checks(): router = _build_router() From 70a4f74a0d65bd89a9e9127076efd4301b987a29 Mon Sep 17 00:00:00 2001 From: mateo Date: Wed, 2 Sep 2026 02:56:23 +0000 Subject: [PATCH 24/48] test: allow callback state fixture mutation Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/test_litellm/test_router_retry_policy_update.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_litellm/test_router_retry_policy_update.py b/tests/test_litellm/test_router_retry_policy_update.py index 9e0bb0b9bef..26251290176 100644 --- a/tests/test_litellm/test_router_retry_policy_update.py +++ b/tests/test_litellm/test_router_retry_policy_update.py @@ -39,7 +39,7 @@ from litellm.types.router import RetryPolicy, UpdateRouterConfig def isolate_litellm_callbacks(): callbacks_before: Final = litellm.callbacks.copy() yield - litellm.callbacks = callbacks_before + litellm.callbacks = callbacks_before # test-quality-ok: required callback-state restoration fixture # --------------------------------------------------------------------------- From 1d3e26fd98b40d40f498e14b2470e8acc79fb9f6 Mon Sep 17 00:00:00 2001 From: mateo Date: Wed, 2 Sep 2026 03:28:04 +0000 Subject: [PATCH 25/48] fix: preserve shared optional callbacks Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/router.py | 19 +++++---- .../test_router_retry_policy_update.py | 41 ++++++++++++++++++- 2 files changed, 50 insertions(+), 10 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index e2865542e89..e7588d8ad5a 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -2099,16 +2099,19 @@ class Router: self.add_optional_pre_call_checks(optional_pre_call_checks) def _remove_optional_callbacks_of_type(self, callback_cls: type[CustomLogger]) -> None: - if self.optional_callbacks is None: + if self.optional_callbacks is None or not any(type(cb) is callback_cls for cb in self.optional_callbacks): return - removed: Final = [cb for cb in self.optional_callbacks if isinstance(cb, callback_cls)] - if not removed: + self.optional_callbacks = [cb for cb in self.optional_callbacks if type(cb) is not callback_cls] + if any( + router is not self and any(type(cb) is callback_cls for cb in (router.optional_callbacks or [])) + for router in tuple(_live_routers) + ): return - self.optional_callbacks = [cb for cb in self.optional_callbacks if not isinstance(cb, callback_cls)] - for cb in removed: - litellm.logging_callback_manager.remove_callback_from_list_by_object( - litellm.callbacks, cb, require_self=False - ) + for cb in tuple(litellm.callbacks): + if type(cb) is callback_cls: + litellm.logging_callback_manager.remove_callback_from_list_by_object( + litellm.callbacks, cb, require_self=False + ) def print_deployment(self, deployment: dict): """ diff --git a/tests/test_litellm/test_router_retry_policy_update.py b/tests/test_litellm/test_router_retry_policy_update.py index 26251290176..be568134763 100644 --- a/tests/test_litellm/test_router_retry_policy_update.py +++ b/tests/test_litellm/test_router_retry_policy_update.py @@ -151,8 +151,45 @@ def test_remove_optional_pre_call_check_removes_local_and_global_callbacks(): router.set_optional_pre_call_checks(["prompt_caching"]) router._remove_optional_callbacks_of_type(PromptCachingDeploymentCheck) - assert not any(isinstance(callback, PromptCachingDeploymentCheck) for callback in (router.optional_callbacks or [])) - assert not any(isinstance(callback, PromptCachingDeploymentCheck) for callback in litellm.callbacks) + assert not any(type(callback) is PromptCachingDeploymentCheck for callback in (router.optional_callbacks or [])) + assert not any(type(callback) is PromptCachingDeploymentCheck for callback in litellm.callbacks) + + +def test_remove_optional_pre_call_check_keeps_global_callback_for_another_router(): + router_a = _build_router() + router_b = _build_router() + + router_a.update_settings(optional_pre_call_checks=["prompt_caching"]) + router_b.update_settings(optional_pre_call_checks=["prompt_caching"]) + + router_a.update_settings(optional_pre_call_checks=[]) + + assert not any(type(callback) is PromptCachingDeploymentCheck for callback in (router_a.optional_callbacks or [])) + assert any(type(callback) is PromptCachingDeploymentCheck for callback in (router_b.optional_callbacks or [])) + assert any(type(callback) is PromptCachingDeploymentCheck for callback in litellm.callbacks) + + router_b.update_settings(optional_pre_call_checks=[]) + + assert not any(type(callback) is PromptCachingDeploymentCheck for callback in (router_b.optional_callbacks or [])) + assert not any(type(callback) is PromptCachingDeploymentCheck for callback in litellm.callbacks) + + +def test_remove_optional_pre_call_check_keeps_global_callback_when_second_router_clears_first(): + router_a = _build_router() + router_b = _build_router() + + router_a.update_settings(optional_pre_call_checks=["prompt_caching"]) + router_b.update_settings(optional_pre_call_checks=["prompt_caching"]) + + router_b.update_settings(optional_pre_call_checks=[]) + + assert any(type(callback) is PromptCachingDeploymentCheck for callback in (router_a.optional_callbacks or [])) + assert not any(type(callback) is PromptCachingDeploymentCheck for callback in (router_b.optional_callbacks or [])) + assert any(type(callback) is PromptCachingDeploymentCheck for callback in litellm.callbacks) + + router_a.update_settings(optional_pre_call_checks=[]) + + assert not any(type(callback) is PromptCachingDeploymentCheck for callback in litellm.callbacks) def test_update_settings_replaces_toggleable_pre_call_checks(): From 974b331a4d09e2883d6fe84bb87ce57cc49ab5bb Mon Sep 17 00:00:00 2001 From: mateo Date: Wed, 2 Sep 2026 06:15:52 +0000 Subject: [PATCH 26/48] fix: accept persistable router settings Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/constants.py | 1 + litellm/proxy/proxy_server.py | 32 +++++++++----- .../proxy/proxy_server/test_routes_config.py | 43 +++++++++++++++++++ 3 files changed, 65 insertions(+), 11 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index 5b44e8b5f51..b6cb6b32187 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -31,6 +31,7 @@ RUNTIME_UPDATABLE_ROUTER_SETTINGS: Final[frozenset[str]] = frozenset( "optional_pre_call_checks", } ) +ROUTER_SETTINGS_MANAGED_OUTSIDE_CONFIG: Final[frozenset[str]] = frozenset({"model_list", "search_tools"}) DEFAULT_BATCH_SIZE: Final = int(os.getenv("DEFAULT_BATCH_SIZE", 512)) DEFAULT_FLUSH_INTERVAL_SECONDS: Final = int(os.getenv("DEFAULT_FLUSH_INTERVAL_SECONDS", 5)) DEFAULT_S3_FLUSH_INTERVAL_SECONDS: Final = int(os.getenv("DEFAULT_S3_FLUSH_INTERVAL_SECONDS", 10)) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index dbeb8486539..ae5e08ab566 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -254,6 +254,7 @@ from litellm.constants import ( PROXY_BUDGET_RESCHEDULER_MAX_TIME, PROXY_BUDGET_RESCHEDULER_MIN_TIME, PROXY_CONFIG_RELOAD_INTERVAL_SECONDS, + ROUTER_SETTINGS_MANAGED_OUTSIDE_CONFIG, USER_SPEND_ALERTS_JOB_ID, WEEKLY_SPEND_REPORT_JOB_ID, ) @@ -5711,13 +5712,9 @@ class ProxyConfig: router_settings: Final = config.get("router_settings", None) if router_settings and isinstance(router_settings, dict): - # model list and search_tools already set - exclude_args: Final = { - "model_list", - "search_tools", - } - - available_args: Final = [x for x in litellm.Router.get_valid_args() if x not in exclude_args] + available_args: Final = [ + x for x in litellm.Router.get_valid_args() if x not in ROUTER_SETTINGS_MANAGED_OUTSIDE_CONFIG + ] for k, v in router_settings.items(): if k in available_args: @@ -16229,14 +16226,17 @@ async def update_config( ) raw_router_settings: Final = request_body.get("router_settings") if isinstance(raw_router_settings, dict): - unsupported_router_settings: Final = sorted(set(raw_router_settings) - RUNTIME_UPDATABLE_ROUTER_SETTINGS) + supported_router_settings: Final = RUNTIME_UPDATABLE_ROUTER_SETTINGS | ( + frozenset(litellm.Router.get_valid_args()) - ROUTER_SETTINGS_MANAGED_OUTSIDE_CONFIG + ) + unsupported_router_settings: Final = sorted(set(raw_router_settings) - supported_router_settings) if unsupported_router_settings: raise HTTPException( status_code=400, detail={ "error": ( f"Unsupported router settings: {', '.join(unsupported_router_settings)} " - "are not runtime-updatable router settings" + "are not valid router settings" ) }, ) @@ -16342,10 +16342,20 @@ async def update_config( ) # router_settings: merge existing + request, request wins. - if config_info.router_settings is not None: + if isinstance(raw_router_settings, dict): existing = await _read_section("router_settings") before_router_settings: Final = copy.deepcopy(existing) - updates = config_info.router_settings.dict(exclude_none=True) + typed_router_settings: Final = ( + config_info.router_settings.dict(exclude_none=True) + if config_info.router_settings is not None + else {} + ) + raw_router_settings_without_none: Final = { + key: value + for key, value in raw_router_settings.items() + if key not in typed_router_settings and value is not None + } + updates: Final = {**typed_router_settings, **raw_router_settings_without_none} new_router_settings: Final = {**existing, **updates} await _upsert_section("router_settings", new_router_settings) asyncio.create_task( diff --git a/tests/test_litellm/proxy/proxy_server/test_routes_config.py b/tests/test_litellm/proxy/proxy_server/test_routes_config.py index 4b0954c350a..6166513d229 100644 --- a/tests/test_litellm/proxy/proxy_server/test_routes_config.py +++ b/tests/test_litellm/proxy/proxy_server/test_routes_config.py @@ -81,6 +81,49 @@ def test_config_update_persists_optional_pre_call_checks(client, auth_as, mock_p assert persisted["optional_pre_call_checks"] == ["prompt_caching"] +def test_config_update_persists_model_group_affinity_config(client, auth_as, mock_prisma, monkeypatch): + from litellm.proxy import proxy_server as ps + from litellm.proxy._types import LitellmUserRoles + + table = _install_litellm_config(mock_prisma) + monkeypatch.setattr(ps, "prisma_client", mock_prisma) + fake_proxy_config = MagicMock() + fake_proxy_config.add_deployment = AsyncMock() + monkeypatch.setattr(ps, "proxy_config", fake_proxy_config) + + model_group_affinity_config = {"gpt-4": ["session_affinity"]} + with auth_as(LitellmUserRoles.PROXY_ADMIN): + response = client.post( + "/config/update", + json={"router_settings": {"model_group_affinity_config": model_group_affinity_config}}, + ) + + assert response.status_code == 200 + persisted = json.loads(table.upsert.call_args.kwargs["data"]["create"]["param_value"]) + assert persisted["model_group_affinity_config"] == model_group_affinity_config + + +def test_config_update_persists_disable_cooldowns(client, auth_as, mock_prisma, monkeypatch): + from litellm.proxy import proxy_server as ps + from litellm.proxy._types import LitellmUserRoles + + table = _install_litellm_config(mock_prisma) + monkeypatch.setattr(ps, "prisma_client", mock_prisma) + fake_proxy_config = MagicMock() + fake_proxy_config.add_deployment = AsyncMock() + monkeypatch.setattr(ps, "proxy_config", fake_proxy_config) + + with auth_as(LitellmUserRoles.PROXY_ADMIN): + response = client.post( + "/config/update", + json={"router_settings": {"disable_cooldowns": True}}, + ) + + assert response.status_code == 200 + persisted = json.loads(table.upsert.call_args.kwargs["data"]["create"]["param_value"]) + assert persisted["disable_cooldowns"] is True + + def test_config_update_rejects_unknown_router_setting(client, auth_as, mock_prisma, monkeypatch): from litellm.proxy import proxy_server as ps from litellm.proxy._types import LitellmUserRoles From 95e09db661471d8fbdaa1a93f79bf8e49ecdd0eb Mon Sep 17 00:00:00 2001 From: mateo Date: Wed, 2 Sep 2026 06:56:27 +0000 Subject: [PATCH 27/48] style: apply ruff format to router settings merge Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/proxy_server.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index ae5e08ab566..1bdf2ba9987 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -16346,9 +16346,7 @@ async def update_config( existing = await _read_section("router_settings") before_router_settings: Final = copy.deepcopy(existing) typed_router_settings: Final = ( - config_info.router_settings.dict(exclude_none=True) - if config_info.router_settings is not None - else {} + config_info.router_settings.dict(exclude_none=True) if config_info.router_settings is not None else {} ) raw_router_settings_without_none: Final = { key: value From aba9644297e4e709232775031cb03f028dd1cbd6 Mon Sep 17 00:00:00 2001 From: mateo Date: Wed, 2 Sep 2026 07:10:04 +0000 Subject: [PATCH 28/48] fix: avoid router settings update name collision Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/proxy_server.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 1bdf2ba9987..6632b209905 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -16353,8 +16353,8 @@ async def update_config( for key, value in raw_router_settings.items() if key not in typed_router_settings and value is not None } - updates: Final = {**typed_router_settings, **raw_router_settings_without_none} - new_router_settings: Final = {**existing, **updates} + router_settings_updates: Final = {**typed_router_settings, **raw_router_settings_without_none} + new_router_settings: Final = {**existing, **router_settings_updates} await _upsert_section("router_settings", new_router_settings) asyncio.create_task( create_config_audit_log( From 385957e830c6b5edae9455dae8361527908cd4bb Mon Sep 17 00:00:00 2001 From: mateo Date: Wed, 2 Sep 2026 07:49:25 +0000 Subject: [PATCH 29/48] fix: reject constructor-managed router settings Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/constants.py | 11 +++++- .../proxy/proxy_server/test_routes_config.py | 36 +++++++++++++++++++ 2 files changed, 46 insertions(+), 1 deletion(-) diff --git a/litellm/constants.py b/litellm/constants.py index b6cb6b32187..1c1939bd350 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -31,7 +31,16 @@ RUNTIME_UPDATABLE_ROUTER_SETTINGS: Final[frozenset[str]] = frozenset( "optional_pre_call_checks", } ) -ROUTER_SETTINGS_MANAGED_OUTSIDE_CONFIG: Final[frozenset[str]] = frozenset({"model_list", "search_tools"}) +ROUTER_SETTINGS_MANAGED_OUTSIDE_CONFIG: Final[frozenset[str]] = frozenset( + { + "model_list", + "search_tools", + "assistants_config", + "router_general_settings", + "ignore_invalid_deployments", + "fallback_access_check", + } +) DEFAULT_BATCH_SIZE: Final = int(os.getenv("DEFAULT_BATCH_SIZE", 512)) DEFAULT_FLUSH_INTERVAL_SECONDS: Final = int(os.getenv("DEFAULT_FLUSH_INTERVAL_SECONDS", 5)) DEFAULT_S3_FLUSH_INTERVAL_SECONDS: Final = int(os.getenv("DEFAULT_S3_FLUSH_INTERVAL_SECONDS", 10)) diff --git a/tests/test_litellm/proxy/proxy_server/test_routes_config.py b/tests/test_litellm/proxy/proxy_server/test_routes_config.py index 6166513d229..dcb63b8ca82 100644 --- a/tests/test_litellm/proxy/proxy_server/test_routes_config.py +++ b/tests/test_litellm/proxy/proxy_server/test_routes_config.py @@ -124,6 +124,42 @@ def test_config_update_persists_disable_cooldowns(client, auth_as, mock_prisma, assert persisted["disable_cooldowns"] is True +def test_config_update_rejects_assistants_config(client, auth_as, mock_prisma, monkeypatch): + from litellm.proxy import proxy_server as ps + from litellm.proxy._types import LitellmUserRoles + + table = _install_litellm_config(mock_prisma) + monkeypatch.setattr(ps, "prisma_client", mock_prisma) + + with auth_as(LitellmUserRoles.PROXY_ADMIN): + response = client.post( + "/config/update", + json={"router_settings": {"assistants_config": {"enabled": True}}}, + ) + + assert response.status_code == 400 + assert "assistants_config" in response.json()["error"]["message"] + table.upsert.assert_not_called() + + +def test_config_update_rejects_router_general_settings(client, auth_as, mock_prisma, monkeypatch): + from litellm.proxy import proxy_server as ps + from litellm.proxy._types import LitellmUserRoles + + table = _install_litellm_config(mock_prisma) + monkeypatch.setattr(ps, "prisma_client", mock_prisma) + + with auth_as(LitellmUserRoles.PROXY_ADMIN): + response = client.post( + "/config/update", + json={"router_settings": {"router_general_settings": {"async_only_mode": True}}}, + ) + + assert response.status_code == 400 + assert "router_general_settings" in response.json()["error"]["message"] + table.upsert.assert_not_called() + + def test_config_update_rejects_unknown_router_setting(client, auth_as, mock_prisma, monkeypatch): from litellm.proxy import proxy_server as ps from litellm.proxy._types import LitellmUserRoles From a7836ede15bb4f62d8f44bdb991402a9829727e3 Mon Sep 17 00:00:00 2001 From: mateo Date: Wed, 2 Sep 2026 14:48:51 +0000 Subject: [PATCH 30/48] fix(models): absorb open registry PRs: govcloud bedrock and mantle, azure gov, openai tiered long-context, scaleway, together qwen3.8, azure ai cache and kimi k2.7 code, azure mai deprecations Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...odel_prices_and_context_window_backup.json | 605 +++++++++++++++++- model_prices_and_context_window.json | 605 +++++++++++++++++- .../llm_cost_calc/test_llm_cost_calc_utils.py | 6 +- ...penai_service_tier_long_context_pricing.py | 156 +++++ whitelisted_bedrock_models.txt | 14 + 5 files changed, 1341 insertions(+), 45 deletions(-) create mode 100644 tests/test_litellm/test_openai_service_tier_long_context_pricing.py diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index c710db1a749..87d348f4752 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -9643,7 +9643,8 @@ "supported_endpoints": [ "/v1/images/generations", "/v1/images/edits" - ] + ], + "deprecation_date": "2026-10-01" }, "azure_ai/MAI-Image-2.5-Flash": { "input_cost_per_image_token": 1.75e-06, @@ -9656,7 +9657,8 @@ "supported_endpoints": [ "/v1/images/generations", "/v1/images/edits" - ] + ], + "deprecation_date": "2026-10-01" }, "azure_ai/MAI-Image-2e": { "deprecation_date": "2026-08-15", @@ -10155,7 +10157,9 @@ "source": "https://azure.microsoft.com/en-us/pricing/details/ai-foundry-models/deepseek/", "supports_function_calling": true, "supports_reasoning": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "cache_read_input_token_cost": 1.45e-07, + "supports_prompt_caching": true }, "azure_ai/deepseek-v4-flash": { "deprecation_date": "2028-02-20", @@ -10169,18 +10173,20 @@ "source": "https://azure.microsoft.com/en-us/pricing/details/ai-foundry-models/deepseek/", "supports_function_calling": true, "supports_reasoning": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "cache_read_input_token_cost": 2.8e-08, + "supports_prompt_caching": true }, "azure_ai/deepseek-v4-flash-0731": { - "cache_read_input_token_cost": 2.8e-08, + "cache_read_input_token_cost": 1.4e-08, "deprecation_date": "2026-12-03", - "input_cost_per_token": 1.9e-07, + "input_cost_per_token": 4.4e-07, "litellm_provider": "azure_ai", "max_input_tokens": 1000000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", - "output_cost_per_token": 5.1e-07, + "output_cost_per_token": 1.32e-06, "source": "https://azure.microsoft.com/en-us/pricing/details/ai-foundry-models/deepseek/", "supports_function_calling": true, "supports_prompt_caching": true, @@ -10400,11 +10406,13 @@ "max_tokens": 262144, "mode": "chat", "output_cost_per_token": 3e-06, - "source": "https://techcommunity.microsoft.com/blog/azure-ai-foundry-blog/kimi-k2-5-now-in-microsoft-foundry/4492321", + "source": "https://azure.microsoft.com/en-us/pricing/details/ai-foundry-models/kimi/", "supports_function_calling": true, "supports_tool_choice": true, "supports_video_input": true, - "supports_vision": true + "supports_vision": true, + "cache_read_input_token_cost": 1e-07, + "supports_prompt_caching": true }, "azure_ai/kimi-k2.6": { "deprecation_date": "2027-04-16", @@ -10415,7 +10423,7 @@ "max_tokens": 262144, "mode": "chat", "output_cost_per_token": 4e-06, - "source": "https://techcommunity.microsoft.com/blog/azure-ai-foundry-blog/introducing-kimi-k2-6-in-microsoft-foundry/4513125", + "source": "https://azure.microsoft.com/en-us/pricing/details/ai-foundry-models/kimi/", "supported_modalities": [ "text", "image" @@ -10426,7 +10434,9 @@ "supports_function_calling": true, "supports_reasoning": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "cache_read_input_token_cost": 1.6e-07, + "supports_prompt_caching": true }, "azure_ai/ministral-3b": { "input_cost_per_token": 4e-08, @@ -12110,7 +12120,7 @@ "max_output_tokens": 2048, "max_tokens": 2048, "mode": "chat", - "output_cost_per_token": 2.65e-06, + "output_cost_per_token": 6e-07, "supports_pdf_input": true }, "bedrock/us-west-1/meta.llama3-70b-instruct-v1:0": { @@ -29098,16 +29108,19 @@ "cache_creation_input_token_cost": 5e-06, "cache_creation_input_token_cost_above_272k_tokens": 1e-05, "cache_creation_input_token_cost_above_272k_tokens_flex": 5e-06, + "cache_creation_input_token_cost_above_272k_tokens_priority": 2e-05, "cache_creation_input_token_cost_flex": 2.5e-06, "cache_creation_input_token_cost_priority": 1e-05, "cache_read_input_token_cost": 4e-07, "cache_read_input_token_cost_above_272k_tokens": 8e-07, "cache_read_input_token_cost_above_272k_tokens_flex": 4e-07, + "cache_read_input_token_cost_above_272k_tokens_priority": 1.6e-06, "cache_read_input_token_cost_flex": 2e-07, "cache_read_input_token_cost_priority": 8e-07, "input_cost_per_token": 4e-06, "input_cost_per_token_above_272k_tokens": 8e-06, "input_cost_per_token_above_272k_tokens_flex": 4e-06, + "input_cost_per_token_above_272k_tokens_priority": 1.6e-05, "input_cost_per_token_batches": 2e-06, "input_cost_per_token_flex": 2e-06, "input_cost_per_token_priority": 8e-06, @@ -29119,6 +29132,7 @@ "output_cost_per_token": 2e-05, "output_cost_per_token_above_272k_tokens": 3e-05, "output_cost_per_token_above_272k_tokens_flex": 1.5e-05, + "output_cost_per_token_above_272k_tokens_priority": 6e-05, "output_cost_per_token_batches": 1e-05, "output_cost_per_token_flex": 1e-05, "output_cost_per_token_priority": 4e-05, @@ -29161,16 +29175,19 @@ "cache_creation_input_token_cost": 5e-06, "cache_creation_input_token_cost_above_272k_tokens": 1e-05, "cache_creation_input_token_cost_above_272k_tokens_flex": 5e-06, + "cache_creation_input_token_cost_above_272k_tokens_priority": 2e-05, "cache_creation_input_token_cost_flex": 2.5e-06, "cache_creation_input_token_cost_priority": 1e-05, "cache_read_input_token_cost": 4e-07, "cache_read_input_token_cost_above_272k_tokens": 8e-07, "cache_read_input_token_cost_above_272k_tokens_flex": 4e-07, + "cache_read_input_token_cost_above_272k_tokens_priority": 1.6e-06, "cache_read_input_token_cost_flex": 2e-07, "cache_read_input_token_cost_priority": 8e-07, "input_cost_per_token": 4e-06, "input_cost_per_token_above_272k_tokens": 8e-06, "input_cost_per_token_above_272k_tokens_flex": 4e-06, + "input_cost_per_token_above_272k_tokens_priority": 1.6e-05, "input_cost_per_token_batches": 2e-06, "input_cost_per_token_flex": 2e-06, "input_cost_per_token_priority": 8e-06, @@ -29182,6 +29199,7 @@ "output_cost_per_token": 2e-05, "output_cost_per_token_above_272k_tokens": 3e-05, "output_cost_per_token_above_272k_tokens_flex": 1.5e-05, + "output_cost_per_token_above_272k_tokens_priority": 6e-05, "output_cost_per_token_batches": 1e-05, "output_cost_per_token_flex": 1e-05, "output_cost_per_token_priority": 4e-05, @@ -29225,16 +29243,19 @@ "cache_creation_input_token_cost": 2.5e-06, "cache_creation_input_token_cost_above_272k_tokens": 5e-06, "cache_creation_input_token_cost_above_272k_tokens_flex": 2.5e-06, + "cache_creation_input_token_cost_above_272k_tokens_priority": 1e-05, "cache_creation_input_token_cost_flex": 1.25e-06, "cache_creation_input_token_cost_priority": 5e-06, "cache_read_input_token_cost": 2e-07, "cache_read_input_token_cost_above_272k_tokens": 4e-07, "cache_read_input_token_cost_above_272k_tokens_flex": 2e-07, + "cache_read_input_token_cost_above_272k_tokens_priority": 8e-07, "cache_read_input_token_cost_flex": 1e-07, "cache_read_input_token_cost_priority": 4e-07, "input_cost_per_token": 2e-06, "input_cost_per_token_above_272k_tokens": 4e-06, "input_cost_per_token_above_272k_tokens_flex": 2e-06, + "input_cost_per_token_above_272k_tokens_priority": 8e-06, "input_cost_per_token_batches": 1e-06, "input_cost_per_token_flex": 1e-06, "input_cost_per_token_priority": 4e-06, @@ -29246,6 +29267,7 @@ "output_cost_per_token": 1.2e-05, "output_cost_per_token_above_272k_tokens": 1.8e-05, "output_cost_per_token_above_272k_tokens_flex": 9e-06, + "output_cost_per_token_above_272k_tokens_priority": 3.6e-05, "output_cost_per_token_batches": 6e-06, "output_cost_per_token_flex": 6e-06, "output_cost_per_token_priority": 2.4e-05, @@ -29288,16 +29310,19 @@ "cache_creation_input_token_cost": 2.5e-07, "cache_creation_input_token_cost_above_272k_tokens": 5e-07, "cache_creation_input_token_cost_above_272k_tokens_flex": 2.5e-07, + "cache_creation_input_token_cost_above_272k_tokens_priority": 1e-06, "cache_creation_input_token_cost_flex": 1.25e-07, "cache_creation_input_token_cost_priority": 5e-07, "cache_read_input_token_cost": 2e-08, "cache_read_input_token_cost_above_272k_tokens": 4e-08, "cache_read_input_token_cost_above_272k_tokens_flex": 2e-08, + "cache_read_input_token_cost_above_272k_tokens_priority": 8e-08, "cache_read_input_token_cost_flex": 1e-08, "cache_read_input_token_cost_priority": 4e-08, "input_cost_per_token": 2e-07, "input_cost_per_token_above_272k_tokens": 4e-07, "input_cost_per_token_above_272k_tokens_flex": 2e-07, + "input_cost_per_token_above_272k_tokens_priority": 8e-07, "input_cost_per_token_batches": 1e-07, "input_cost_per_token_flex": 1e-07, "input_cost_per_token_priority": 4e-07, @@ -29309,6 +29334,7 @@ "output_cost_per_token": 1.2e-06, "output_cost_per_token_above_272k_tokens": 1.8e-06, "output_cost_per_token_above_272k_tokens_flex": 9e-07, + "output_cost_per_token_above_272k_tokens_priority": 3.6e-06, "output_cost_per_token_batches": 6e-07, "output_cost_per_token_flex": 6e-07, "output_cost_per_token_priority": 2.4e-06, @@ -29548,7 +29574,10 @@ "supports_web_search": true, "supports_none_reasoning_effort": true, "supports_xhigh_reasoning_effort": true, - "supports_minimal_reasoning_effort": false + "supports_minimal_reasoning_effort": false, + "input_cost_per_token_above_272k_tokens_flex": 5e-06, + "output_cost_per_token_above_272k_tokens_flex": 2.25e-05, + "cache_read_input_token_cost_above_272k_tokens_flex": 5e-07 }, "gpt-5.5-2026-04-23": { "cache_read_input_token_cost": 5e-07, @@ -29602,7 +29631,10 @@ "supports_web_search": true, "supports_none_reasoning_effort": true, "supports_xhigh_reasoning_effort": true, - "supports_minimal_reasoning_effort": false + "supports_minimal_reasoning_effort": false, + "input_cost_per_token_above_272k_tokens_flex": 5e-06, + "output_cost_per_token_above_272k_tokens_flex": 2.25e-05, + "cache_read_input_token_cost_above_272k_tokens_flex": 5e-07 }, "gpt-5.5-pro": { "cache_read_input_token_cost": 3e-06, @@ -29751,7 +29783,10 @@ "supports_none_reasoning_effort": true, "default_reasoning_effort": "none", "supports_xhigh_reasoning_effort": true, - "supports_minimal_reasoning_effort": true + "supports_minimal_reasoning_effort": true, + "input_cost_per_token_above_272k_tokens_flex": 2.5e-06, + "output_cost_per_token_above_272k_tokens_flex": 1.125e-05, + "cache_read_input_token_cost_above_272k_tokens_flex": 2.5e-07 }, "gpt-5.4-2026-03-05": { "cache_read_input_token_cost": 2.5e-07, @@ -29800,7 +29835,10 @@ "supports_none_reasoning_effort": true, "default_reasoning_effort": "none", "supports_xhigh_reasoning_effort": true, - "supports_minimal_reasoning_effort": true + "supports_minimal_reasoning_effort": true, + "input_cost_per_token_above_272k_tokens_flex": 2.5e-06, + "output_cost_per_token_above_272k_tokens_flex": 1.125e-05, + "cache_read_input_token_cost_above_272k_tokens_flex": 2.5e-07 }, "gpt-5.4-pro": { "cache_read_input_token_cost": 3e-06, @@ -29849,7 +29887,9 @@ "supports_web_search": true, "supports_none_reasoning_effort": false, "supports_xhigh_reasoning_effort": true, - "supports_minimal_reasoning_effort": true + "supports_minimal_reasoning_effort": true, + "input_cost_per_token_above_272k_tokens_flex": 3e-05, + "output_cost_per_token_above_272k_tokens_flex": 0.000135 }, "gpt-5.4-pro-2026-03-05": { "cache_read_input_token_cost": 3e-06, @@ -29898,7 +29938,9 @@ "supports_web_search": true, "supports_none_reasoning_effort": false, "supports_xhigh_reasoning_effort": true, - "supports_minimal_reasoning_effort": true + "supports_minimal_reasoning_effort": true, + "input_cost_per_token_above_272k_tokens_flex": 3e-05, + "output_cost_per_token_above_272k_tokens_flex": 0.000135 }, "gpt-5.4-mini": { "cache_read_input_token_cost": 7.5e-08, @@ -41350,13 +41392,13 @@ "source": "https://docs.together.ai/docs/serverless-models" }, "together_ai/Qwen/Qwen3.8-2.4T-A95B": { - "cache_read_input_token_cost": 5e-07, - "input_cost_per_token": 2.5e-06, + "cache_read_input_token_cost": 2.5e-07, + "input_cost_per_token": 2e-06, "litellm_provider": "together_ai", "max_input_tokens": 1010000, "max_tokens": 1010000, "mode": "chat", - "output_cost_per_token": 6.25e-06, + "output_cost_per_token": 6e-06, "source": "https://docs.together.ai/docs/serverless-models", "supports_prompt_caching": true }, @@ -57558,5 +57600,526 @@ "supported_endpoints": [ "/v1/audio/transcriptions" ] + }, + "scaleway/glm-5.2": { + "input_cost_per_token": 1.8e-06, + "litellm_provider": "scaleway", + "max_input_tokens": 256000, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "output_cost_per_token": 5.5e-06, + "source": "https://www.scaleway.com/en/pricing/model-as-a-service/", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_vision": false + }, + "scaleway/deepseek-v4-flash-0731": { + "cache_read_input_token_cost": 8e-08, + "input_cost_per_token": 4e-07, + "litellm_provider": "scaleway", + "max_input_tokens": 256000, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 8e-07, + "source": "https://www.scaleway.com/en/pricing/model-as-a-service/", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_vision": false + }, + "azure_ai/kimi-k2.7-code": { + "cache_read_input_token_cost": 1.9e-07, + "deprecation_date": "2026-10-03", + "input_cost_per_token": 9.5e-07, + "litellm_provider": "azure_ai", + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 4e-06, + "source": "https://azure.microsoft.com/en-us/pricing/details/ai-foundry-models/kimi/", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "bedrock/us-gov-west-1/nvidia.nemotron-nano-3-30b": { + "input_cost_per_token": 7.2e-08, + "litellm_provider": "bedrock", + "max_input_tokens": 262144, + "max_output_tokens": 8192, + "max_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 2.88e-07, + "source": "https://aws.amazon.com/bedrock/pricing/", + "supports_function_calling": true, + "supports_native_structured_output": true, + "supports_system_messages": true, + "supports_tool_choice": true + }, + "bedrock/us-gov-west-1/nvidia.nemotron-nano-12b-v2": { + "input_cost_per_token": 2.4e-07, + "litellm_provider": "bedrock", + "max_input_tokens": 128000, + "max_output_tokens": 8192, + "max_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 7.2e-07, + "supports_system_messages": true, + "supports_vision": true + }, + "bedrock/us-gov-west-1/nvidia.nemotron-super-3-120b": { + "input_cost_per_token": 1.8e-07, + "litellm_provider": "bedrock", + "max_input_tokens": 256000, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 7.8e-07, + "source": "https://aws.amazon.com/bedrock/pricing/", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_system_messages": true, + "supports_tool_choice": true + }, + "bedrock/us-gov-west-1/openai.gpt-oss-20b-1:0": { + "input_cost_per_token": 8.4e-08, + "litellm_provider": "bedrock", + "max_input_tokens": 128000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 3.6e-07, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "bedrock/us-gov-west-1/openai.gpt-oss-120b-1:0": { + "input_cost_per_token": 1.8e-07, + "litellm_provider": "bedrock", + "max_input_tokens": 128000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 7.2e-07, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "bedrock/us-gov-west-1/anthropic.claude-sonnet-5": { + "bedrock_converse_supports_strict_tools": false, + "bedrock_output_config_effort_ceiling": "xhigh", + "cache_creation_input_token_cost": 3e-06, + "cache_creation_input_token_cost_above_1hr": 4.8e-06, + "cache_read_input_token_cost": 2.4e-07, + "input_cost_per_token": 2.4e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 1.2e-05, + "prompt_cache_min_tokens": 1024, + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_mid_conversation_system": true, + "supports_native_structured_output": false, + "supports_output_config": true, + "supports_parallel_tool_use_config": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true + }, + "bedrock/us-gov-west-1/anthropic.claude-opus-4-8": { + "bedrock_converse_supports_strict_tools": false, + "bedrock_output_config_effort_ceiling": "xhigh", + "cache_creation_input_token_cost": 7.5e-06, + "cache_creation_input_token_cost_above_1hr": 1.2e-05, + "cache_read_input_token_cost": 6e-07, + "input_cost_per_token": 6e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 3e-05, + "prompt_cache_min_tokens": 1024, + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_mid_conversation_system": true, + "supports_native_structured_output": true, + "supports_output_config": true, + "supports_parallel_tool_use_config": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true + }, + "bedrock/us-gov-east-1/nvidia.nemotron-nano-3-30b": { + "input_cost_per_token": 7.2e-08, + "litellm_provider": "bedrock", + "max_input_tokens": 262144, + "max_output_tokens": 8192, + "max_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 2.88e-07, + "source": "https://aws.amazon.com/bedrock/pricing/", + "supports_function_calling": true, + "supports_native_structured_output": true, + "supports_system_messages": true, + "supports_tool_choice": true + }, + "bedrock/us-gov-east-1/nvidia.nemotron-nano-12b-v2": { + "input_cost_per_token": 2.4e-07, + "litellm_provider": "bedrock", + "max_input_tokens": 128000, + "max_output_tokens": 8192, + "max_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 7.2e-07, + "supports_system_messages": true, + "supports_vision": true + }, + "bedrock/us-gov-east-1/nvidia.nemotron-super-3-120b": { + "input_cost_per_token": 1.8e-07, + "litellm_provider": "bedrock", + "max_input_tokens": 256000, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 7.8e-07, + "source": "https://aws.amazon.com/bedrock/pricing/", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_system_messages": true, + "supports_tool_choice": true + }, + "bedrock/us-gov-east-1/openai.gpt-oss-20b-1:0": { + "input_cost_per_token": 8.4e-08, + "litellm_provider": "bedrock", + "max_input_tokens": 128000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 3.6e-07, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "bedrock/us-gov-east-1/openai.gpt-oss-120b-1:0": { + "input_cost_per_token": 1.8e-07, + "litellm_provider": "bedrock", + "max_input_tokens": 128000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 7.2e-07, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "bedrock/us-gov-east-1/anthropic.claude-sonnet-5": { + "bedrock_converse_supports_strict_tools": false, + "bedrock_output_config_effort_ceiling": "xhigh", + "cache_creation_input_token_cost": 3e-06, + "cache_creation_input_token_cost_above_1hr": 4.8e-06, + "cache_read_input_token_cost": 2.4e-07, + "input_cost_per_token": 2.4e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 1.2e-05, + "prompt_cache_min_tokens": 1024, + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_mid_conversation_system": true, + "supports_native_structured_output": false, + "supports_output_config": true, + "supports_parallel_tool_use_config": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true + }, + "bedrock/us-gov-east-1/anthropic.claude-opus-4-8": { + "bedrock_converse_supports_strict_tools": false, + "bedrock_output_config_effort_ceiling": "xhigh", + "cache_creation_input_token_cost": 7.5e-06, + "cache_creation_input_token_cost_above_1hr": 1.2e-05, + "cache_read_input_token_cost": 6e-07, + "input_cost_per_token": 6e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 3e-05, + "prompt_cache_min_tokens": 1024, + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_mid_conversation_system": true, + "supports_native_structured_output": true, + "supports_output_config": true, + "supports_parallel_tool_use_config": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true + }, + "bedrock_mantle/us-gov-west-1/openai.gpt-5.6-terra": { + "litellm_provider": "bedrock_mantle", + "max_input_tokens": 1050000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "responses", + "use_openai_responses_path": true, + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true, + "input_cost_per_token": 2.64e-06, + "input_cost_per_token_above_272k_tokens": 5.28e-06, + "cache_creation_input_token_cost": 3.3e-06, + "cache_creation_input_token_cost_above_272k_tokens": 6.6e-06, + "cache_read_input_token_cost": 2.64e-07, + "cache_read_input_token_cost_above_272k_tokens": 5.28e-07, + "output_cost_per_token": 1.584e-05, + "output_cost_per_token_above_272k_tokens": 2.376e-05 + }, + "bedrock_mantle/us-gov-west-1/openai.gpt-5.6-luna": { + "litellm_provider": "bedrock_mantle", + "max_input_tokens": 1050000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "responses", + "use_openai_responses_path": true, + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true, + "input_cost_per_token": 2.64e-07, + "input_cost_per_token_above_272k_tokens": 5.28e-07, + "cache_creation_input_token_cost": 3.3e-07, + "cache_creation_input_token_cost_above_272k_tokens": 6.6e-07, + "cache_read_input_token_cost": 2.64e-08, + "cache_read_input_token_cost_above_272k_tokens": 5.28e-08, + "output_cost_per_token": 1.584e-06, + "output_cost_per_token_above_272k_tokens": 2.376e-06 + }, + "bedrock_mantle/us-gov-west-1/openai.gpt-5.4": { + "litellm_provider": "bedrock_mantle", + "max_input_tokens": 1050000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "responses", + "use_openai_responses_path": true, + "supported_endpoints": [ + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true, + "input_cost_per_token": 3.3e-06, + "cache_read_input_token_cost": 3.3e-07, + "output_cost_per_token": 1.98e-05 + }, + "bedrock_mantle/us-gov-west-1/xai.grok-4.3": { + "use_openai_responses_path": true, + "litellm_provider": "bedrock_mantle", + "max_input_tokens": 131072, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true, + "source": "https://aws.amazon.com/bedrock/pricing/", + "input_cost_per_token": 1.5e-06, + "output_cost_per_token": 3e-06, + "cache_read_input_token_cost": 2.4e-07 + }, + "bedrock_mantle/us-gov-east-1/openai.gpt-5.4": { + "litellm_provider": "bedrock_mantle", + "max_input_tokens": 1050000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "responses", + "use_openai_responses_path": true, + "supported_endpoints": [ + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true, + "input_cost_per_token": 3.3e-06, + "cache_read_input_token_cost": 3.3e-07, + "output_cost_per_token": 1.98e-05 + }, + "azure/us-gov/gpt-5.1": { + "cache_read_input_token_cost": 1.71875e-07, + "default_reasoning_effort": "none", + "deprecation_date": "2027-05-15", + "input_cost_per_token": 1.71875e-06, + "litellm_provider": "azure", + "max_input_tokens": 272000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 1.375e-05, + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/batch", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text", + "image" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_none_reasoning_effort": true, + "supports_parallel_function_calling": 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 + }, + "azure/us-gov/o3-mini": { + "cache_read_input_token_cost": 7.57e-07, + "deprecation_date": "2026-10-01", + "input_cost_per_token": 1.513e-06, + "litellm_provider": "azure", + "max_input_tokens": 200000, + "max_output_tokens": 100000, + "max_tokens": 100000, + "mode": "chat", + "output_cost_per_token": 6.05e-06, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": false + }, + "azure/us-gov/text-embedding-3-large": { + "deprecation_date": "2028-02-09", + "input_cost_per_token": 1.63e-07, + "litellm_provider": "azure", + "max_input_tokens": 8191, + "max_tokens": 8191, + "mode": "embedding", + "output_cost_per_token": 0.0 + }, + "azure/us-gov/text-embedding-3-small": { + "deprecation_date": "2028-02-09", + "input_cost_per_token": 2.5e-08, + "litellm_provider": "azure", + "max_input_tokens": 8191, + "max_tokens": 8191, + "mode": "embedding", + "output_cost_per_token": 0.0 } } diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index c710db1a749..87d348f4752 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -9643,7 +9643,8 @@ "supported_endpoints": [ "/v1/images/generations", "/v1/images/edits" - ] + ], + "deprecation_date": "2026-10-01" }, "azure_ai/MAI-Image-2.5-Flash": { "input_cost_per_image_token": 1.75e-06, @@ -9656,7 +9657,8 @@ "supported_endpoints": [ "/v1/images/generations", "/v1/images/edits" - ] + ], + "deprecation_date": "2026-10-01" }, "azure_ai/MAI-Image-2e": { "deprecation_date": "2026-08-15", @@ -10155,7 +10157,9 @@ "source": "https://azure.microsoft.com/en-us/pricing/details/ai-foundry-models/deepseek/", "supports_function_calling": true, "supports_reasoning": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "cache_read_input_token_cost": 1.45e-07, + "supports_prompt_caching": true }, "azure_ai/deepseek-v4-flash": { "deprecation_date": "2028-02-20", @@ -10169,18 +10173,20 @@ "source": "https://azure.microsoft.com/en-us/pricing/details/ai-foundry-models/deepseek/", "supports_function_calling": true, "supports_reasoning": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "cache_read_input_token_cost": 2.8e-08, + "supports_prompt_caching": true }, "azure_ai/deepseek-v4-flash-0731": { - "cache_read_input_token_cost": 2.8e-08, + "cache_read_input_token_cost": 1.4e-08, "deprecation_date": "2026-12-03", - "input_cost_per_token": 1.9e-07, + "input_cost_per_token": 4.4e-07, "litellm_provider": "azure_ai", "max_input_tokens": 1000000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", - "output_cost_per_token": 5.1e-07, + "output_cost_per_token": 1.32e-06, "source": "https://azure.microsoft.com/en-us/pricing/details/ai-foundry-models/deepseek/", "supports_function_calling": true, "supports_prompt_caching": true, @@ -10400,11 +10406,13 @@ "max_tokens": 262144, "mode": "chat", "output_cost_per_token": 3e-06, - "source": "https://techcommunity.microsoft.com/blog/azure-ai-foundry-blog/kimi-k2-5-now-in-microsoft-foundry/4492321", + "source": "https://azure.microsoft.com/en-us/pricing/details/ai-foundry-models/kimi/", "supports_function_calling": true, "supports_tool_choice": true, "supports_video_input": true, - "supports_vision": true + "supports_vision": true, + "cache_read_input_token_cost": 1e-07, + "supports_prompt_caching": true }, "azure_ai/kimi-k2.6": { "deprecation_date": "2027-04-16", @@ -10415,7 +10423,7 @@ "max_tokens": 262144, "mode": "chat", "output_cost_per_token": 4e-06, - "source": "https://techcommunity.microsoft.com/blog/azure-ai-foundry-blog/introducing-kimi-k2-6-in-microsoft-foundry/4513125", + "source": "https://azure.microsoft.com/en-us/pricing/details/ai-foundry-models/kimi/", "supported_modalities": [ "text", "image" @@ -10426,7 +10434,9 @@ "supports_function_calling": true, "supports_reasoning": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "cache_read_input_token_cost": 1.6e-07, + "supports_prompt_caching": true }, "azure_ai/ministral-3b": { "input_cost_per_token": 4e-08, @@ -12110,7 +12120,7 @@ "max_output_tokens": 2048, "max_tokens": 2048, "mode": "chat", - "output_cost_per_token": 2.65e-06, + "output_cost_per_token": 6e-07, "supports_pdf_input": true }, "bedrock/us-west-1/meta.llama3-70b-instruct-v1:0": { @@ -29098,16 +29108,19 @@ "cache_creation_input_token_cost": 5e-06, "cache_creation_input_token_cost_above_272k_tokens": 1e-05, "cache_creation_input_token_cost_above_272k_tokens_flex": 5e-06, + "cache_creation_input_token_cost_above_272k_tokens_priority": 2e-05, "cache_creation_input_token_cost_flex": 2.5e-06, "cache_creation_input_token_cost_priority": 1e-05, "cache_read_input_token_cost": 4e-07, "cache_read_input_token_cost_above_272k_tokens": 8e-07, "cache_read_input_token_cost_above_272k_tokens_flex": 4e-07, + "cache_read_input_token_cost_above_272k_tokens_priority": 1.6e-06, "cache_read_input_token_cost_flex": 2e-07, "cache_read_input_token_cost_priority": 8e-07, "input_cost_per_token": 4e-06, "input_cost_per_token_above_272k_tokens": 8e-06, "input_cost_per_token_above_272k_tokens_flex": 4e-06, + "input_cost_per_token_above_272k_tokens_priority": 1.6e-05, "input_cost_per_token_batches": 2e-06, "input_cost_per_token_flex": 2e-06, "input_cost_per_token_priority": 8e-06, @@ -29119,6 +29132,7 @@ "output_cost_per_token": 2e-05, "output_cost_per_token_above_272k_tokens": 3e-05, "output_cost_per_token_above_272k_tokens_flex": 1.5e-05, + "output_cost_per_token_above_272k_tokens_priority": 6e-05, "output_cost_per_token_batches": 1e-05, "output_cost_per_token_flex": 1e-05, "output_cost_per_token_priority": 4e-05, @@ -29161,16 +29175,19 @@ "cache_creation_input_token_cost": 5e-06, "cache_creation_input_token_cost_above_272k_tokens": 1e-05, "cache_creation_input_token_cost_above_272k_tokens_flex": 5e-06, + "cache_creation_input_token_cost_above_272k_tokens_priority": 2e-05, "cache_creation_input_token_cost_flex": 2.5e-06, "cache_creation_input_token_cost_priority": 1e-05, "cache_read_input_token_cost": 4e-07, "cache_read_input_token_cost_above_272k_tokens": 8e-07, "cache_read_input_token_cost_above_272k_tokens_flex": 4e-07, + "cache_read_input_token_cost_above_272k_tokens_priority": 1.6e-06, "cache_read_input_token_cost_flex": 2e-07, "cache_read_input_token_cost_priority": 8e-07, "input_cost_per_token": 4e-06, "input_cost_per_token_above_272k_tokens": 8e-06, "input_cost_per_token_above_272k_tokens_flex": 4e-06, + "input_cost_per_token_above_272k_tokens_priority": 1.6e-05, "input_cost_per_token_batches": 2e-06, "input_cost_per_token_flex": 2e-06, "input_cost_per_token_priority": 8e-06, @@ -29182,6 +29199,7 @@ "output_cost_per_token": 2e-05, "output_cost_per_token_above_272k_tokens": 3e-05, "output_cost_per_token_above_272k_tokens_flex": 1.5e-05, + "output_cost_per_token_above_272k_tokens_priority": 6e-05, "output_cost_per_token_batches": 1e-05, "output_cost_per_token_flex": 1e-05, "output_cost_per_token_priority": 4e-05, @@ -29225,16 +29243,19 @@ "cache_creation_input_token_cost": 2.5e-06, "cache_creation_input_token_cost_above_272k_tokens": 5e-06, "cache_creation_input_token_cost_above_272k_tokens_flex": 2.5e-06, + "cache_creation_input_token_cost_above_272k_tokens_priority": 1e-05, "cache_creation_input_token_cost_flex": 1.25e-06, "cache_creation_input_token_cost_priority": 5e-06, "cache_read_input_token_cost": 2e-07, "cache_read_input_token_cost_above_272k_tokens": 4e-07, "cache_read_input_token_cost_above_272k_tokens_flex": 2e-07, + "cache_read_input_token_cost_above_272k_tokens_priority": 8e-07, "cache_read_input_token_cost_flex": 1e-07, "cache_read_input_token_cost_priority": 4e-07, "input_cost_per_token": 2e-06, "input_cost_per_token_above_272k_tokens": 4e-06, "input_cost_per_token_above_272k_tokens_flex": 2e-06, + "input_cost_per_token_above_272k_tokens_priority": 8e-06, "input_cost_per_token_batches": 1e-06, "input_cost_per_token_flex": 1e-06, "input_cost_per_token_priority": 4e-06, @@ -29246,6 +29267,7 @@ "output_cost_per_token": 1.2e-05, "output_cost_per_token_above_272k_tokens": 1.8e-05, "output_cost_per_token_above_272k_tokens_flex": 9e-06, + "output_cost_per_token_above_272k_tokens_priority": 3.6e-05, "output_cost_per_token_batches": 6e-06, "output_cost_per_token_flex": 6e-06, "output_cost_per_token_priority": 2.4e-05, @@ -29288,16 +29310,19 @@ "cache_creation_input_token_cost": 2.5e-07, "cache_creation_input_token_cost_above_272k_tokens": 5e-07, "cache_creation_input_token_cost_above_272k_tokens_flex": 2.5e-07, + "cache_creation_input_token_cost_above_272k_tokens_priority": 1e-06, "cache_creation_input_token_cost_flex": 1.25e-07, "cache_creation_input_token_cost_priority": 5e-07, "cache_read_input_token_cost": 2e-08, "cache_read_input_token_cost_above_272k_tokens": 4e-08, "cache_read_input_token_cost_above_272k_tokens_flex": 2e-08, + "cache_read_input_token_cost_above_272k_tokens_priority": 8e-08, "cache_read_input_token_cost_flex": 1e-08, "cache_read_input_token_cost_priority": 4e-08, "input_cost_per_token": 2e-07, "input_cost_per_token_above_272k_tokens": 4e-07, "input_cost_per_token_above_272k_tokens_flex": 2e-07, + "input_cost_per_token_above_272k_tokens_priority": 8e-07, "input_cost_per_token_batches": 1e-07, "input_cost_per_token_flex": 1e-07, "input_cost_per_token_priority": 4e-07, @@ -29309,6 +29334,7 @@ "output_cost_per_token": 1.2e-06, "output_cost_per_token_above_272k_tokens": 1.8e-06, "output_cost_per_token_above_272k_tokens_flex": 9e-07, + "output_cost_per_token_above_272k_tokens_priority": 3.6e-06, "output_cost_per_token_batches": 6e-07, "output_cost_per_token_flex": 6e-07, "output_cost_per_token_priority": 2.4e-06, @@ -29548,7 +29574,10 @@ "supports_web_search": true, "supports_none_reasoning_effort": true, "supports_xhigh_reasoning_effort": true, - "supports_minimal_reasoning_effort": false + "supports_minimal_reasoning_effort": false, + "input_cost_per_token_above_272k_tokens_flex": 5e-06, + "output_cost_per_token_above_272k_tokens_flex": 2.25e-05, + "cache_read_input_token_cost_above_272k_tokens_flex": 5e-07 }, "gpt-5.5-2026-04-23": { "cache_read_input_token_cost": 5e-07, @@ -29602,7 +29631,10 @@ "supports_web_search": true, "supports_none_reasoning_effort": true, "supports_xhigh_reasoning_effort": true, - "supports_minimal_reasoning_effort": false + "supports_minimal_reasoning_effort": false, + "input_cost_per_token_above_272k_tokens_flex": 5e-06, + "output_cost_per_token_above_272k_tokens_flex": 2.25e-05, + "cache_read_input_token_cost_above_272k_tokens_flex": 5e-07 }, "gpt-5.5-pro": { "cache_read_input_token_cost": 3e-06, @@ -29751,7 +29783,10 @@ "supports_none_reasoning_effort": true, "default_reasoning_effort": "none", "supports_xhigh_reasoning_effort": true, - "supports_minimal_reasoning_effort": true + "supports_minimal_reasoning_effort": true, + "input_cost_per_token_above_272k_tokens_flex": 2.5e-06, + "output_cost_per_token_above_272k_tokens_flex": 1.125e-05, + "cache_read_input_token_cost_above_272k_tokens_flex": 2.5e-07 }, "gpt-5.4-2026-03-05": { "cache_read_input_token_cost": 2.5e-07, @@ -29800,7 +29835,10 @@ "supports_none_reasoning_effort": true, "default_reasoning_effort": "none", "supports_xhigh_reasoning_effort": true, - "supports_minimal_reasoning_effort": true + "supports_minimal_reasoning_effort": true, + "input_cost_per_token_above_272k_tokens_flex": 2.5e-06, + "output_cost_per_token_above_272k_tokens_flex": 1.125e-05, + "cache_read_input_token_cost_above_272k_tokens_flex": 2.5e-07 }, "gpt-5.4-pro": { "cache_read_input_token_cost": 3e-06, @@ -29849,7 +29887,9 @@ "supports_web_search": true, "supports_none_reasoning_effort": false, "supports_xhigh_reasoning_effort": true, - "supports_minimal_reasoning_effort": true + "supports_minimal_reasoning_effort": true, + "input_cost_per_token_above_272k_tokens_flex": 3e-05, + "output_cost_per_token_above_272k_tokens_flex": 0.000135 }, "gpt-5.4-pro-2026-03-05": { "cache_read_input_token_cost": 3e-06, @@ -29898,7 +29938,9 @@ "supports_web_search": true, "supports_none_reasoning_effort": false, "supports_xhigh_reasoning_effort": true, - "supports_minimal_reasoning_effort": true + "supports_minimal_reasoning_effort": true, + "input_cost_per_token_above_272k_tokens_flex": 3e-05, + "output_cost_per_token_above_272k_tokens_flex": 0.000135 }, "gpt-5.4-mini": { "cache_read_input_token_cost": 7.5e-08, @@ -41350,13 +41392,13 @@ "source": "https://docs.together.ai/docs/serverless-models" }, "together_ai/Qwen/Qwen3.8-2.4T-A95B": { - "cache_read_input_token_cost": 5e-07, - "input_cost_per_token": 2.5e-06, + "cache_read_input_token_cost": 2.5e-07, + "input_cost_per_token": 2e-06, "litellm_provider": "together_ai", "max_input_tokens": 1010000, "max_tokens": 1010000, "mode": "chat", - "output_cost_per_token": 6.25e-06, + "output_cost_per_token": 6e-06, "source": "https://docs.together.ai/docs/serverless-models", "supports_prompt_caching": true }, @@ -57558,5 +57600,526 @@ "supported_endpoints": [ "/v1/audio/transcriptions" ] + }, + "scaleway/glm-5.2": { + "input_cost_per_token": 1.8e-06, + "litellm_provider": "scaleway", + "max_input_tokens": 256000, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "output_cost_per_token": 5.5e-06, + "source": "https://www.scaleway.com/en/pricing/model-as-a-service/", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_vision": false + }, + "scaleway/deepseek-v4-flash-0731": { + "cache_read_input_token_cost": 8e-08, + "input_cost_per_token": 4e-07, + "litellm_provider": "scaleway", + "max_input_tokens": 256000, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 8e-07, + "source": "https://www.scaleway.com/en/pricing/model-as-a-service/", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_vision": false + }, + "azure_ai/kimi-k2.7-code": { + "cache_read_input_token_cost": 1.9e-07, + "deprecation_date": "2026-10-03", + "input_cost_per_token": 9.5e-07, + "litellm_provider": "azure_ai", + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 4e-06, + "source": "https://azure.microsoft.com/en-us/pricing/details/ai-foundry-models/kimi/", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "bedrock/us-gov-west-1/nvidia.nemotron-nano-3-30b": { + "input_cost_per_token": 7.2e-08, + "litellm_provider": "bedrock", + "max_input_tokens": 262144, + "max_output_tokens": 8192, + "max_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 2.88e-07, + "source": "https://aws.amazon.com/bedrock/pricing/", + "supports_function_calling": true, + "supports_native_structured_output": true, + "supports_system_messages": true, + "supports_tool_choice": true + }, + "bedrock/us-gov-west-1/nvidia.nemotron-nano-12b-v2": { + "input_cost_per_token": 2.4e-07, + "litellm_provider": "bedrock", + "max_input_tokens": 128000, + "max_output_tokens": 8192, + "max_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 7.2e-07, + "supports_system_messages": true, + "supports_vision": true + }, + "bedrock/us-gov-west-1/nvidia.nemotron-super-3-120b": { + "input_cost_per_token": 1.8e-07, + "litellm_provider": "bedrock", + "max_input_tokens": 256000, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 7.8e-07, + "source": "https://aws.amazon.com/bedrock/pricing/", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_system_messages": true, + "supports_tool_choice": true + }, + "bedrock/us-gov-west-1/openai.gpt-oss-20b-1:0": { + "input_cost_per_token": 8.4e-08, + "litellm_provider": "bedrock", + "max_input_tokens": 128000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 3.6e-07, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "bedrock/us-gov-west-1/openai.gpt-oss-120b-1:0": { + "input_cost_per_token": 1.8e-07, + "litellm_provider": "bedrock", + "max_input_tokens": 128000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 7.2e-07, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "bedrock/us-gov-west-1/anthropic.claude-sonnet-5": { + "bedrock_converse_supports_strict_tools": false, + "bedrock_output_config_effort_ceiling": "xhigh", + "cache_creation_input_token_cost": 3e-06, + "cache_creation_input_token_cost_above_1hr": 4.8e-06, + "cache_read_input_token_cost": 2.4e-07, + "input_cost_per_token": 2.4e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 1.2e-05, + "prompt_cache_min_tokens": 1024, + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_mid_conversation_system": true, + "supports_native_structured_output": false, + "supports_output_config": true, + "supports_parallel_tool_use_config": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true + }, + "bedrock/us-gov-west-1/anthropic.claude-opus-4-8": { + "bedrock_converse_supports_strict_tools": false, + "bedrock_output_config_effort_ceiling": "xhigh", + "cache_creation_input_token_cost": 7.5e-06, + "cache_creation_input_token_cost_above_1hr": 1.2e-05, + "cache_read_input_token_cost": 6e-07, + "input_cost_per_token": 6e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 3e-05, + "prompt_cache_min_tokens": 1024, + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_mid_conversation_system": true, + "supports_native_structured_output": true, + "supports_output_config": true, + "supports_parallel_tool_use_config": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true + }, + "bedrock/us-gov-east-1/nvidia.nemotron-nano-3-30b": { + "input_cost_per_token": 7.2e-08, + "litellm_provider": "bedrock", + "max_input_tokens": 262144, + "max_output_tokens": 8192, + "max_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 2.88e-07, + "source": "https://aws.amazon.com/bedrock/pricing/", + "supports_function_calling": true, + "supports_native_structured_output": true, + "supports_system_messages": true, + "supports_tool_choice": true + }, + "bedrock/us-gov-east-1/nvidia.nemotron-nano-12b-v2": { + "input_cost_per_token": 2.4e-07, + "litellm_provider": "bedrock", + "max_input_tokens": 128000, + "max_output_tokens": 8192, + "max_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 7.2e-07, + "supports_system_messages": true, + "supports_vision": true + }, + "bedrock/us-gov-east-1/nvidia.nemotron-super-3-120b": { + "input_cost_per_token": 1.8e-07, + "litellm_provider": "bedrock", + "max_input_tokens": 256000, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 7.8e-07, + "source": "https://aws.amazon.com/bedrock/pricing/", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_system_messages": true, + "supports_tool_choice": true + }, + "bedrock/us-gov-east-1/openai.gpt-oss-20b-1:0": { + "input_cost_per_token": 8.4e-08, + "litellm_provider": "bedrock", + "max_input_tokens": 128000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 3.6e-07, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "bedrock/us-gov-east-1/openai.gpt-oss-120b-1:0": { + "input_cost_per_token": 1.8e-07, + "litellm_provider": "bedrock", + "max_input_tokens": 128000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 7.2e-07, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "bedrock/us-gov-east-1/anthropic.claude-sonnet-5": { + "bedrock_converse_supports_strict_tools": false, + "bedrock_output_config_effort_ceiling": "xhigh", + "cache_creation_input_token_cost": 3e-06, + "cache_creation_input_token_cost_above_1hr": 4.8e-06, + "cache_read_input_token_cost": 2.4e-07, + "input_cost_per_token": 2.4e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 1.2e-05, + "prompt_cache_min_tokens": 1024, + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_mid_conversation_system": true, + "supports_native_structured_output": false, + "supports_output_config": true, + "supports_parallel_tool_use_config": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true + }, + "bedrock/us-gov-east-1/anthropic.claude-opus-4-8": { + "bedrock_converse_supports_strict_tools": false, + "bedrock_output_config_effort_ceiling": "xhigh", + "cache_creation_input_token_cost": 7.5e-06, + "cache_creation_input_token_cost_above_1hr": 1.2e-05, + "cache_read_input_token_cost": 6e-07, + "input_cost_per_token": 6e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 3e-05, + "prompt_cache_min_tokens": 1024, + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_mid_conversation_system": true, + "supports_native_structured_output": true, + "supports_output_config": true, + "supports_parallel_tool_use_config": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true + }, + "bedrock_mantle/us-gov-west-1/openai.gpt-5.6-terra": { + "litellm_provider": "bedrock_mantle", + "max_input_tokens": 1050000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "responses", + "use_openai_responses_path": true, + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true, + "input_cost_per_token": 2.64e-06, + "input_cost_per_token_above_272k_tokens": 5.28e-06, + "cache_creation_input_token_cost": 3.3e-06, + "cache_creation_input_token_cost_above_272k_tokens": 6.6e-06, + "cache_read_input_token_cost": 2.64e-07, + "cache_read_input_token_cost_above_272k_tokens": 5.28e-07, + "output_cost_per_token": 1.584e-05, + "output_cost_per_token_above_272k_tokens": 2.376e-05 + }, + "bedrock_mantle/us-gov-west-1/openai.gpt-5.6-luna": { + "litellm_provider": "bedrock_mantle", + "max_input_tokens": 1050000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "responses", + "use_openai_responses_path": true, + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true, + "input_cost_per_token": 2.64e-07, + "input_cost_per_token_above_272k_tokens": 5.28e-07, + "cache_creation_input_token_cost": 3.3e-07, + "cache_creation_input_token_cost_above_272k_tokens": 6.6e-07, + "cache_read_input_token_cost": 2.64e-08, + "cache_read_input_token_cost_above_272k_tokens": 5.28e-08, + "output_cost_per_token": 1.584e-06, + "output_cost_per_token_above_272k_tokens": 2.376e-06 + }, + "bedrock_mantle/us-gov-west-1/openai.gpt-5.4": { + "litellm_provider": "bedrock_mantle", + "max_input_tokens": 1050000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "responses", + "use_openai_responses_path": true, + "supported_endpoints": [ + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true, + "input_cost_per_token": 3.3e-06, + "cache_read_input_token_cost": 3.3e-07, + "output_cost_per_token": 1.98e-05 + }, + "bedrock_mantle/us-gov-west-1/xai.grok-4.3": { + "use_openai_responses_path": true, + "litellm_provider": "bedrock_mantle", + "max_input_tokens": 131072, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true, + "source": "https://aws.amazon.com/bedrock/pricing/", + "input_cost_per_token": 1.5e-06, + "output_cost_per_token": 3e-06, + "cache_read_input_token_cost": 2.4e-07 + }, + "bedrock_mantle/us-gov-east-1/openai.gpt-5.4": { + "litellm_provider": "bedrock_mantle", + "max_input_tokens": 1050000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "responses", + "use_openai_responses_path": true, + "supported_endpoints": [ + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true, + "input_cost_per_token": 3.3e-06, + "cache_read_input_token_cost": 3.3e-07, + "output_cost_per_token": 1.98e-05 + }, + "azure/us-gov/gpt-5.1": { + "cache_read_input_token_cost": 1.71875e-07, + "default_reasoning_effort": "none", + "deprecation_date": "2027-05-15", + "input_cost_per_token": 1.71875e-06, + "litellm_provider": "azure", + "max_input_tokens": 272000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 1.375e-05, + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/batch", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text", + "image" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_none_reasoning_effort": true, + "supports_parallel_function_calling": 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 + }, + "azure/us-gov/o3-mini": { + "cache_read_input_token_cost": 7.57e-07, + "deprecation_date": "2026-10-01", + "input_cost_per_token": 1.513e-06, + "litellm_provider": "azure", + "max_input_tokens": 200000, + "max_output_tokens": 100000, + "max_tokens": 100000, + "mode": "chat", + "output_cost_per_token": 6.05e-06, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": false + }, + "azure/us-gov/text-embedding-3-large": { + "deprecation_date": "2028-02-09", + "input_cost_per_token": 1.63e-07, + "litellm_provider": "azure", + "max_input_tokens": 8191, + "max_tokens": 8191, + "mode": "embedding", + "output_cost_per_token": 0.0 + }, + "azure/us-gov/text-embedding-3-small": { + "deprecation_date": "2028-02-09", + "input_cost_per_token": 2.5e-08, + "litellm_provider": "azure", + "max_input_tokens": 8191, + "max_tokens": 8191, + "mode": "embedding", + "output_cost_per_token": 0.0 } } 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 0e1c832ebf5..5c8de19a7e9 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 @@ -1522,7 +1522,7 @@ def test_gpt_5_6_alias_prices_match_sol(local_model_cost_map): sol = litellm.model_cost["gpt-5.6-sol"] cost_fields = sorted(field for field in sol if "cost" in field) - assert len(cost_fields) == 23 + assert len(cost_fields) == 27 for field in cost_fields: assert alias.get(field) == sol.get(field), field @@ -4039,8 +4039,8 @@ def test_fast_service_tier_matches_priority_above_the_context_threshold(_local_m ) assert fast == priority - assert fast[0] == pytest.approx(300_000 * 8e-06, rel=1e-9) - assert fast[1] == pytest.approx(1_000 * 3e-05, rel=1e-9) + assert fast[0] == pytest.approx(300_000 * 1.6e-05, rel=1e-9) + assert fast[1] == pytest.approx(1_000 * 6e-05, rel=1e-9) def test_priority_reasoning_tokens_bill_at_the_priority_output_rate(_local_model_cost_map): diff --git a/tests/test_litellm/test_openai_service_tier_long_context_pricing.py b/tests/test_litellm/test_openai_service_tier_long_context_pricing.py new file mode 100644 index 00000000000..c0860a5b55f --- /dev/null +++ b/tests/test_litellm/test_openai_service_tier_long_context_pricing.py @@ -0,0 +1,156 @@ +import json +from functools import lru_cache +from pathlib import Path + +import pytest + +import litellm + +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" + +FLEX_LONG_CONTEXT = { + "gpt-5.4": { + "input_cost_per_token_above_272k_tokens_flex": 2.5e-06, + "output_cost_per_token_above_272k_tokens_flex": 1.125e-05, + "cache_read_input_token_cost_above_272k_tokens_flex": 2.5e-07, + }, + "gpt-5.4-pro": { + "input_cost_per_token_above_272k_tokens_flex": 3e-05, + "output_cost_per_token_above_272k_tokens_flex": 0.000135, + }, + "gpt-5.5": { + "input_cost_per_token_above_272k_tokens_flex": 5e-06, + "output_cost_per_token_above_272k_tokens_flex": 2.25e-05, + "cache_read_input_token_cost_above_272k_tokens_flex": 5e-07, + }, +} + +PRIORITY_LONG_CONTEXT = { + "gpt-5.6": { + "input_cost_per_token_above_272k_tokens_priority": 1.6e-05, + "output_cost_per_token_above_272k_tokens_priority": 6e-05, + "cache_read_input_token_cost_above_272k_tokens_priority": 1.6e-06, + "cache_creation_input_token_cost_above_272k_tokens_priority": 2e-05, + }, + "gpt-5.6-sol": { + "input_cost_per_token_above_272k_tokens_priority": 1.6e-05, + "output_cost_per_token_above_272k_tokens_priority": 6e-05, + "cache_read_input_token_cost_above_272k_tokens_priority": 1.6e-06, + "cache_creation_input_token_cost_above_272k_tokens_priority": 2e-05, + }, + "gpt-5.6-terra": { + "input_cost_per_token_above_272k_tokens_priority": 8e-06, + "output_cost_per_token_above_272k_tokens_priority": 3.6e-05, + "cache_read_input_token_cost_above_272k_tokens_priority": 8e-07, + "cache_creation_input_token_cost_above_272k_tokens_priority": 1e-05, + }, + "gpt-5.6-luna": { + "input_cost_per_token_above_272k_tokens_priority": 8e-07, + "output_cost_per_token_above_272k_tokens_priority": 3.6e-06, + "cache_read_input_token_cost_above_272k_tokens_priority": 8e-08, + "cache_creation_input_token_cost_above_272k_tokens_priority": 1e-06, + }, +} + +EXPECTED = {**FLEX_LONG_CONTEXT, **PRIORITY_LONG_CONTEXT} + +NO_PUBLISHED_PRIORITY_LONG_CONTEXT = ("gpt-5.4", "gpt-5.5") + + +@pytest.fixture(autouse=True) +def _local_model_cost_map(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) + + +@lru_cache(maxsize=2) +def _load(path: Path) -> dict[str, dict[str, object]]: + with open(path) as f: + return json.load(f) + + +@pytest.mark.parametrize("path", [MAIN_PATH, BACKUP_PATH], ids=["main", "backup"]) +@pytest.mark.parametrize("model", sorted(EXPECTED)) +def test_service_tier_long_context_rates_are_published(model: str, path: Path) -> None: + """Each tier must carry its own above-272K rates, in both price files.""" + info = _load(path).get(model) + assert info is not None, f"{model} not found in {path.name}" + for key, expected in EXPECTED[model].items(): + assert info.get(key) == pytest.approx(expected), f"{model}.{key} is {info.get(key)!r}, expected {expected!r}" + + +@pytest.mark.parametrize("model", sorted(EXPECTED)) +def test_tier_long_context_rate_is_half_or_double_the_standard(model: str) -> None: + """Flex is half the standard long-context rate; priority is double it.""" + info = _load(MAIN_PATH)[model] + tier = "flex" if model in FLEX_LONG_CONTEXT else "priority" + ratio = 0.5 if tier == "flex" else 2.0 + for base in ("input_cost_per_token", "output_cost_per_token"): + standard = info[f"{base}_above_272k_tokens"] + tiered = info[f"{base}_above_272k_tokens_{tier}"] + assert tiered == pytest.approx(standard * ratio), ( + f"{model}.{base}_above_272k_tokens_{tier} is {tiered!r}, " + f"expected {ratio}x the standard long-context rate {standard!r}" + ) + + +@pytest.mark.parametrize("model", NO_PUBLISHED_PRIORITY_LONG_CONTEXT) +def test_no_priority_long_context_rates_where_openai_publishes_none(model: str) -> None: + """Guard against back-filling a rate OpenAI does not publish.""" + info = _load(MAIN_PATH)[model] + assert "input_cost_per_token_above_272k_tokens_priority" not in info + + +LONG_CONTEXT_PROMPT_TOKENS = 300_000 +COMPLETION_TOKENS = 1_000 + +TIERED_COST_CASES = [ + ("gpt-5.4", "flex", 2.5e-06, 1.125e-05), + ("gpt-5.4-pro", "flex", 3e-05, 0.000135), + ("gpt-5.5", "flex", 5e-06, 2.25e-05), + ("gpt-5.6", "priority", 1.6e-05, 6e-05), + ("gpt-5.6-sol", "priority", 1.6e-05, 6e-05), + ("gpt-5.6-terra", "priority", 8e-06, 3.6e-05), + ("gpt-5.6-luna", "priority", 8e-07, 3.6e-06), +] + + +@pytest.mark.parametrize("model,tier,input_rate,output_rate", TIERED_COST_CASES) +def test_cost_per_token_bills_long_context_at_the_tier_rate( + model: str, tier: str, input_rate: float, output_rate: float +) -> None: + """A prompt over 272K on flex or priority must bill at that tier's long-context rate.""" + input_cost, output_cost = litellm.cost_per_token( + model=model, + prompt_tokens=LONG_CONTEXT_PROMPT_TOKENS, + completion_tokens=COMPLETION_TOKENS, + service_tier=tier, + ) + assert input_cost == pytest.approx(LONG_CONTEXT_PROMPT_TOKENS * input_rate) + assert output_cost == pytest.approx(COMPLETION_TOKENS * output_rate) + + +@pytest.mark.parametrize("model,tier,input_rate,output_rate", TIERED_COST_CASES) +def test_cost_per_token_tier_differs_from_the_standard_long_context_cost( + model: str, tier: str, input_rate: float, output_rate: float +) -> None: + """Flex halves the standard long-context bill and priority doubles it.""" + ratio = 0.5 if tier == "flex" else 2.0 + standard = sum( + litellm.cost_per_token( + model=model, + prompt_tokens=LONG_CONTEXT_PROMPT_TOKENS, + completion_tokens=COMPLETION_TOKENS, + ) + ) + tiered = sum( + litellm.cost_per_token( + model=model, + prompt_tokens=LONG_CONTEXT_PROMPT_TOKENS, + completion_tokens=COMPLETION_TOKENS, + service_tier=tier, + ) + ) + assert tiered == pytest.approx(standard * ratio) diff --git a/whitelisted_bedrock_models.txt b/whitelisted_bedrock_models.txt index 7e20081988d..8753d7c3c77 100644 --- a/whitelisted_bedrock_models.txt +++ b/whitelisted_bedrock_models.txt @@ -217,3 +217,17 @@ bedrock/us-east-1/zai.glm-5 bedrock/us-west-2/zai.glm-5 bedrock/us-gov-east-1/anthropic.claude-haiku-4-5-20251001-v1:0 bedrock/us-gov-west-1/anthropic.claude-haiku-4-5-20251001-v1:0 +bedrock/us-gov-west-1/nvidia.nemotron-nano-3-30b +bedrock/us-gov-west-1/nvidia.nemotron-nano-12b-v2 +bedrock/us-gov-west-1/nvidia.nemotron-super-3-120b +bedrock/us-gov-west-1/openai.gpt-oss-20b-1:0 +bedrock/us-gov-west-1/openai.gpt-oss-120b-1:0 +bedrock/us-gov-west-1/anthropic.claude-sonnet-5 +bedrock/us-gov-west-1/anthropic.claude-opus-4-8 +bedrock/us-gov-east-1/nvidia.nemotron-nano-3-30b +bedrock/us-gov-east-1/nvidia.nemotron-nano-12b-v2 +bedrock/us-gov-east-1/nvidia.nemotron-super-3-120b +bedrock/us-gov-east-1/openai.gpt-oss-20b-1:0 +bedrock/us-gov-east-1/openai.gpt-oss-120b-1:0 +bedrock/us-gov-east-1/anthropic.claude-sonnet-5 +bedrock/us-gov-east-1/anthropic.claude-opus-4-8 From 8588a2ea42f7fac2a19e37123ebac5a7327b182a Mon Sep 17 00:00:00 2001 From: Oliver Jensen Date: Wed, 2 Sep 2026 17:01:54 +0200 Subject: [PATCH 31/48] fix(docker): install saml extra in litellm-backend image (#39291) The monolithic images install the saml extra but the split backend image did not, so /sso/saml/* returned 501 on Helm split-image deployments. The gateway image is unchanged since /sso/ routes are backend-only. --- backend/Dockerfile | 2 ++ 1 file changed, 2 insertions(+) diff --git a/backend/Dockerfile b/backend/Dockerfile index aa01b9fba8b..622fedcd70d 100644 --- a/backend/Dockerfile +++ b/backend/Dockerfile @@ -46,6 +46,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \ --extra proxy-runtime \ --extra extra_proxy \ --extra semantic-router \ + --extra saml \ --python python3.13 # Stage 2 — copy source and install the project + workspace members. @@ -57,6 +58,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \ --extra proxy-runtime \ --extra extra_proxy \ --extra semantic-router \ + --extra saml \ --python python3.13 RUN HOME=/opt/prisma XDG_CACHE_HOME=/opt/prisma/.cache PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries \ From 6b83b16559e5ceb4904121bcb90623a5f9f7115c Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 2 Sep 2026 08:04:14 -0700 Subject: [PATCH 32/48] feat(gemini): day-0 pricing for gemini-3.8-flash Gemini 3.8 Flash launches today with the same promotional pricing, limits, and thinking settings as Gemini 3.7 Flash, so the gemini/, vertex_ai/, and bare cost map entries mirror the 3.7 Flash ones. Regression tests lock the launch prices, the 4096-token cache minimum, and the gemini-3 thought signature gate in for the new model. --- ...odel_prices_and_context_window_backup.json | 173 ++++++++++++++++++ model_prices_and_context_window.json | 173 ++++++++++++++++++ .../llm_cost_calc/test_llm_cost_calc_utils.py | 45 +++++ .../test_vertex_ai_gemini_transformation.py | 3 + tests/test_litellm/test_utils.py | 1 + 5 files changed, 395 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index a3cfb300ea6..cc828e126ad 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -23514,6 +23514,63 @@ "web_search_billing_unit": "per_query", "google_maps_grounding_cost_per_query": 0.014 }, + "vertex_ai/gemini-3.8-flash": { + "prompt_cache_min_tokens": 4096, + "cache_read_input_token_cost": 7.5e-08, + "cache_read_input_token_cost_flex": 3.75e-08, + "input_cost_per_token": 7.5e-07, + "input_cost_per_token_batches": 3.75e-07, + "input_cost_per_token_flex": 3.75e-07, + "litellm_provider": "vertex_ai", + "max_input_tokens": 1048576, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "output_cost_per_reasoning_token": 3.75e-06, + "output_cost_per_token": 3.75e-06, + "output_cost_per_token_batches": 1.875e-06, + "output_cost_per_token_flex": 1.875e-06, + "regional_endpoint_uplift_multiplier": 1.1, + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_input": true, + "supports_function_calling": true, + "supports_parallel_function_calling": 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_url_context": true, + "supports_video_input": true, + "supports_vision": true, + "supports_web_search": true, + "supports_native_streaming": true, + "input_cost_per_token_priority": 1.35e-06, + "output_cost_per_token_priority": 6.75e-06, + "cache_read_input_token_cost_priority": 1.35e-07, + "search_context_cost_per_query": { + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014, + "search_context_size_high": 0.014 + }, + "web_search_billing_unit": "per_query", + "google_maps_grounding_cost_per_query": 0.014 + }, "vertex_ai/gemini-3.1-pro-preview": { "prompt_cache_min_tokens": 4096, "cache_read_input_token_cost": 2e-07, @@ -25351,6 +25408,65 @@ "web_search_billing_unit": "per_query", "google_maps_grounding_cost_per_query": 0.014 }, + "gemini/gemini-3.8-flash": { + "prompt_cache_min_tokens": 4096, + "cache_read_input_token_cost": 7.5e-08, + "cache_read_input_token_cost_flex": 3.75e-08, + "input_cost_per_token": 7.5e-07, + "input_cost_per_token_batches": 3.75e-07, + "input_cost_per_token_flex": 3.75e-07, + "litellm_provider": "gemini", + "max_input_tokens": 1048576, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "output_cost_per_reasoning_token": 3.75e-06, + "output_cost_per_token": 3.75e-06, + "output_cost_per_token_batches": 1.875e-06, + "output_cost_per_token_flex": 1.875e-06, + "rpm": 2000, + "source": "https://ai.google.dev/gemini-api/docs/pricing", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_output": false, + "supports_audio_input": true, + "supports_function_calling": true, + "supports_parallel_function_calling": 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_url_context": true, + "supports_video_input": true, + "supports_vision": true, + "supports_web_search": true, + "supports_native_streaming": true, + "tpm": 800000, + "input_cost_per_token_priority": 1.35e-06, + "output_cost_per_token_priority": 6.75e-06, + "cache_read_input_token_cost_priority": 1.35e-07, + "search_context_cost_per_query": { + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014, + "search_context_size_high": 0.014 + }, + "web_search_billing_unit": "per_query", + "google_maps_grounding_cost_per_query": 0.014 + }, "gemini/gemini-omni-flash-preview": { "input_cost_per_audio_token": 1.5e-06, "input_cost_per_token": 1.5e-06, @@ -25759,6 +25875,63 @@ "web_search_billing_unit": "per_query", "google_maps_grounding_cost_per_query": 0.014 }, + "gemini-3.8-flash": { + "prompt_cache_min_tokens": 4096, + "cache_read_input_token_cost": 7.5e-08, + "cache_read_input_token_cost_flex": 3.75e-08, + "input_cost_per_token": 7.5e-07, + "input_cost_per_token_batches": 3.75e-07, + "input_cost_per_token_flex": 3.75e-07, + "litellm_provider": "vertex_ai-language-models", + "max_input_tokens": 1048576, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "output_cost_per_reasoning_token": 3.75e-06, + "output_cost_per_token": 3.75e-06, + "output_cost_per_token_batches": 1.875e-06, + "output_cost_per_token_flex": 1.875e-06, + "source": "https://ai.google.dev/gemini-api/docs/pricing", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_output": false, + "supports_audio_input": true, + "supports_function_calling": true, + "supports_parallel_function_calling": 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_url_context": true, + "supports_video_input": true, + "supports_vision": true, + "supports_web_search": true, + "supports_native_streaming": true, + "input_cost_per_token_priority": 1.35e-06, + "output_cost_per_token_priority": 6.75e-06, + "cache_read_input_token_cost_priority": 1.35e-07, + "search_context_cost_per_query": { + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014, + "search_context_size_high": 0.014 + }, + "web_search_billing_unit": "per_query", + "google_maps_grounding_cost_per_query": 0.014 + }, "gemini/gemini-2.5-pro-preview-tts": { "cache_read_input_token_cost": 1.25e-07, "input_cost_per_audio_token": 7e-07, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index a3cfb300ea6..cc828e126ad 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -23514,6 +23514,63 @@ "web_search_billing_unit": "per_query", "google_maps_grounding_cost_per_query": 0.014 }, + "vertex_ai/gemini-3.8-flash": { + "prompt_cache_min_tokens": 4096, + "cache_read_input_token_cost": 7.5e-08, + "cache_read_input_token_cost_flex": 3.75e-08, + "input_cost_per_token": 7.5e-07, + "input_cost_per_token_batches": 3.75e-07, + "input_cost_per_token_flex": 3.75e-07, + "litellm_provider": "vertex_ai", + "max_input_tokens": 1048576, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "output_cost_per_reasoning_token": 3.75e-06, + "output_cost_per_token": 3.75e-06, + "output_cost_per_token_batches": 1.875e-06, + "output_cost_per_token_flex": 1.875e-06, + "regional_endpoint_uplift_multiplier": 1.1, + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_input": true, + "supports_function_calling": true, + "supports_parallel_function_calling": 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_url_context": true, + "supports_video_input": true, + "supports_vision": true, + "supports_web_search": true, + "supports_native_streaming": true, + "input_cost_per_token_priority": 1.35e-06, + "output_cost_per_token_priority": 6.75e-06, + "cache_read_input_token_cost_priority": 1.35e-07, + "search_context_cost_per_query": { + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014, + "search_context_size_high": 0.014 + }, + "web_search_billing_unit": "per_query", + "google_maps_grounding_cost_per_query": 0.014 + }, "vertex_ai/gemini-3.1-pro-preview": { "prompt_cache_min_tokens": 4096, "cache_read_input_token_cost": 2e-07, @@ -25351,6 +25408,65 @@ "web_search_billing_unit": "per_query", "google_maps_grounding_cost_per_query": 0.014 }, + "gemini/gemini-3.8-flash": { + "prompt_cache_min_tokens": 4096, + "cache_read_input_token_cost": 7.5e-08, + "cache_read_input_token_cost_flex": 3.75e-08, + "input_cost_per_token": 7.5e-07, + "input_cost_per_token_batches": 3.75e-07, + "input_cost_per_token_flex": 3.75e-07, + "litellm_provider": "gemini", + "max_input_tokens": 1048576, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "output_cost_per_reasoning_token": 3.75e-06, + "output_cost_per_token": 3.75e-06, + "output_cost_per_token_batches": 1.875e-06, + "output_cost_per_token_flex": 1.875e-06, + "rpm": 2000, + "source": "https://ai.google.dev/gemini-api/docs/pricing", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_output": false, + "supports_audio_input": true, + "supports_function_calling": true, + "supports_parallel_function_calling": 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_url_context": true, + "supports_video_input": true, + "supports_vision": true, + "supports_web_search": true, + "supports_native_streaming": true, + "tpm": 800000, + "input_cost_per_token_priority": 1.35e-06, + "output_cost_per_token_priority": 6.75e-06, + "cache_read_input_token_cost_priority": 1.35e-07, + "search_context_cost_per_query": { + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014, + "search_context_size_high": 0.014 + }, + "web_search_billing_unit": "per_query", + "google_maps_grounding_cost_per_query": 0.014 + }, "gemini/gemini-omni-flash-preview": { "input_cost_per_audio_token": 1.5e-06, "input_cost_per_token": 1.5e-06, @@ -25759,6 +25875,63 @@ "web_search_billing_unit": "per_query", "google_maps_grounding_cost_per_query": 0.014 }, + "gemini-3.8-flash": { + "prompt_cache_min_tokens": 4096, + "cache_read_input_token_cost": 7.5e-08, + "cache_read_input_token_cost_flex": 3.75e-08, + "input_cost_per_token": 7.5e-07, + "input_cost_per_token_batches": 3.75e-07, + "input_cost_per_token_flex": 3.75e-07, + "litellm_provider": "vertex_ai-language-models", + "max_input_tokens": 1048576, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "output_cost_per_reasoning_token": 3.75e-06, + "output_cost_per_token": 3.75e-06, + "output_cost_per_token_batches": 1.875e-06, + "output_cost_per_token_flex": 1.875e-06, + "source": "https://ai.google.dev/gemini-api/docs/pricing", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_output": false, + "supports_audio_input": true, + "supports_function_calling": true, + "supports_parallel_function_calling": 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_url_context": true, + "supports_video_input": true, + "supports_vision": true, + "supports_web_search": true, + "supports_native_streaming": true, + "input_cost_per_token_priority": 1.35e-06, + "output_cost_per_token_priority": 6.75e-06, + "cache_read_input_token_cost_priority": 1.35e-07, + "search_context_cost_per_query": { + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014, + "search_context_size_high": 0.014 + }, + "web_search_billing_unit": "per_query", + "google_maps_grounding_cost_per_query": 0.014 + }, "gemini/gemini-2.5-pro-preview-tts": { "cache_read_input_token_cost": 1.25e-07, "input_cost_per_audio_token": 7e-07, 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 0e1c832ebf5..0ccb05c67a3 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 @@ -4200,6 +4200,51 @@ def test_generic_cost_per_token_gemini_37_flash(_local_model_cost_map): assert completion_cost == pytest.approx(0.001875) +GEMINI_38_FLASH_LAUNCH_PRICING = [ + ("gemini-3.8-flash", 7.5e-07, 3.75e-06, 7.5e-08), + ("gemini/gemini-3.8-flash", 7.5e-07, 3.75e-06, 7.5e-08), + ("vertex_ai/gemini-3.8-flash", 7.5e-07, 3.75e-06, 7.5e-08), +] + + +@pytest.mark.parametrize("model,input_cost,output_cost,cache_read_cost", GEMINI_38_FLASH_LAUNCH_PRICING) +def test_gemini_38_flash_launch_pricing(model, input_cost, output_cost, cache_read_cost, _local_model_cost_map): + model_cost_map = litellm.model_cost[model] + assert model_cost_map["input_cost_per_token"] == input_cost + assert model_cost_map["output_cost_per_token"] == output_cost + assert model_cost_map["output_cost_per_reasoning_token"] == output_cost + assert model_cost_map["cache_read_input_token_cost"] == cache_read_cost + assert model_cost_map["mode"] == "chat" + assert model_cost_map["supports_reasoning"] is True + assert model_cost_map["supports_function_calling"] is True + assert model_cost_map["max_input_tokens"] == 1048576 + + +def test_gemini_38_flash_matches_37_flash_promotional_pricing(_local_model_cost_map): + for prefix in ("", "gemini/", "vertex_ai/"): + assert litellm.model_cost[f"{prefix}gemini-3.8-flash"] == litellm.model_cost[f"{prefix}gemini-3.7-flash"] + + +def test_generic_cost_per_token_gemini_38_flash(_local_model_cost_map): + usage = Usage( + prompt_tokens=1000, + completion_tokens=500, + total_tokens=1500, + completion_tokens_details=CompletionTokensDetailsWrapper( + reasoning_tokens=200, + text_tokens=300, + ), + prompt_tokens_details=PromptTokensDetailsWrapper(text_tokens=1000), + ) + prompt_cost, completion_cost = generic_cost_per_token( + model="gemini-3.8-flash", + usage=usage, + custom_llm_provider="gemini", + ) + assert prompt_cost == pytest.approx(0.00075) + assert completion_cost == pytest.approx(0.001875) + + def test_grok_46_launch_pricing(_local_model_cost_map): model_cost_map = litellm.model_cost["xai/grok-4.6"] assert model_cost_map["input_cost_per_token"] == 2e-06 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 8c1de12e7d9..4679b978f78 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 @@ -1096,10 +1096,13 @@ def test_natively_signed_parallel_turn_never_carries_a_placeholder(model): "gemini-3.5-flash", "gemini-3.6-flash", "gemini-3.7-flash", + "gemini-3.8-flash", "vertex_ai/gemini-3.5-flash", "vertex_ai/gemini-3.7-flash", + "vertex_ai/gemini-3.8-flash", "gemini/gemini-3.5-flash", "gemini/gemini-3.7-flash", + "gemini/gemini-3.8-flash", ], ) def test_placeholder_scoped_to_first_call_across_gemini_3_variants(model): diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index 521e91daded..0790b41c349 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -4655,6 +4655,7 @@ GEMINI_4096_CACHE_MIN_MODELS: Final = tuple( "gemini-3.5-flash", "gemini-3.6-flash", "gemini-3.7-flash", + "gemini-3.8-flash", "gemini-3.1-pro-preview", "gemini-3.1-pro-preview-customtools", ) From b76127774059d577229bbc9f74b3bf1b9fef812c Mon Sep 17 00:00:00 2001 From: mateo Date: Wed, 2 Sep 2026 15:09:49 +0000 Subject: [PATCH 33/48] fix(models): drop inherited retirement dates from azure/us-gov entries pending a Government schedule source Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/model_prices_and_context_window_backup.json | 4 ---- model_prices_and_context_window.json | 4 ---- 2 files changed, 8 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 87d348f4752..63b88c2a7b4 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -58055,7 +58055,6 @@ "azure/us-gov/gpt-5.1": { "cache_read_input_token_cost": 1.71875e-07, "default_reasoning_effort": "none", - "deprecation_date": "2027-05-15", "input_cost_per_token": 1.71875e-06, "litellm_provider": "azure", "max_input_tokens": 272000, @@ -58090,7 +58089,6 @@ }, "azure/us-gov/o3-mini": { "cache_read_input_token_cost": 7.57e-07, - "deprecation_date": "2026-10-01", "input_cost_per_token": 1.513e-06, "litellm_provider": "azure", "max_input_tokens": 200000, @@ -58105,7 +58103,6 @@ "supports_vision": false }, "azure/us-gov/text-embedding-3-large": { - "deprecation_date": "2028-02-09", "input_cost_per_token": 1.63e-07, "litellm_provider": "azure", "max_input_tokens": 8191, @@ -58114,7 +58111,6 @@ "output_cost_per_token": 0.0 }, "azure/us-gov/text-embedding-3-small": { - "deprecation_date": "2028-02-09", "input_cost_per_token": 2.5e-08, "litellm_provider": "azure", "max_input_tokens": 8191, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 87d348f4752..63b88c2a7b4 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -58055,7 +58055,6 @@ "azure/us-gov/gpt-5.1": { "cache_read_input_token_cost": 1.71875e-07, "default_reasoning_effort": "none", - "deprecation_date": "2027-05-15", "input_cost_per_token": 1.71875e-06, "litellm_provider": "azure", "max_input_tokens": 272000, @@ -58090,7 +58089,6 @@ }, "azure/us-gov/o3-mini": { "cache_read_input_token_cost": 7.57e-07, - "deprecation_date": "2026-10-01", "input_cost_per_token": 1.513e-06, "litellm_provider": "azure", "max_input_tokens": 200000, @@ -58105,7 +58103,6 @@ "supports_vision": false }, "azure/us-gov/text-embedding-3-large": { - "deprecation_date": "2028-02-09", "input_cost_per_token": 1.63e-07, "litellm_provider": "azure", "max_input_tokens": 8191, @@ -58114,7 +58111,6 @@ "output_cost_per_token": 0.0 }, "azure/us-gov/text-embedding-3-small": { - "deprecation_date": "2028-02-09", "input_cost_per_token": 2.5e-08, "litellm_provider": "azure", "max_input_tokens": 8191, From 07cf9dc46f5a4fd3b506f89cc77744f723eff190 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Wed, 2 Sep 2026 08:10:52 -0700 Subject: [PATCH 34/48] feat(models): add Azure DeepSeek V4 Flash 0731 --- .../model_prices_and_context_window_backup.json | 16 ++++++++++++++++ model_prices_and_context_window.json | 16 ++++++++++++++++ 2 files changed, 32 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index a3cfb300ea6..893ac49f5c7 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -10171,6 +10171,22 @@ "supports_reasoning": true, "supports_tool_choice": true }, + "azure_ai/DeepSeek-V4-Flash-0731": { + "cache_read_input_token_cost": 2.8e-08, + "deprecation_date": "2026-12-03", + "input_cost_per_token": 1.9e-07, + "litellm_provider": "azure_ai", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5.1e-07, + "source": "https://azure.microsoft.com/en-us/pricing/details/ai-foundry-models/deepseek/", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, "azure_ai/deepseek-v4-flash-0731": { "cache_read_input_token_cost": 2.8e-08, "deprecation_date": "2026-12-03", diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index a3cfb300ea6..893ac49f5c7 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -10171,6 +10171,22 @@ "supports_reasoning": true, "supports_tool_choice": true }, + "azure_ai/DeepSeek-V4-Flash-0731": { + "cache_read_input_token_cost": 2.8e-08, + "deprecation_date": "2026-12-03", + "input_cost_per_token": 1.9e-07, + "litellm_provider": "azure_ai", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5.1e-07, + "source": "https://azure.microsoft.com/en-us/pricing/details/ai-foundry-models/deepseek/", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, "azure_ai/deepseek-v4-flash-0731": { "cache_read_input_token_cost": 2.8e-08, "deprecation_date": "2026-12-03", From da23e0241dc82649ff56f2e73e6e156c4e098129 Mon Sep 17 00:00:00 2001 From: mateo Date: Wed, 2 Sep 2026 15:28:53 +0000 Subject: [PATCH 35/48] fix(models): add cloudflare whisper transcription pricing and pin govcloud pricing tests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...odel_prices_and_context_window_backup.json | 20 ++ model_prices_and_context_window.json | 20 ++ .../test_bedrock_usgov_pricing.py | 200 +++++++++++++++--- ...st_cloudflare_workers_ai_model_metadata.py | 16 ++ 4 files changed, 230 insertions(+), 26 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 63b88c2a7b4..30621a17df3 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -58117,5 +58117,25 @@ "max_tokens": 8191, "mode": "embedding", "output_cost_per_token": 0.0 + }, + "cloudflare/@cf/openai/whisper": { + "input_cost_per_second": 7.5e-06, + "litellm_provider": "cloudflare", + "mode": "audio_transcription", + "output_cost_per_second": 0.0, + "source": "https://developers.cloudflare.com/workers-ai/models/whisper/", + "supported_endpoints": [ + "/v1/audio/transcriptions" + ] + }, + "cloudflare/@cf/openai/whisper-large-v3-turbo": { + "input_cost_per_second": 8.5e-06, + "litellm_provider": "cloudflare", + "mode": "audio_transcription", + "output_cost_per_second": 0.0, + "source": "https://developers.cloudflare.com/workers-ai/models/whisper-large-v3-turbo/", + "supported_endpoints": [ + "/v1/audio/transcriptions" + ] } } diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 63b88c2a7b4..30621a17df3 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -58117,5 +58117,25 @@ "max_tokens": 8191, "mode": "embedding", "output_cost_per_token": 0.0 + }, + "cloudflare/@cf/openai/whisper": { + "input_cost_per_second": 7.5e-06, + "litellm_provider": "cloudflare", + "mode": "audio_transcription", + "output_cost_per_second": 0.0, + "source": "https://developers.cloudflare.com/workers-ai/models/whisper/", + "supported_endpoints": [ + "/v1/audio/transcriptions" + ] + }, + "cloudflare/@cf/openai/whisper-large-v3-turbo": { + "input_cost_per_second": 8.5e-06, + "litellm_provider": "cloudflare", + "mode": "audio_transcription", + "output_cost_per_second": 0.0, + "source": "https://developers.cloudflare.com/workers-ai/models/whisper-large-v3-turbo/", + "supported_endpoints": [ + "/v1/audio/transcriptions" + ] } } diff --git a/tests/test_litellm/test_bedrock_usgov_pricing.py b/tests/test_litellm/test_bedrock_usgov_pricing.py index 6b3312b5cc4..f9e8fd4c46c 100644 --- a/tests/test_litellm/test_bedrock_usgov_pricing.py +++ b/tests/test_litellm/test_bedrock_usgov_pricing.py @@ -26,9 +26,7 @@ import pytest @pytest.fixture(scope="module") def model_data(): - json_path = os.path.join( - os.path.dirname(__file__), "../../model_prices_and_context_window.json" - ) + json_path = os.path.join(os.path.dirname(__file__), "../../model_prices_and_context_window.json") with open(json_path) as f: return json.load(f) @@ -51,21 +49,14 @@ def test_usgov_sonnet_4_5_pricing(model_data, model_key): info = model_data[model_key] assert info["input_cost_per_token"] == 3.6e-06, ( - f"{model_key}: input_cost_per_token should be $3.60/MTok " - f"(got {info['input_cost_per_token']})" + f"{model_key}: input_cost_per_token should be $3.60/MTok (got {info['input_cost_per_token']})" ) - assert ( - info["output_cost_per_token"] == 1.8e-05 - ), f"{model_key}: output_cost_per_token should be $18.00/MTok" - assert ( - info["cache_creation_input_token_cost"] == 4.5e-06 - ), f"{model_key}: 5m cache write should be $4.50/MTok" - assert ( - info["cache_creation_input_token_cost_above_1hr"] == 7.2e-06 - ), f"{model_key}: 1h cache write should be $7.20/MTok" - assert ( - info["cache_read_input_token_cost"] == 3.6e-07 - ), f"{model_key}: cache read should be $0.36/MTok" + assert info["output_cost_per_token"] == 1.8e-05, f"{model_key}: output_cost_per_token should be $18.00/MTok" + assert info["cache_creation_input_token_cost"] == 4.5e-06, f"{model_key}: 5m cache write should be $4.50/MTok" + assert info["cache_creation_input_token_cost_above_1hr"] == 7.2e-06, ( + f"{model_key}: 1h cache write should be $7.20/MTok" + ) + assert info["cache_read_input_token_cost"] == 3.6e-07, f"{model_key}: cache read should be $0.36/MTok" def test_usgov_carries_20_percent_premium_over_global(model_data): @@ -84,9 +75,7 @@ def test_usgov_carries_20_percent_premium_over_global(model_data): "cache_read_input_token_cost", ): ratio = usgov_info[field] / global_info[field] - assert ( - abs(ratio - 1.2) < 1e-9 - ), f"{field}: us-gov / global ratio is {ratio}, expected 1.2" + assert abs(ratio - 1.2) < 1e-9, f"{field}: us-gov / global ratio is {ratio}, expected 1.2" # The us-gov.anthropic.* cross-region inference profile is the only us-gov @@ -112,9 +101,7 @@ def test_usgov_cross_region_above_200k_carries_gov_premium(model_data, field, ex """ info = model_data[USGOV_CROSS_REGION_KEY] assert field in info, f"{USGOV_CROSS_REGION_KEY}: missing field {field}" - assert ( - info[field] == expected - ), f"{USGOV_CROSS_REGION_KEY}: {field} should be {expected} (got {info[field]})" + assert info[field] == expected, f"{USGOV_CROSS_REGION_KEY}: {field} should be {expected} (got {info[field]})" def test_usgov_cross_region_above_200k_ratio_to_global(model_data): @@ -127,6 +114,167 @@ def test_usgov_cross_region_above_200k_ratio_to_global(model_data): usgov_info = model_data[USGOV_CROSS_REGION_KEY] for field in EXPECTED_USGOV_ABOVE_200K: ratio = usgov_info[field] / global_info[field] - assert ( - abs(ratio - 1.2) < 1e-9 - ), f"{field}: us-gov / global ratio is {ratio}, expected 1.2" + assert abs(ratio - 1.2) < 1e-9, f"{field}: us-gov / global ratio is {ratio}, expected 1.2" + + +CLAUDE_GOV_EXPECTED = { + "anthropic.claude-sonnet-5": { + "input_cost_per_token": 2.4e-06, + "output_cost_per_token": 1.2e-05, + "cache_creation_input_token_cost": 3e-06, + "cache_creation_input_token_cost_above_1hr": 4.8e-06, + "cache_read_input_token_cost": 2.4e-07, + }, + "anthropic.claude-opus-4-8": { + "input_cost_per_token": 6e-06, + "output_cost_per_token": 3e-05, + "cache_creation_input_token_cost": 7.5e-06, + "cache_creation_input_token_cost_above_1hr": 1.2e-05, + "cache_read_input_token_cost": 6e-07, + }, +} + + +@pytest.mark.parametrize("base_key", CLAUDE_GOV_EXPECTED) +@pytest.mark.parametrize("region", ["us-gov-east-1", "us-gov-west-1"]) +def test_usgov_claude_sonnet5_opus48_pricing(model_data, region, base_key): + """Sonnet 5 and Opus 4.8 gov entries must match the rates AWS publishes + for both GovCloud regions on the Bedrock pricing page (1.2x global). + """ + gov_key = f"bedrock/{region}/{base_key}" + assert gov_key in model_data, f"Missing model entry: {gov_key}" + info = model_data[gov_key] + for field, expected in CLAUDE_GOV_EXPECTED[base_key].items(): + assert info[field] == expected, f"{gov_key}: {field} should be {expected} (got {info[field]})" + ratio = info[field] / model_data[base_key][field] + assert abs(ratio - 1.2) < 1e-9, f"{gov_key}: {field} gov/global ratio is {ratio}, expected 1.2" + + +CONVERSE_GOV_EXPECTED = { + "nvidia.nemotron-nano-3-30b": (7.2e-08, 2.88e-07), + "nvidia.nemotron-nano-12b-v2": (2.4e-07, 7.2e-07), + "nvidia.nemotron-super-3-120b": (1.8e-07, 7.8e-07), + "openai.gpt-oss-20b-1:0": (8.4e-08, 3.6e-07), + "openai.gpt-oss-120b-1:0": (1.8e-07, 7.2e-07), +} + + +@pytest.mark.parametrize("base_key", CONVERSE_GOV_EXPECTED) +@pytest.mark.parametrize("region", ["us-gov-east-1", "us-gov-west-1"]) +def test_usgov_converse_model_pricing(model_data, region, base_key): + """Nemotron and gpt-oss gov entries must match the AWS Bedrock offer file, + which prices both GovCloud regions identically at 1.2x commercial. + """ + gov_key = f"bedrock/{region}/{base_key}" + assert gov_key in model_data, f"Missing model entry: {gov_key}" + info = model_data[gov_key] + expected_input, expected_output = CONVERSE_GOV_EXPECTED[base_key] + assert info["input_cost_per_token"] == expected_input + assert info["output_cost_per_token"] == expected_output + assert info["litellm_provider"] == "bedrock" + base = model_data[base_key] + assert abs(info["input_cost_per_token"] / base["input_cost_per_token"] - 1.2) < 1e-9 + assert abs(info["output_cost_per_token"] / base["output_cost_per_token"] - 1.2) < 1e-9 + + +def test_usgov_west_llama3_8b_output_price_fixed(model_data): + """The us-gov-west-1 llama3-8b entry carried the 70B output rate ($2.65/MTok); + the AWS Bedrock offer file prices output at $0.60/MTok. AWS lists the model + in us-gov-west-1 only, so there is no east entry to check. + """ + info = model_data["bedrock/us-gov-west-1/meta.llama3-8b-instruct-v1:0"] + assert info["input_cost_per_token"] == 3e-07 + assert info["output_cost_per_token"] == 6e-07 + + +MANTLE_GOV_TIERED_EXPECTED = { + "openai.gpt-5.6-luna": { + "input_cost_per_token": 2.64e-07, + "input_cost_per_token_above_272k_tokens": 5.28e-07, + "cache_creation_input_token_cost": 3.3e-07, + "cache_creation_input_token_cost_above_272k_tokens": 6.6e-07, + "cache_read_input_token_cost": 2.64e-08, + "cache_read_input_token_cost_above_272k_tokens": 5.28e-08, + "output_cost_per_token": 1.584e-06, + "output_cost_per_token_above_272k_tokens": 2.376e-06, + }, + "openai.gpt-5.6-terra": { + "input_cost_per_token": 2.64e-06, + "input_cost_per_token_above_272k_tokens": 5.28e-06, + "cache_creation_input_token_cost": 3.3e-06, + "cache_creation_input_token_cost_above_272k_tokens": 6.6e-06, + "cache_read_input_token_cost": 2.64e-07, + "cache_read_input_token_cost_above_272k_tokens": 5.28e-07, + "output_cost_per_token": 1.584e-05, + "output_cost_per_token_above_272k_tokens": 2.376e-05, + }, +} + + +@pytest.mark.parametrize("model", MANTLE_GOV_TIERED_EXPECTED) +def test_usgov_west_mantle_terra_luna_pricing(model_data, model): + """Terra and Luna carry 1.2x commercial across every tier in the + us-gov-west-1 offer file; the us-gov-east-1 offer file has no SKUs for them. + """ + gov_key = f"bedrock_mantle/us-gov-west-1/{model}" + assert gov_key in model_data, f"Missing model entry: {gov_key}" + info = model_data[gov_key] + for field, expected in MANTLE_GOV_TIERED_EXPECTED[model].items(): + assert info[field] == expected, f"{gov_key}: {field} should be {expected} (got {info[field]})" + assert info["litellm_provider"] == "bedrock_mantle" + assert f"bedrock_mantle/us-gov-east-1/{model}" not in model_data + + +@pytest.mark.parametrize("region", ["us-gov-east-1", "us-gov-west-1"]) +def test_usgov_mantle_gpt_5_4_pricing_has_no_long_context_tier(model_data, region): + """gpt-5.4 gov rates come from the offer file, which publishes only the + standard tier in GovCloud: no long-context SKUs exist there, unlike commercial. + """ + gov_key = f"bedrock_mantle/{region}/openai.gpt-5.4" + assert gov_key in model_data, f"Missing model entry: {gov_key}" + info = model_data[gov_key] + assert info["input_cost_per_token"] == 3.3e-06 + assert info["cache_read_input_token_cost"] == 3.3e-07 + assert info["output_cost_per_token"] == 1.98e-05 + assert not any(field.endswith("_above_272k_tokens") for field in info) + + +def test_usgov_mantle_grok_4_3_west_only(model_data): + """grok-4.3 is priced in the us-gov-west-1 offer file only; the east offer + file carries grok-4.6 instead. + """ + info = model_data["bedrock_mantle/us-gov-west-1/xai.grok-4.3"] + assert info["input_cost_per_token"] == 1.5e-06 + assert info["output_cost_per_token"] == 3e-06 + assert info["cache_read_input_token_cost"] == 2.4e-07 + assert "bedrock_mantle/us-gov-east-1/xai.grok-4.3" not in model_data + + +AZURE_GOV_EXPECTED = { + "azure/us-gov/gpt-5.1": { + "input_cost_per_token": 1.71875e-06, + "cache_read_input_token_cost": 1.71875e-07, + "output_cost_per_token": 1.375e-05, + }, + "azure/us-gov/o3-mini": { + "input_cost_per_token": 1.513e-06, + "cache_read_input_token_cost": 7.57e-07, + "output_cost_per_token": 6.05e-06, + }, + "azure/us-gov/text-embedding-3-large": {"input_cost_per_token": 1.63e-07}, + "azure/us-gov/text-embedding-3-small": {"input_cost_per_token": 2.5e-08}, +} + + +@pytest.mark.parametrize("gov_key", AZURE_GOV_EXPECTED) +def test_azure_usgov_pricing(model_data, gov_key): + """Azure Government meters from the Azure retail prices API + (usgovvirginia/usgovarizona, serviceName 'Foundry Models'). No Government + retirement schedule is published, so these entries carry no deprecation_date. + """ + assert gov_key in model_data, f"Missing model entry: {gov_key}" + info = model_data[gov_key] + for field, expected in AZURE_GOV_EXPECTED[gov_key].items(): + assert info[field] == expected, f"{gov_key}: {field} should be {expected} (got {info[field]})" + assert info["litellm_provider"] == "azure" + assert "deprecation_date" not in info diff --git a/tests/test_litellm/test_cloudflare_workers_ai_model_metadata.py b/tests/test_litellm/test_cloudflare_workers_ai_model_metadata.py index 9ca4515239a..e33bcfb8378 100644 --- a/tests/test_litellm/test_cloudflare_workers_ai_model_metadata.py +++ b/tests/test_litellm/test_cloudflare_workers_ai_model_metadata.py @@ -75,6 +75,22 @@ def test_additional_current_models_are_present(): assert entry["output_cost_per_token"] > 0 +@pytest.mark.parametrize( + "key, published_price_per_audio_minute", + [ + ("cloudflare/@cf/openai/whisper", 0.00045), + ("cloudflare/@cf/openai/whisper-large-v3-turbo", 0.00051), + ], +) +def test_whisper_transcription_pricing_is_stored_per_second(key, published_price_per_audio_minute): + entry = litellm.model_cost[key] + assert entry["litellm_provider"] == "cloudflare" + assert entry["mode"] == "audio_transcription" + assert entry["supported_endpoints"] == ["/v1/audio/transcriptions"] + assert entry["output_cost_per_second"] == 0.0 + assert entry["input_cost_per_second"] == pytest.approx(published_price_per_audio_minute / 60) + + def test_root_and_backup_have_identical_cloudflare_keys(): if not os.path.exists(ROOT_MAP): pytest.skip("root cost map only ships in source checkouts") From 2ce4e3f8a99e12efce9433640059d9fca7bfb448 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 2 Sep 2026 08:32:49 -0700 Subject: [PATCH 36/48] fix(guardrails): run apply_guardrail-only providers in logging_only mode (#39297) * fix(guardrails): run apply_guardrail-only providers in logging_only mode A CustomGuardrail that implements only apply_guardrail inherited the CustomLogger no-op async_logging_hook, so mode: logging_only never scanned anything and never recorded guardrail_information. CustomGuardrail.async_logging_hook now routes the logged request and response through the call type's guardrail translation on copies and appends the verdict to standard_logging_object.guardrail_information. Resolves LIT-4876 Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(guardrails): keep logging_only scan copies inside the error boundary and return a fresh logging payload Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(guardrails): cover embedding scan, native-hook bypass, and unmapped call type in logging_only Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: yassin Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/integrations/custom_guardrail.py | 65 ++++++ .../integrations/test_custom_guardrail.py | 199 ++++++++++++++++++ 2 files changed, 264 insertions(+) diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index e87ac9521ae..372c9bf6b91 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -1,4 +1,5 @@ import contextvars +import copy import hashlib import os import secrets @@ -39,6 +40,7 @@ except ImportError: if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation dc: Final = DualCache() @@ -852,6 +854,69 @@ class CustomGuardrail(CustomLogger): return result + async def async_logging_hook( + self, + kwargs: dict, # mutable-ok: CustomLogger.async_logging_hook contract + result: object, + call_type: str, + ) -> tuple[dict, object]: # mutable-ok: CustomLogger.async_logging_hook contract + """logging_only: run apply_guardrail on copies of the logged request/response and record the verdict.""" + from litellm.llms import get_guardrail_translation_mapping + + if not self.uses_apply_guardrail_interface() or self.use_native_lifecycle_hooks: + return kwargs, result + try: + translation: Final = get_guardrail_translation_mapping(CallTypes(call_type))() + except ValueError: + verbose_logger.debug( + "Guardrail %s: no guardrail translation for call_type=%s, skipping logging_only scan", + self.guardrail_name, + call_type, + ) + return kwargs, result + litellm_params: Final = kwargs.get("litellm_params") or {} + scratch_metadata: Final = { + key: value + for key, value in (litellm_params.get("metadata") or {}).items() + if key != "standard_logging_guardrail_information" + } + try: + await self._scan_logged_call(kwargs, result, translation, scratch_metadata) + except Exception as e: + verbose_logger.warning("Guardrail %s: logging_only scan raised: %s", self.guardrail_name, e) + recorded: Final = scratch_metadata.get("standard_logging_guardrail_information") + standard_logging_object: Final = kwargs.get("standard_logging_object") + if not recorded or not isinstance(standard_logging_object, dict): + return kwargs, result + entries: Final = recorded if isinstance(recorded, list) else [recorded] + existing: Final = standard_logging_object.get("guardrail_information") or [] + return { + **kwargs, + "standard_logging_object": {**standard_logging_object, "guardrail_information": [*existing, *entries]}, + }, result + + async def _scan_logged_call( + self, + kwargs: dict, # mutable-ok: CustomLogger.async_logging_hook contract + result: object, + translation: "BaseTranslation", + scratch_metadata: dict, # mutable-ok: apply_guardrail records its verdict into request metadata + ) -> None: + optional_params: Final = kwargs.get("optional_params") or {} + scratch_input: Final = copy.deepcopy(kwargs.get("messages") or kwargs.get("input")) + scratch_request: Final = { + "model": kwargs.get("model"), + "messages": scratch_input, + "input": scratch_input, + "tools": copy.deepcopy(optional_params.get("tools")), + "litellm_call_id": kwargs.get("litellm_call_id"), + "metadata": scratch_metadata, + } + await translation.process_input_messages(data=scratch_request, guardrail_to_apply=self) + await translation.process_output_response( + response=copy.deepcopy(result), guardrail_to_apply=self, request_data=scratch_request + ) + def supports_scan_only_tool_results(self) -> bool: """Whether this guardrail can scan tool-result content. diff --git a/tests/test_litellm/integrations/test_custom_guardrail.py b/tests/test_litellm/integrations/test_custom_guardrail.py index d978eb48c12..7d70b9a8862 100644 --- a/tests/test_litellm/integrations/test_custom_guardrail.py +++ b/tests/test_litellm/integrations/test_custom_guardrail.py @@ -2237,3 +2237,202 @@ class TestRecordsOwnGuardrailInformation: ) assert _guardrail_entries(request_data) == [] + + +class _ApplyOnlyObserver(CustomGuardrail): + """Overrides only apply_guardrail, like panw_prisma_airs; inherits async_logging_hook.""" + + def __init__(self, block: bool = False): + from litellm.types.guardrails import GuardrailEventHooks + + super().__init__(guardrail_name="apply-only-observer", event_hook=GuardrailEventHooks.logging_only) + self.block = block + self.calls: list = [] + + @log_guardrail_information + async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): + from fastapi import HTTPException + + self.calls.append((input_type, list(inputs.get("texts") or []))) + if self.block: + raise HTTPException(status_code=400, detail={"error": "flagged"}) + return GenericGuardrailAPIInputs(texts=["[MASKED]" for _ in inputs.get("texts") or []]) + + +def _logged_call(messages: list | str) -> tuple[dict, object]: + from litellm.types.utils import Choices, Message, ModelResponse + + response = ModelResponse(choices=[Choices(message=Message(role="assistant", content="general kenobi"))]) + kwargs = { + "model": "gpt-5.4-mini", + "messages": messages, + "litellm_call_id": "call-1", + "litellm_params": {"metadata": {"user_api_key_user_id": "u1"}}, + "optional_params": {}, + "standard_logging_object": {"guardrail_information": None}, + } + return kwargs, response + + +class TestLoggingOnlyApplyGuardrail: + """LIT-4876 regression: a guardrail in mode logging_only that implements only + apply_guardrail must still run against the logged request and response and + record guardrail_information, instead of inheriting the CustomLogger no-op.""" + + @pytest.mark.asyncio + async def test_runs_apply_guardrail_observe_only_and_records_verdict(self): + guardrail = _ApplyOnlyObserver() + messages = [{"role": "user", "content": "hello there"}] + kwargs, response = _logged_call(messages) + + out_kwargs, out_response = await guardrail.async_logging_hook(kwargs, response, CallTypes.acompletion.value) + + assert guardrail.calls == [("request", ["hello there"]), ("response", ["general kenobi"])] + assert out_kwargs["messages"] == [{"role": "user", "content": "hello there"}] + assert out_response.choices[0].message.content == "general kenobi" + entries = out_kwargs["standard_logging_object"]["guardrail_information"] + assert [e["guardrail_name"] for e in entries] == ["apply-only-observer", "apply-only-observer"] + assert {e["guardrail_mode"] for e in entries} == {"logging_only"} + assert {e["guardrail_status"] for e in entries} == {"success"} + assert "standard_logging_guardrail_information" not in kwargs["litellm_params"]["metadata"] + assert kwargs["standard_logging_object"] == {"guardrail_information": None} + + @pytest.mark.asyncio + async def test_appends_to_pre_call_verdicts_without_duplicating_them(self): + guardrail = _ApplyOnlyObserver() + kwargs, response = _logged_call([{"role": "user", "content": "hello there"}]) + pre_call_entry = {"guardrail_name": "pii-blocker", "guardrail_mode": "pre_call", "guardrail_status": "success"} + kwargs["litellm_params"]["metadata"]["standard_logging_guardrail_information"] = [pre_call_entry] + kwargs["standard_logging_object"]["guardrail_information"] = [pre_call_entry] + + out_kwargs, _ = await guardrail.async_logging_hook(kwargs, response, CallTypes.acompletion.value) + + entries = out_kwargs["standard_logging_object"]["guardrail_information"] + assert [e["guardrail_name"] for e in entries] == ["pii-blocker", "apply-only-observer", "apply-only-observer"] + assert kwargs["litellm_params"]["metadata"]["standard_logging_guardrail_information"] == [pre_call_entry] + + @pytest.mark.asyncio + async def test_request_copy_failure_is_swallowed(self): + import threading + + guardrail = _ApplyOnlyObserver() + kwargs, response = _logged_call([{"role": "user", "content": "hello there", "lock": threading.Lock()}]) + + out_kwargs, out_response = await guardrail.async_logging_hook(kwargs, response, CallTypes.acompletion.value) + + assert guardrail.calls == [] + assert out_kwargs is kwargs + assert out_response is response + + @pytest.mark.asyncio + async def test_block_verdict_is_recorded_without_raising(self): + guardrail = _ApplyOnlyObserver(block=True) + kwargs, response = _logged_call([{"role": "user", "content": "flagged content"}]) + + out_kwargs, _ = await guardrail.async_logging_hook(kwargs, response, CallTypes.acompletion.value) + + assert guardrail.calls == [("request", ["flagged content"])] + entries = out_kwargs["standard_logging_object"]["guardrail_information"] + assert [e["guardrail_status"] for e in entries] == ["guardrail_intervened"] + + @pytest.mark.asyncio + async def test_call_type_without_translation_is_skipped(self): + guardrail = _ApplyOnlyObserver() + kwargs, response = _logged_call([{"role": "user", "content": "hello there"}]) + + out_kwargs, _ = await guardrail.async_logging_hook(kwargs, response, CallTypes.amoderation.value) + + assert guardrail.calls == [] + assert out_kwargs["standard_logging_object"]["guardrail_information"] is None + + @pytest.mark.asyncio + async def test_aembedding_scans_logged_input(self): + from litellm.types.utils import EmbeddingResponse + + guardrail = _ApplyOnlyObserver() + kwargs, _ = _logged_call("hello there") + response = EmbeddingResponse(data=[{"embedding": [0.1], "index": 0, "object": "embedding"}]) + + out_kwargs, out_response = await guardrail.async_logging_hook(kwargs, response, CallTypes.aembedding.value) + + assert guardrail.calls == [("request", ["hello there"])] + assert out_kwargs["messages"] == "hello there" + assert out_response is response + entries = out_kwargs["standard_logging_object"]["guardrail_information"] + assert [e["guardrail_status"] for e in entries] == ["success"] + + @pytest.mark.asyncio + async def test_native_lifecycle_hook_guardrail_is_left_alone(self): + class _NativeHooks(_ApplyOnlyObserver): + use_native_lifecycle_hooks = True + + guardrail = _NativeHooks() + kwargs, response = _logged_call([{"role": "user", "content": "hello there"}]) + + out_kwargs, out_response = await guardrail.async_logging_hook(kwargs, response, CallTypes.acompletion.value) + + assert guardrail.calls == [] + assert out_kwargs is kwargs + assert out_response is response + + @pytest.mark.asyncio + async def test_aresponses_scans_logged_messages_when_input_is_cleared(self): + from litellm.types.llms.openai import ResponsesAPIResponse + + guardrail = _ApplyOnlyObserver() + kwargs, _ = _logged_call([{"role": "user", "content": "hello there"}]) + kwargs["input"] = None + response = ResponsesAPIResponse( + id="resp_1", + created_at=1, + model="gpt-5.4-mini", + object="response", + status="completed", + output=[ + { + "type": "message", + "id": "msg_1", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "general kenobi"}], + } + ], + ) + + out_kwargs, _ = await guardrail.async_logging_hook(kwargs, response, CallTypes.aresponses.value) + + assert guardrail.calls == [("request", ["hello there"]), ("response", ["general kenobi"])] + entries = out_kwargs["standard_logging_object"]["guardrail_information"] + assert [e["guardrail_status"] for e in entries] == ["success", "success"] + + @pytest.mark.asyncio + async def test_async_success_handler_records_verdict_in_standard_logging_object(self): + import datetime as dt + + from litellm.litellm_core_utils.litellm_logging import Logging + + guardrail = _ApplyOnlyObserver() + guardrail.default_on = True + messages = [{"role": "user", "content": "hello there"}] + _, response = _logged_call(messages) + logging_obj = Logging( + model="gpt-5.4-mini", + messages=messages, + stream=False, + call_type=CallTypes.acompletion.value, + start_time=dt.datetime.now(), + litellm_call_id="call-1", + function_id="fn-1", + dynamic_async_success_callbacks=[guardrail], + ) + logging_obj.update_environment_variables( + litellm_params={"metadata": {}}, optional_params={}, model="gpt-5.4-mini", custom_llm_provider="openai" + ) + + await logging_obj.async_success_handler( + result=response, start_time=dt.datetime.now(), end_time=dt.datetime.now() + ) + + assert guardrail.calls == [("request", ["hello there"]), ("response", ["general kenobi"])] + entries = logging_obj.model_call_details["standard_logging_object"]["guardrail_information"] + assert [e["guardrail_status"] for e in entries] == ["success", "success"] From 69cd1bada6249889a9412155a6086faea65bfe6c Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 2 Sep 2026 08:46:28 -0700 Subject: [PATCH 37/48] test(gemini): compare gemini-3.8-flash to 3.7 flash field by field --- .../llm_cost_calc/test_llm_cost_calc_utils.py | 41 +++++++++++++++++-- 1 file changed, 38 insertions(+), 3 deletions(-) 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 0ccb05c67a3..9b3e60764e3 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 @@ -4220,9 +4220,44 @@ def test_gemini_38_flash_launch_pricing(model, input_cost, output_cost, cache_re assert model_cost_map["max_input_tokens"] == 1048576 -def test_gemini_38_flash_matches_37_flash_promotional_pricing(_local_model_cost_map): - for prefix in ("", "gemini/", "vertex_ai/"): - assert litellm.model_cost[f"{prefix}gemini-3.8-flash"] == litellm.model_cost[f"{prefix}gemini-3.7-flash"] +GEMINI_38_FLASH_FIELDS_SHARED_WITH_37_FLASH = ( + "input_cost_per_token", + "output_cost_per_token", + "output_cost_per_reasoning_token", + "cache_read_input_token_cost", + "input_cost_per_token_batches", + "output_cost_per_token_batches", + "input_cost_per_token_flex", + "output_cost_per_token_flex", + "cache_read_input_token_cost_flex", + "input_cost_per_token_priority", + "output_cost_per_token_priority", + "cache_read_input_token_cost_priority", + "search_context_cost_per_query", + "google_maps_grounding_cost_per_query", + "prompt_cache_min_tokens", + "max_input_tokens", + "max_output_tokens", + "supports_reasoning", + "supports_function_calling", + "supports_prompt_caching", + "supports_vision", + "supports_pdf_input", + "supports_audio_input", + "supports_video_input", + "supports_response_schema", + "supports_tool_choice", + "supports_web_search", + "supports_url_context", +) + + +@pytest.mark.parametrize("prefix", ["", "gemini/", "vertex_ai/"]) +def test_gemini_38_flash_matches_37_flash_promotional_pricing(prefix, _local_model_cost_map): + new_model = litellm.model_cost[f"{prefix}gemini-3.8-flash"] + old_model = litellm.model_cost[f"{prefix}gemini-3.7-flash"] + for field in GEMINI_38_FLASH_FIELDS_SHARED_WITH_37_FLASH: + assert new_model[field] == old_model[field], field def test_generic_cost_per_token_gemini_38_flash(_local_model_cost_map): From de80e3afe448237c69a535517ca88c12d079ee6d Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Wed, 2 Sep 2026 10:19:03 -0700 Subject: [PATCH 38/48] fix(helm): scale the classic chart's HPA out at the documented 60 percent CPU (#35975) * fix(helm): scale the classic chart's HPA out at the documented 60 percent CPU The litellm-helm chart shipped targetCPUUtilizationPercentage: 80, which is unexamined helm create scaffold rather than a chosen number. It arrived packaged with the stock minReplicas: 1, maxReplicas: 100, a commented-out targetMemoryUtilizationPercentage: 80, and the boilerplate "such as Minikube" comment, the same provenance as the 128Mi resource example this file just corrected. 60 is the documented recommendation. The mechanism behind it is scale-up lag: the chart's own startupProbe is failureThreshold: 30 times periodSeconds: 10, so a replica can take up to 300 seconds to become ready, and a pod added at 80 percent utilization arrives minutes after saturation. The memory target stays commented out on purpose. The prisma query engine's resident memory is a high-water mark that ratchets to the pod's worst-ever write and is never returned, so a memory-target HPA reads the largest write a pod ever did rather than what it is doing now, and replicas ratchet up without scaling back in. hpa_tests.yaml carried its second suite after a YAML document separator, and helm-unittest loads only the first document per file, so that suite never ran; an assertion planted in it still passed. Fold it into the one live suite and add coverage pinning the rendered CPU target, the absence of a memory metric by default, and that overrides still take effect. Bump the chart to 1.1.2, since rendered output changes for anyone running with autoscaling enabled. * fix(helm): bump litellm-helm to 1.1.3 after rebase onto 1.1.2 Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- helm/litellm-helm/Chart.yaml | 2 +- helm/litellm-helm/tests/hpa_tests.yaml | 42 ++++++++++++++++++++++---- helm/litellm-helm/values.yaml | 11 ++++++- 3 files changed, 47 insertions(+), 8 deletions(-) diff --git a/helm/litellm-helm/Chart.yaml b/helm/litellm-helm/Chart.yaml index 3959d85edf3..a3cb388ffc6 100644 --- a/helm/litellm-helm/Chart.yaml +++ b/helm/litellm-helm/Chart.yaml @@ -18,7 +18,7 @@ type: application # This is the chart version. This version number should be incremented each time you make changes # to the chart and its templates, including the app version. # Versions are expected to follow Semantic Versioning (https://semver.org/) -version: 1.1.2 +version: 1.1.3 # This is the version number of the application being deployed. This version number should be # incremented each time you make changes to the application. Versions are not expected to diff --git a/helm/litellm-helm/tests/hpa_tests.yaml b/helm/litellm-helm/tests/hpa_tests.yaml index ec18c3591d3..cd062dd5971 100644 --- a/helm/litellm-helm/tests/hpa_tests.yaml +++ b/helm/litellm-helm/tests/hpa_tests.yaml @@ -1,4 +1,4 @@ -suite: "hpa with behavior" +suite: "hpa" templates: - hpa.yaml tests: @@ -23,14 +23,44 @@ tests: - equal: { path: spec.behavior.scaleUp.stabilizationWindowSeconds, value: 60 } - equal: { path: spec.behavior.scaleDown.stabilizationWindowSeconds, value: 90 } ---- -suite: "hpa without behavior" -templates: - - hpa.yaml -tests: - it: "does not render behavior when not set" set: autoscaling.enabled: true asserts: - isKind: { of: HorizontalPodAutoscaler } - isNull: { path: spec.behavior } + + - it: "scales on cpu at the documented 60 percent by default" + set: + autoscaling.enabled: true + asserts: + - isKind: { of: HorizontalPodAutoscaler } + - equal: { path: "spec.metrics[0].resource.name", value: cpu } + - equal: { path: "spec.metrics[0].resource.target.type", value: Utilization } + - equal: { path: "spec.metrics[0].resource.target.averageUtilization", value: 60 } + + - it: "does not scale on memory by default" + set: + autoscaling.enabled: true + asserts: + - lengthEqual: { path: spec.metrics, count: 1 } + + - it: "honours an explicit cpu target override" + set: + autoscaling.enabled: true + autoscaling.targetCPUUtilizationPercentage: 75 + asserts: + - equal: { path: "spec.metrics[0].resource.target.averageUtilization", value: 75 } + + - it: "renders a memory metric only when a memory target is set" + set: + autoscaling.enabled: true + autoscaling.targetMemoryUtilizationPercentage: 80 + asserts: + - lengthEqual: { path: spec.metrics, count: 2 } + - equal: { path: "spec.metrics[1].resource.name", value: memory } + - equal: { path: "spec.metrics[1].resource.target.averageUtilization", value: 80 } + + - it: "renders no hpa when autoscaling is disabled" + asserts: + - hasDocuments: { count: 0 } diff --git a/helm/litellm-helm/values.yaml b/helm/litellm-helm/values.yaml index f8df98de102..637be2322e3 100644 --- a/helm/litellm-helm/values.yaml +++ b/helm/litellm-helm/values.yaml @@ -200,7 +200,16 @@ autoscaling: enabled: false minReplicas: 1 maxReplicas: 100 - targetCPUUtilizationPercentage: 80 + # 60 is the documented recommendation. See "Recommended Machine Specifications" + # in https://docs.litellm.ai/docs/proxy/prod. A new replica clears the startupProbe + # above only after up to failureThreshold x periodSeconds = 300 seconds, so a target + # high enough to trip near saturation adds capacity minutes after it was needed. + targetCPUUtilizationPercentage: 60 + # Deliberately left unset rather than given a value. The prisma query engine's + # resident memory is a high-water mark that ratchets to the pod's worst-ever write + # and is never returned, so a memory target reads the largest write a pod ever did + # rather than what it is doing now, and replicas ratchet up without scaling back in. + # Memory is a floor to provision under 'resources', not a signal to scale on. # targetMemoryUtilizationPercentage: 80 # behavior: {} From 7a35c34303e944f68e69299346ec04371b5f59c7 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 2 Sep 2026 10:30:14 -0700 Subject: [PATCH 39/48] fix(models): add the us-gov. geo inference profile keys for Claude Sonnet 5 and Opus 4.8 --- ...odel_prices_and_context_window_backup.json | 64 +++++++++++++++++++ model_prices_and_context_window.json | 64 +++++++++++++++++++ .../test_bedrock_usgov_pricing.py | 19 ++++-- 3 files changed, 142 insertions(+), 5 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 30621a17df3..28c503966f3 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -42008,6 +42008,70 @@ "supports_parallel_tool_use_config": true, "prompt_cache_min_tokens": 1024 }, + "us-gov.anthropic.claude-sonnet-5": { + "bedrock_converse_supports_strict_tools": false, + "bedrock_output_config_effort_ceiling": "xhigh", + "cache_creation_input_token_cost": 3e-06, + "cache_creation_input_token_cost_above_1hr": 4.8e-06, + "cache_read_input_token_cost": 2.4e-07, + "input_cost_per_token": 2.4e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 1.2e-05, + "prompt_cache_min_tokens": 1024, + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_mid_conversation_system": true, + "supports_native_structured_output": false, + "supports_output_config": true, + "supports_parallel_tool_use_config": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true + }, + "us-gov.anthropic.claude-opus-4-8": { + "bedrock_converse_supports_strict_tools": false, + "bedrock_output_config_effort_ceiling": "xhigh", + "cache_creation_input_token_cost": 7.5e-06, + "cache_creation_input_token_cost_above_1hr": 1.2e-05, + "cache_read_input_token_cost": 6e-07, + "input_cost_per_token": 6e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 3e-05, + "prompt_cache_min_tokens": 1024, + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_mid_conversation_system": true, + "supports_native_structured_output": true, + "supports_output_config": true, + "supports_parallel_tool_use_config": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true + }, "au.anthropic.claude-haiku-4-5-20251001-v1:0": { "cache_creation_input_token_cost": 1.375e-06, "cache_creation_input_token_cost_above_1hr": 2.2e-06, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 30621a17df3..28c503966f3 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -42008,6 +42008,70 @@ "supports_parallel_tool_use_config": true, "prompt_cache_min_tokens": 1024 }, + "us-gov.anthropic.claude-sonnet-5": { + "bedrock_converse_supports_strict_tools": false, + "bedrock_output_config_effort_ceiling": "xhigh", + "cache_creation_input_token_cost": 3e-06, + "cache_creation_input_token_cost_above_1hr": 4.8e-06, + "cache_read_input_token_cost": 2.4e-07, + "input_cost_per_token": 2.4e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 1.2e-05, + "prompt_cache_min_tokens": 1024, + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_mid_conversation_system": true, + "supports_native_structured_output": false, + "supports_output_config": true, + "supports_parallel_tool_use_config": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true + }, + "us-gov.anthropic.claude-opus-4-8": { + "bedrock_converse_supports_strict_tools": false, + "bedrock_output_config_effort_ceiling": "xhigh", + "cache_creation_input_token_cost": 7.5e-06, + "cache_creation_input_token_cost_above_1hr": 1.2e-05, + "cache_read_input_token_cost": 6e-07, + "input_cost_per_token": 6e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 3e-05, + "prompt_cache_min_tokens": 1024, + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_mid_conversation_system": true, + "supports_native_structured_output": true, + "supports_output_config": true, + "supports_parallel_tool_use_config": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true + }, "au.anthropic.claude-haiku-4-5-20251001-v1:0": { "cache_creation_input_token_cost": 1.375e-06, "cache_creation_input_token_cost_above_1hr": 2.2e-06, diff --git a/tests/test_litellm/test_bedrock_usgov_pricing.py b/tests/test_litellm/test_bedrock_usgov_pricing.py index f9e8fd4c46c..f7d95ecda01 100644 --- a/tests/test_litellm/test_bedrock_usgov_pricing.py +++ b/tests/test_litellm/test_bedrock_usgov_pricing.py @@ -135,15 +135,24 @@ CLAUDE_GOV_EXPECTED = { } +USGOV_CLAUDE_KEY_TEMPLATES = { + "bedrock/us-gov-east-1/{base_key}": "bedrock", + "bedrock/us-gov-west-1/{base_key}": "bedrock", + "us-gov.{base_key}": "bedrock_converse", +} + + @pytest.mark.parametrize("base_key", CLAUDE_GOV_EXPECTED) -@pytest.mark.parametrize("region", ["us-gov-east-1", "us-gov-west-1"]) -def test_usgov_claude_sonnet5_opus48_pricing(model_data, region, base_key): - """Sonnet 5 and Opus 4.8 gov entries must match the rates AWS publishes - for both GovCloud regions on the Bedrock pricing page (1.2x global). +@pytest.mark.parametrize("key_template,expected_provider", USGOV_CLAUDE_KEY_TEMPLATES.items()) +def test_usgov_claude_sonnet5_opus48_pricing(model_data, key_template, expected_provider, base_key): + """Sonnet 5 and Opus 4.8 gov entries, both in-region keys and the us-gov. + geo inference profile the model cards list for GovCloud, must match the + rates AWS publishes on the Bedrock pricing page (1.2x global). """ - gov_key = f"bedrock/{region}/{base_key}" + gov_key = key_template.format(base_key=base_key) assert gov_key in model_data, f"Missing model entry: {gov_key}" info = model_data[gov_key] + assert info["litellm_provider"] == expected_provider for field, expected in CLAUDE_GOV_EXPECTED[base_key].items(): assert info[field] == expected, f"{gov_key}: {field} should be {expected} (got {info[field]})" ratio = info[field] / model_data[base_key][field] From 0c539614451869b790897c0576d202820d85ad59 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 2 Sep 2026 10:44:51 -0700 Subject: [PATCH 40/48] fix(proxy-extras): give prisma migrate deploy its own timeout budget --- .../litellm_proxy_extras/prisma_toolchain.py | 25 +++++- .../litellm_proxy_extras/utils.py | 24 ++++-- .../test_prisma_toolchain.py | 78 ++++++++++++++++++- 3 files changed, 114 insertions(+), 13 deletions(-) diff --git a/litellm-proxy-extras/litellm_proxy_extras/prisma_toolchain.py b/litellm-proxy-extras/litellm_proxy_extras/prisma_toolchain.py index f3b55fd4d96..5feb7a953b4 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/prisma_toolchain.py +++ b/litellm-proxy-extras/litellm_proxy_extras/prisma_toolchain.py @@ -14,9 +14,19 @@ then fails on a Node binary that was never written. Deleting a cache directory that exists without a Node binary is what turns a killed bootstrap back into a recoverable one. -Both budgets are overridable so an operator can widen them without a release: -``LITELLM_PRISMA_BOOTSTRAP_TIMEOUT`` for the toolchain install and -``LITELLM_PRISMA_COMMAND_TIMEOUT`` for every individual Prisma command. +``prisma migrate deploy`` is the other command whose runtime is not a +constant: it grows with the number of pending migrations, so a fresh database +that has to replay every migration this package ships overruns a per-command +budget sized for the short bookkeeping commands, on a laptop as much as on a +slow CI runner. The Python ``prisma`` wrapper spawns Node and the schema engine +as separate children, so killing the wrapper on timeout leaves them running: +the retry then contends with that orphan for Prisma's advisory lock and cannot +finish any sooner. Migrate deploy therefore runs under its own budget. + +All three budgets are overridable so an operator can widen them without a +release: ``LITELLM_PRISMA_BOOTSTRAP_TIMEOUT`` for the toolchain install, +``LITELLM_PRISMA_MIGRATE_DEPLOY_TIMEOUT`` for ``prisma migrate deploy`` and +``LITELLM_PRISMA_COMMAND_TIMEOUT`` for every other Prisma command. """ import math @@ -36,10 +46,12 @@ except ImportError: PRISMA_COMMAND_TIMEOUT_ENV_VAR = "LITELLM_PRISMA_COMMAND_TIMEOUT" PRISMA_BOOTSTRAP_TIMEOUT_ENV_VAR = "LITELLM_PRISMA_BOOTSTRAP_TIMEOUT" +PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR = "LITELLM_PRISMA_MIGRATE_DEPLOY_TIMEOUT" NODEENV_CACHE_DIR_ENV_VAR = "PRISMA_NODEENV_CACHE_DIR" DEFAULT_PRISMA_COMMAND_TIMEOUT = 60.0 DEFAULT_PRISMA_BOOTSTRAP_TIMEOUT = 600.0 +DEFAULT_PRISMA_MIGRATE_DEPLOY_TIMEOUT = 600.0 BOOTSTRAP_ARG = "--version" @@ -88,6 +100,13 @@ def prisma_bootstrap_timeout() -> float: ) +def prisma_migrate_deploy_timeout() -> float: + """Seconds one ``prisma migrate deploy`` may run for, however many migrations are pending.""" + return _timeout_from_env( + PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR, DEFAULT_PRISMA_MIGRATE_DEPLOY_TIMEOUT + ) + + def nodeenv_cache_dir() -> Optional[Path]: """Where Prisma installs its private Node runtime, or None if unknowable.""" override = os.getenv(NODEENV_CACHE_DIR_ENV_VAR) diff --git a/litellm-proxy-extras/litellm_proxy_extras/utils.py b/litellm-proxy-extras/litellm_proxy_extras/utils.py index b8032dd0d28..ab9ec1e8a3a 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/utils.py +++ b/litellm-proxy-extras/litellm_proxy_extras/utils.py @@ -15,8 +15,10 @@ from litellm_proxy_extras.replica_identity import ( apply_replica_identity_full, ) from litellm_proxy_extras.prisma_toolchain import ( + PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR, ensure_prisma_toolchain, prisma_command_timeout, + prisma_migrate_deploy_timeout, ) @@ -698,12 +700,13 @@ class ProxyExtrasDBManager: original_dir = os.getcwd() os.chdir(migrations_dir) + deploy_timeout = prisma_migrate_deploy_timeout() try: for attempt in range(4): try: result = subprocess.run( [_get_prisma_command(), "migrate", "deploy"], - timeout=prisma_command_timeout(), + timeout=deploy_timeout, check=True, capture_output=True, text=True, @@ -713,8 +716,12 @@ class ProxyExtrasDBManager: return True except subprocess.TimeoutExpired: - logger.info( - f"prisma migrate deploy attempt {attempt + 1} timed out, retrying" + logger.warning( + "prisma migrate deploy attempt %s timed out after %ss, retrying. " + "Raise %s if this database needs longer to apply its pending migrations.", + attempt + 1, + deploy_timeout, + PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR, ) time.sleep(random.randrange(5, 15)) continue @@ -823,7 +830,8 @@ class ProxyExtrasDBManager: "Database migration failed after 4 attempts (retry loop " "exhausted by timeouts or repeated idempotent-recovery " "continues). Check database connectivity, load, and " - "_prisma_migrations ledger state." + "_prisma_migrations ledger state, and raise " + f"{PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR} if the attempts timed out." ) finally: os.chdir(original_dir) @@ -908,7 +916,7 @@ class ProxyExtrasDBManager: # Set migrations directory for Prisma result = subprocess.run( [_get_prisma_command(), "migrate", "deploy"], - timeout=prisma_command_timeout(), + timeout=prisma_migrate_deploy_timeout(), check=True, capture_output=True, text=True, @@ -1126,7 +1134,11 @@ class ProxyExtrasDBManager: ) return True except subprocess.TimeoutExpired: - logger.info(f"Attempt {attempt + 1} timed out") + logger.warning( + "Attempt %s timed out. Raise %s if this database needs longer to apply its pending migrations.", + attempt + 1, + PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR, + ) time.sleep(random.randrange(5, 15)) except subprocess.CalledProcessError as e: attempts_left = 3 - attempt diff --git a/tests/proxy_migration_tests/test_prisma_toolchain.py b/tests/proxy_migration_tests/test_prisma_toolchain.py index d12a2c4dd4e..1bd440d37b4 100644 --- a/tests/proxy_migration_tests/test_prisma_toolchain.py +++ b/tests/proxy_migration_tests/test_prisma_toolchain.py @@ -7,6 +7,11 @@ attempt fails identically. These tests pin the two behaviours that keep a container recoverable: an incomplete cache is deleted before Prisma is invoked, and the install gets a budget of its own rather than sharing the one that bounds each migration command. + +``prisma migrate deploy`` gets a budget of its own for the same reason: its +runtime grows with the number of pending migrations, so a fresh database that +replays every migration overran the per-command budget on slow machines and +the proxy gave up after four identical timeouts. """ import ast @@ -14,19 +19,23 @@ import json import os import sys import time +from collections.abc import Callable from pathlib import Path import pytest from litellm_proxy_extras.prisma_toolchain import ( DEFAULT_PRISMA_COMMAND_TIMEOUT, + DEFAULT_PRISMA_MIGRATE_DEPLOY_TIMEOUT, PRISMA_BOOTSTRAP_TIMEOUT_ENV_VAR, PRISMA_COMMAND_TIMEOUT_ENV_VAR, + PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR, ensure_prisma_toolchain, heal_incomplete_nodeenv_cache, node_binary_path, prisma_bootstrap_timeout, prisma_command_timeout, + prisma_migrate_deploy_timeout, ) from litellm_proxy_extras.utils import ProxyExtrasDBManager @@ -42,13 +51,24 @@ import time args = sys.argv[1:] cache_dir = os.environ["PRISMA_NODEENV_CACHE_DIR"] -with pathlib.Path(os.environ["FAKE_PRISMA_LOG"]).open("a") as log: +log_path = pathlib.Path(os.environ["FAKE_PRISMA_LOG"]) +earlier_deploys = sum( + 1 + for line in (log_path.read_text().splitlines() if log_path.exists() else []) + if json.loads(line)["args"][:2] == ["migrate", "deploy"] +) +with log_path.open("a") as log: log.write( json.dumps({{"args": args, "cache_dir_present": os.path.isdir(cache_dir)}}) + "\\n" ) time.sleep(float(os.environ.get("FAKE_PRISMA_SLEEP", "0"))) if args[:2] == ["migrate", "deploy"]: + if earlier_deploys == 0: + time.sleep(float(os.environ.get("FAKE_PRISMA_FIRST_DEPLOY_SLEEP", "0"))) + elif os.environ.get("FAKE_PRISMA_LATER_DEPLOY_STDERR"): + print(os.environ["FAKE_PRISMA_LATER_DEPLOY_STDERR"], file=sys.stderr) + sys.exit(1) print("No pending migrations to apply") sys.exit(0) """ @@ -80,9 +100,14 @@ def toolchain_env(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> tuple[Path monkeypatch.setenv("PATH", f"{bin_dir}{os.pathsep}{os.environ['PATH']}") monkeypatch.delenv(PRISMA_COMMAND_TIMEOUT_ENV_VAR, raising=False) monkeypatch.delenv(PRISMA_BOOTSTRAP_TIMEOUT_ENV_VAR, raising=False) + monkeypatch.delenv(PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR, raising=False) return cache_dir, log_path +def _deploy_calls(log_path: Path) -> list[list[str]]: + return [call["args"] for call in _fake_prisma_calls(log_path) if call["args"][:2] == ["migrate", "deploy"]] + + def _make_incomplete_cache(cache_dir: Path) -> None: (cache_dir / "lib").mkdir(parents=True) (cache_dir / "bin").mkdir() @@ -209,25 +234,70 @@ def test_setup_database_prepares_the_toolchain_before_migrating( assert calls[0]["cache_dir_present"] is False +@pytest.mark.parametrize("use_v2_resolver", [False, True], ids=["v1", "v2"]) +def test_migrate_deploy_is_not_bounded_by_the_per_command_timeout( + toolchain_env: tuple[Path, Path], + monkeypatch: pytest.MonkeyPatch, + use_v2_resolver: bool, +) -> None: + """A fresh database replays every migration, which takes longer than any bookkeeping command.""" + _, log_path = toolchain_env + monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@localhost:9/x") + monkeypatch.setenv(PRISMA_COMMAND_TIMEOUT_ENV_VAR, "1") + monkeypatch.setenv("FAKE_PRISMA_FIRST_DEPLOY_SLEEP", "3") + + assert ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=use_v2_resolver) is True + assert _deploy_calls(log_path) == [["migrate", "deploy"]] + + +def test_migrate_deploy_stops_at_its_own_timeout( + toolchain_env: tuple[Path, Path], monkeypatch: pytest.MonkeyPatch +) -> None: + """The deploy budget still bounds a deploy that hangs, so boot cannot wait forever.""" + _, log_path = toolchain_env + monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@localhost:9/x") + monkeypatch.setenv(PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR, "1") + monkeypatch.setenv("FAKE_PRISMA_FIRST_DEPLOY_SLEEP", "60") + monkeypatch.setenv("FAKE_PRISMA_LATER_DEPLOY_STDERR", "Error: P3018 permission denied for schema public") + + started = time.monotonic() + with pytest.raises(RuntimeError, match="insufficient permissions"): + ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True) + elapsed = time.monotonic() - started + + assert len(_deploy_calls(log_path)) == 2 + assert elapsed < 30 + + @pytest.mark.parametrize( "raw", ["", "0", "-5", "not-a-number", "nan", "inf", "-inf", "1e400"], ) +@pytest.mark.parametrize( + ("env_var", "read_timeout", "default"), + [ + (PRISMA_COMMAND_TIMEOUT_ENV_VAR, prisma_command_timeout, DEFAULT_PRISMA_COMMAND_TIMEOUT), + (PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR, prisma_migrate_deploy_timeout, DEFAULT_PRISMA_MIGRATE_DEPLOY_TIMEOUT), + ], + ids=["command", "migrate_deploy"], +) def test_unusable_timeout_override_falls_back_to_the_default( - raw: str, monkeypatch: pytest.MonkeyPatch + raw: str, env_var: str, read_timeout: Callable[[], float], default: float, monkeypatch: pytest.MonkeyPatch ) -> None: """A non-finite override would silently disable the timeout it configures.""" - monkeypatch.setenv(PRISMA_COMMAND_TIMEOUT_ENV_VAR, raw) + monkeypatch.setenv(env_var, raw) - assert prisma_command_timeout() == DEFAULT_PRISMA_COMMAND_TIMEOUT + assert read_timeout() == default def test_timeout_overrides_are_independent(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setenv(PRISMA_COMMAND_TIMEOUT_ENV_VAR, "12") monkeypatch.setenv(PRISMA_BOOTSTRAP_TIMEOUT_ENV_VAR, "900") + monkeypatch.setenv(PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR, "1200") assert prisma_command_timeout() == 12 assert prisma_bootstrap_timeout() == 900 + assert prisma_migrate_deploy_timeout() == 1200 @pytest.mark.parametrize("module", ["utils.py", "replica_identity.py"]) From ffc0a8e428a4d8af7e5b130bddb4f0f9cf0cb229 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 2 Sep 2026 10:51:17 -0700 Subject: [PATCH 41/48] fix: run access group key sync UPDATEs on the writer, not the read replica (#39128) Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../access_group_key_sync.py | 6 +- .../test_access_group_key_sync.py | 57 +++++++++++++++++++ 2 files changed, 61 insertions(+), 2 deletions(-) create mode 100644 tests/test_litellm/proxy/management_helpers/test_access_group_key_sync.py diff --git a/litellm/proxy/management_helpers/access_group_key_sync.py b/litellm/proxy/management_helpers/access_group_key_sync.py index 5d43cb29978..c9f93fae0d9 100644 --- a/litellm/proxy/management_helpers/access_group_key_sync.py +++ b/litellm/proxy/management_helpers/access_group_key_sync.py @@ -38,6 +38,7 @@ from litellm.proxy._types import ( from litellm.proxy.auth.auth_checks import ( _delete_cache_access_object, # pyright: ignore[reportPrivateUsage] # the access-group endpoints reach for this same cache primitive ) +from litellm.proxy.db.routing_prisma_wrapper import WriterPinnedClient from litellm.repositories.table_repositories import AccessGroupRepository @@ -72,8 +73,9 @@ _REPOINT_KEY_SQL: Final = ( def _raw_executor(prisma_client: object) -> _RawExecutor: - """Narrow the untyped Prisma client down to the raw-query call this module makes.""" - return AccessGroupRepository(prisma_client).prisma_client.db # pyright: ignore[reportAny] # untyped Prisma client + """Narrow the untyped Prisma client down to the raw-query call this module makes, pinned to the writer.""" + db: Final = AccessGroupRepository(prisma_client).prisma_client.db # pyright: ignore[reportAny] # untyped Prisma client + return WriterPinnedClient(db).db # pyright: ignore[reportAny, reportReturnType] # untyped Prisma client behind the pin async def _invalidate_access_group_cache(access_group_id: str) -> None: diff --git a/tests/test_litellm/proxy/management_helpers/test_access_group_key_sync.py b/tests/test_litellm/proxy/management_helpers/test_access_group_key_sync.py new file mode 100644 index 00000000000..60c36e33e09 --- /dev/null +++ b/tests/test_litellm/proxy/management_helpers/test_access_group_key_sync.py @@ -0,0 +1,57 @@ +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from litellm.proxy.db.prisma_client import PrismaWrapper +from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper +from litellm.proxy.management_helpers.access_group_key_sync import ( + sync_key_access_group_membership, + sync_key_regeneration_access_group_membership, +) + + +def _routed_prisma_client(): + writer_inner = MagicMock(name="writer_prisma") + reader_inner = MagicMock(name="reader_prisma") + writer_inner.query_raw = AsyncMock(return_value=[]) + reader_inner.query_raw = AsyncMock(return_value=[]) + writer = PrismaWrapper(original_prisma=writer_inner, iam_token_db_auth=False) + reader = PrismaWrapper(original_prisma=reader_inner, iam_token_db_auth=False) + routing = RoutingPrismaWrapper(writer=writer, reader=reader) + return SimpleNamespace(db=routing), writer_inner, reader_inner + + +@pytest.mark.asyncio +async def test_regeneration_repoint_update_runs_on_the_writer(): + prisma_client, writer_inner, reader_inner = _routed_prisma_client() + + await sync_key_regeneration_access_group_membership( + prisma_client=prisma_client, + previous_key_token="old-token", + new_key_token="new-token", + data=None, + existing_key_row=MagicMock(), + ) + + writer_inner.query_raw.assert_awaited_once() + assert writer_inner.query_raw.await_args.args[0].startswith('UPDATE "LiteLLM_AccessGroupTable"') + reader_inner.query_raw.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_membership_attach_and_detach_updates_run_on_the_writer(): + prisma_client, writer_inner, reader_inner = _routed_prisma_client() + + await sync_key_access_group_membership( + prisma_client=prisma_client, + key_token="token", + previous_access_group_ids=["ag-old"], + updated_access_group_ids=["ag-new"], + ) + + assert writer_inner.query_raw.await_count == 2 + assert all( + call.args[0].startswith('UPDATE "LiteLLM_AccessGroupTable"') for call in writer_inner.query_raw.await_args_list + ) + reader_inner.query_raw.assert_not_awaited() From a9d3a0746c582de8910d0a7078489ac961e436e8 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 2 Sep 2026 10:59:11 -0700 Subject: [PATCH 42/48] fix(models): price Azure DeepSeek V4 Flash 0731 from its own meters under the catalog id --- ...odel_prices_and_context_window_backup.json | 22 +++---------------- model_prices_and_context_window.json | 22 +++---------------- 2 files changed, 6 insertions(+), 38 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 893ac49f5c7..0af5a89742e 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -10172,31 +10172,15 @@ "supports_tool_choice": true }, "azure_ai/DeepSeek-V4-Flash-0731": { - "cache_read_input_token_cost": 2.8e-08, + "cache_read_input_token_cost": 1.4e-08, "deprecation_date": "2026-12-03", - "input_cost_per_token": 1.9e-07, + "input_cost_per_token": 4.4e-07, "litellm_provider": "azure_ai", "max_input_tokens": 1000000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", - "output_cost_per_token": 5.1e-07, - "source": "https://azure.microsoft.com/en-us/pricing/details/ai-foundry-models/deepseek/", - "supports_function_calling": true, - "supports_prompt_caching": true, - "supports_reasoning": true, - "supports_tool_choice": true - }, - "azure_ai/deepseek-v4-flash-0731": { - "cache_read_input_token_cost": 2.8e-08, - "deprecation_date": "2026-12-03", - "input_cost_per_token": 1.9e-07, - "litellm_provider": "azure_ai", - "max_input_tokens": 1000000, - "max_output_tokens": 128000, - "max_tokens": 128000, - "mode": "chat", - "output_cost_per_token": 5.1e-07, + "output_cost_per_token": 1.32e-06, "source": "https://azure.microsoft.com/en-us/pricing/details/ai-foundry-models/deepseek/", "supports_function_calling": true, "supports_prompt_caching": true, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 893ac49f5c7..0af5a89742e 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -10172,31 +10172,15 @@ "supports_tool_choice": true }, "azure_ai/DeepSeek-V4-Flash-0731": { - "cache_read_input_token_cost": 2.8e-08, + "cache_read_input_token_cost": 1.4e-08, "deprecation_date": "2026-12-03", - "input_cost_per_token": 1.9e-07, + "input_cost_per_token": 4.4e-07, "litellm_provider": "azure_ai", "max_input_tokens": 1000000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", - "output_cost_per_token": 5.1e-07, - "source": "https://azure.microsoft.com/en-us/pricing/details/ai-foundry-models/deepseek/", - "supports_function_calling": true, - "supports_prompt_caching": true, - "supports_reasoning": true, - "supports_tool_choice": true - }, - "azure_ai/deepseek-v4-flash-0731": { - "cache_read_input_token_cost": 2.8e-08, - "deprecation_date": "2026-12-03", - "input_cost_per_token": 1.9e-07, - "litellm_provider": "azure_ai", - "max_input_tokens": 1000000, - "max_output_tokens": 128000, - "max_tokens": 128000, - "mode": "chat", - "output_cost_per_token": 5.1e-07, + "output_cost_per_token": 1.32e-06, "source": "https://azure.microsoft.com/en-us/pricing/details/ai-foundry-models/deepseek/", "supports_function_calling": true, "supports_prompt_caching": true, From e0be9a35e6838505b5a2ae2ecf93c01579a02f56 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 2 Sep 2026 11:02:43 -0700 Subject: [PATCH 43/48] fix(deps): raise the pypdf floor to 6.16.1 for three new advisories GHSA-jp53-mhqp-8xcg (fixed in 6.16.0), GHSA-23w6-3w8w-8484 and GHSA-763m-79hh-57f2 (fixed in 6.16.1) flag pypdf 6.15.0 in uv.lock and keep osv-scan red alongside the tornado advisories. The proxy-runtime extra now requires pypdf>=6.16.1 and the lock resolves 6.16.2. --- pyproject.toml | 2 +- uv.lock | 10 +++++----- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index d0f14722acd..60162544612 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -161,7 +161,7 @@ proxy-runtime = [ "mangum>=0.17.0,<1.0", "azure-ai-contentsafety>=1.0.0,<2.0", "azure-storage-file-datalake>=12.20.0,<13.0", - "pypdf>=6.12.0,<7.0", + "pypdf>=6.16.1,<7.0", "llm-sandbox>=0.3.39,<1.0", "detect-secrets>=1.5.0,<2.0", ] diff --git a/uv.lock b/uv.lock index 8bac024d49e..aa59ff7b229 100644 --- a/uv.lock +++ b/uv.lock @@ -10,7 +10,7 @@ resolution-markers = [ ] [options] -exclude-newer = "2026-08-29T20:52:40.322465Z" +exclude-newer = "2026-08-30T17:51:25.171404Z" exclude-newer-span = "P3D" [manifest] @@ -4552,7 +4552,7 @@ requires-dist = [ { name = "pydantic-settings", specifier = ">=2.14.1,<3.0" }, { name = "pyjwt", marker = "extra == 'proxy'", specifier = ">=2.13.0,<3.0" }, { name = "pynacl", marker = "extra == 'proxy'", specifier = ">=1.6.2,<2.0" }, - { name = "pypdf", marker = "extra == 'proxy-runtime'", specifier = ">=6.12.0,<7.0" }, + { name = "pypdf", marker = "extra == 'proxy-runtime'", specifier = ">=6.16.1,<7.0" }, { name = "pyroscope-io", marker = "sys_platform != 'win32' and extra == 'proxy'", specifier = ">=0.8.16,<1.0" }, { name = "python-dotenv", specifier = ">=1.0.0,<2.0" }, { name = "python-multipart", marker = "extra == 'proxy'", specifier = ">=0.0.27,<1.0" }, @@ -7564,14 +7564,14 @@ wheels = [ [[package]] name = "pypdf" -version = "6.15.0" +version = "6.16.2" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "typing-extensions", marker = "python_full_version < '3.11'" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/17/17/ee75a92718ec7212de831e71454d702225aa5e474a805cce169806044453/pypdf-6.15.0.tar.gz", hash = "sha256:d39c4d955a76409284a905e2d65b40076d77ab76129e0faaeeb6612403ecfc79", size = 6993794, upload-time = "2026-08-06T13:06:49.929Z" } +sdist = { url = "https://files.pythonhosted.org/packages/44/66/54212e75406afd9f3e933d0dda23072f6aecc55c5a273077dc2e0b028b23/pypdf-6.16.2.tar.gz", hash = "sha256:595647f6191de6f402cfde1d0c455d6cbccbd509aac32b34783009c032de5d6e", size = 7008996, upload-time = "2026-08-23T13:50:07.135Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/af/72/ce3067ac31e214a66388159f8462ddb8c13dd00170f24d555a1f1ae8ee91/pypdf-6.15.0-py3-none-any.whl", hash = "sha256:14e001d6504822cb1ca9c7ed9a69bccb320f59b320730f55af804361abe4d5ee", size = 378123, upload-time = "2026-08-06T13:06:47.709Z" }, + { url = "https://files.pythonhosted.org/packages/13/f1/a2da3b55acd4ab737bf728c97edaaed5ec1d3c1236acb639dcdfa97e42c7/pypdf-6.16.2-py3-none-any.whl", hash = "sha256:c8b09a59399062fb45a1b8156c18a787a10a3dae03ac9674397a226712c94604", size = 385060, upload-time = "2026-08-23T13:50:05.349Z" }, ] [[package]] From a677242d6f07af683b9c146133287f07b4e1459c Mon Sep 17 00:00:00 2001 From: Ali Ahmed <128928915+QuantumBreakz@users.noreply.github.com> Date: Thu, 3 Sep 2026 00:07:13 +0500 Subject: [PATCH 44/48] fix(headroom): stop re-compressing retrieved CCR content in client tool loops (#38591) When the headroom_retrieve tool is exposed to a client that runs its own tool-execution loop (the LiteLLM MCP gateway path), the client executes the retrieve call and sends the recovered original content back as a tool result on the next turn. The guardrail then compressed that row again, and because CCR is content-addressed it collapsed back to the exact same hash it was just retrieved from. The model never saw the expansion and the agent looped. Hold tool-result rows that carry headroom_retrieve output back from the compression service, the same way the live turn and trailing tool exchange are already protected, so the expansion survives. Retrieve calls are matched by the direct headroom_retrieve name and the mcp____headroom_retrieve gateway name. Because a long gateway name is truncated past 64 chars in the OpenAI-translated view the guardrail scans, the pairing also falls back to the tool-call id read from the request's own untranslated messages, which is never truncated. Fixes #38558 --- .../guardrail_hooks/headroom/headroom.py | 120 ++++++++++++- .../guardrail_hooks/test_headroom.py | 158 ++++++++++++++++++ 2 files changed, 274 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py b/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py index d8c8c2f4974..fc881a60f43 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py +++ b/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py @@ -10,6 +10,7 @@ from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, TypeGuard import httpx from fastapi import HTTPException from httpx import Response as HttpxResponse +from pydantic import TypeAdapter import litellm from litellm._logging import verbose_proxy_logger @@ -52,6 +53,10 @@ BYPASS_HEADER: Final = "x-headroom-bypass" HEADROOM_RETRIEVE_TOOL_NAME: Final = "headroom_retrieve" _HASH_PATTERN: Final = re.compile(r"hash=([a-f0-9]{24})") _HASH_CACHE_TTL_SECONDS: Final = 15 * 60 +# Narrows the base class's bare-dict ``request_data`` at the boundary so its +# untranslated messages can be read with concrete types (values pass through by +# reference, so this is a shallow top-level reconstruction). +_REQUEST_DATA_ADAPTER: Final = TypeAdapter(dict[str, object]) def _is_str_object_dict(value: object) -> TypeGuard[dict[str, object]]: # guard-ok: isinstance narrows correctly; predicate is trivially correct # fmt: skip @@ -116,16 +121,119 @@ def _restore_content_shapes( return restored -def _protected_indices(messages: Sequence[Mapping[str, object]]) -> frozenset[int]: +def _tool_call_name(tool_call: Mapping[str, object]) -> str | None: + function: Final = tool_call.get("function") + if not _is_str_object_dict(function): + return None + name: Final = function.get("name") + return name if isinstance(name, str) else None + + +def _is_retrieve_tool_name(name: str | None) -> bool: + """Match the retrieve tool whether called directly or via the MCP gateway. + + Server-side the tool is ``headroom_retrieve``; exposed through LiteLLM's MCP + gateway a client calls it as ``mcp____headroom_retrieve``. + """ + return name is not None and ( + name == HEADROOM_RETRIEVE_TOOL_NAME or name.endswith(f"__{HEADROOM_RETRIEVE_TOOL_NAME}") + ) + + +def _retrieve_call_ids_in_message(message: Mapping[str, object]) -> frozenset[str]: + if message.get("role") != "assistant": + return frozenset() + tool_calls: Final = message.get("tool_calls") + if not _is_object_list(tool_calls): + return frozenset() + return frozenset( + str(tool_call["id"]) + for tool_call in tool_calls + if _is_str_object_dict(tool_call) and tool_call.get("id") and _is_retrieve_tool_name(_tool_call_name(tool_call)) + ) + + +def _anthropic_tool_use_retrieve_id(block: object) -> str | None: + if not _is_str_object_dict(block) or block.get("type") != "tool_use": + return None + name: Final = block.get("name") + call_id: Final = block.get("id") + if isinstance(name, str) and call_id is not None and _is_retrieve_tool_name(name): + return str(call_id) + return None + + +def _anthropic_retrieve_ids_in_message(message: Mapping[str, object]) -> frozenset[str]: + content: Final = message.get("content") + if not _is_object_list(content): + return frozenset() + return frozenset(call_id for block in content if (call_id := _anthropic_tool_use_retrieve_id(block)) is not None) + + +def _raw_retrieve_call_ids(messages: object) -> frozenset[str]: + """Retrieve-tool call ids read from the request's own, untranslated messages. + + The guardrail otherwise scans an OpenAI-translated view where a tool name + over 64 chars is truncated to ``{prefix}_{hash}``, which drops the + ``__headroom_retrieve`` suffix a long ``mcp____`` prefix pushes past + the limit. Tool-call ids are never truncated, so pairing the tool result to + an id read from the original request keeps the match intact. Both wire + shapes are handled: OpenAI ``tool_calls`` and Anthropic ``tool_use`` blocks. + """ + if not _is_object_list(messages): + return frozenset() + return frozenset( + call_id + for message in messages + if _is_str_object_dict(message) + for call_id in _retrieve_call_ids_in_message(message) | _anthropic_retrieve_ids_in_message(message) + ) + + +def _retrieval_result_indices( + messages: Sequence[Mapping[str, object]], extra_retrieve_call_ids: frozenset[str] = frozenset() +) -> frozenset[int]: + """Indices of tool-result rows that carry ``headroom_retrieve`` output. + + When the retrieve tool is exposed to a client that runs its own tool loop + (the LiteLLM MCP gateway path), the client executes the call and sends the + recovered original content back as a tool result on the next turn. That + content is exactly what a prior compression stubbed, so compressing it again + re-derives the identical content hash: a no-op that strands the model on the + marker and loops the agent. Hold those rows back so the expansion survives. + + ``extra_retrieve_call_ids`` carries ids recovered from the untruncated + request so the pairing survives tool-name truncation (see + ``_raw_retrieve_call_ids``). + """ + retrieve_call_ids: Final = extra_retrieve_call_ids | frozenset( + call_id for message in messages for call_id in _retrieve_call_ids_in_message(message) + ) + if not retrieve_call_ids: + return frozenset() + return frozenset( + index + for index, message in enumerate(messages) + if message.get("role") in ("tool", "function") and str(message.get("tool_call_id")) in retrieve_call_ids + ) + + +def _protected_indices( + messages: Sequence[Mapping[str, object]], extra_retrieve_call_ids: frozenset[str] = frozenset() +) -> frozenset[int]: """Indices headroom must not send to the compression service. ``get_protected_indices`` is litellm's own compression policy: the system - rows, the last user row, the last assistant row. It is expanded over whole + rows, the last user row, the last assistant row. Rows carrying just-retrieved + ``headroom_retrieve`` output are added so re-compression can't collapse them + back to the marker they were expanded from. The union is expanded over whole tool exchanges the way ``compress()`` expands it, so a protected assistant tool call cannot end up answered by a marker standing in for the result the model just asked for. """ - protected: Final = frozenset(get_protected_indices(messages)) + protected: Final = frozenset(get_protected_indices(messages)) | _retrieval_result_indices( + messages, extra_retrieve_call_ids + ) return protected | frozenset( index for group in group_tool_exchanges(messages) @@ -634,7 +742,11 @@ class HeadroomGuardrail(CustomGuardrail): # /v1/compress grows a field for sending the live turn as the retrieval # query without compressing it: query-aware compression reads the newest # user message, so it is withheld here at some cost to history ranking. - protected_indices: Final = _protected_indices(messages) + # request_data is a bare dict on the base signature; narrow it before + # reading the untranslated messages so long tool names can be recovered. + raw_messages: Final = _REQUEST_DATA_ADAPTER.validate_python(request_data).get("messages") + raw_retrieve_call_ids: Final = _raw_retrieve_call_ids(raw_messages) + protected_indices: Final = _protected_indices(messages, raw_retrieve_call_ids) compressible: Final = [m for i, m in enumerate(messages) if i not in protected_indices] if not compressible: return inputs diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py index 1fbc975e40a..c04fb7b30ec 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py @@ -2199,6 +2199,164 @@ async def test_history_is_still_compressed(guardrail: HeadroomGuardrail): assert has_headroom_retrieve_tool(result.get("tools") or []) +# --------------------------------------------------------------------------- +# #38558: a client that runs its own tool loop (e.g. Claude Code via the MCP +# gateway) executes headroom_retrieve and echoes the recovered original content +# back as a tool result. Compressing that row re-derives the same content hash +# it was just retrieved from -- the marker returns and the agent loops. The +# retrieved row must be held back from the compression service. +# --------------------------------------------------------------------------- + +RETRIEVE_ECHO_MESSAGES = [ + {"role": "system", "content": "You are Claude Code. " + "S" * 5000}, + {"role": "user", "content": "H" * 5000}, + { + "role": "assistant", + "content": "Expanding the marker.", + "tool_calls": [ + { + "id": "hr_1", + "type": "function", + "function": { + "name": "mcp__headroom__headroom_retrieve", + "arguments": '{"hash": "b573993006976af767214fac"}', + }, + } + ], + }, + {"role": "tool", "tool_call_id": "hr_1", "content": "RETRIEVED BODY " + "R" * 5000}, + {"role": "assistant", "content": "Older answer. " + "O" * 5000}, + {"role": "user", "content": "now summarize the description"}, +] + + +@pytest.mark.asyncio +async def test_retrieved_content_is_never_recompressed(guardrail: HeadroomGuardrail): + """The tool result carrying headroom_retrieve output is held back, so it can + never collapse back to the hash it was just retrieved from.""" + wire, result = await _wire_and_result(guardrail, RETRIEVE_ECHO_MESSAGES) + + assert not any(row.get("tool_call_id") == "hr_1" for row in wire) + assert not any("RETRIEVED BODY" in json.dumps(row) for row in wire) + # Reaches the model byte-identical, so no marker stands in for the expansion. + assert result["structured_messages"][3] == RETRIEVE_ECHO_MESSAGES[3] + # Negative control: unrelated history is still compressed, not a no-op. + assert any(row.get("content") == "H" * 5000 for row in wire) + + +@pytest.mark.asyncio +async def test_retrieved_content_guard_matches_direct_tool_name(guardrail: HeadroomGuardrail): + """Server-side the tool is named headroom_retrieve (no MCP prefix); its + result must be protected the same way.""" + messages = [ + {"role": "system", "content": "sys " + "S" * 5000}, + {"role": "user", "content": "H" * 5000}, + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "hr_direct", + "type": "function", + "function": {"name": HEADROOM_RETRIEVE_TOOL_NAME, "arguments": "{}"}, + } + ], + }, + {"role": "tool", "tool_call_id": "hr_direct", "content": "RETRIEVED BODY " + "R" * 5000}, + {"role": "assistant", "content": "Older. " + "O" * 5000}, + {"role": "user", "content": "summarize"}, + ] + wire, result = await _wire_and_result(guardrail, messages) + + assert not any(row.get("tool_call_id") == "hr_direct" for row in wire) + assert result["structured_messages"][3] == messages[3] + + +@pytest.mark.asyncio +async def test_retrieved_content_protected_when_mcp_tool_name_is_truncated(guardrail: HeadroomGuardrail): + """A long mcp____headroom_retrieve name is truncated past 64 chars in + the OpenAI-translated view the guardrail scans, dropping the suffix. The call + id read from the request's own Anthropic tool_use (never truncated) still + pairs the retrieved row so it is held back.""" + from litellm.llms.anthropic.experimental_pass_through.adapters.transformation import ( + truncate_tool_name, + ) + + long_name = "mcp__" + "s" * 45 + "__" + HEADROOM_RETRIEVE_TOOL_NAME + assert len(long_name) > 64 + truncated = truncate_tool_name(long_name) + assert not truncated.endswith(HEADROOM_RETRIEVE_TOOL_NAME) + + # What the guardrail scans: OpenAI-translated messages with the truncated name. + structured = [ + {"role": "system", "content": "sys " + "S" * 5000}, + {"role": "user", "content": "H" * 5000}, + { + "role": "assistant", + "content": "", + "tool_calls": [{"id": "hr_long", "type": "function", "function": {"name": truncated, "arguments": "{}"}}], + }, + {"role": "tool", "tool_call_id": "hr_long", "content": "RETRIEVED BODY " + "R" * 5000}, + {"role": "assistant", "content": "Older. " + "O" * 5000}, + {"role": "user", "content": "summarize"}, + ] + # The request's own messages, untranslated: Anthropic tool_use carries the full name. + raw_messages = [ + {"role": "assistant", "content": [{"type": "tool_use", "id": "hr_long", "name": long_name, "input": {}}]}, + {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "hr_long", "content": "RETRIEVED BODY"}]}, + ] + + inputs = GenericGuardrailAPIInputs(texts=["x"], structured_messages=json.loads(json.dumps(structured))) + sent: dict = {} + + def _echo(**kwargs): + sent["messages"] = kwargs["json"]["messages"] + return _make_compress_response(json.loads(json.dumps(kwargs["json"]["messages"]))) + + with patch.object(guardrail.async_handler, "post", new_callable=AsyncMock, side_effect=_echo): + result = await guardrail.apply_guardrail( + inputs=inputs, + request_data={"model": "claude-sonnet-4-5-20250929", "messages": raw_messages}, + input_type="request", + ) + + assert not any(row.get("tool_call_id") == "hr_long" for row in sent["messages"]) + assert result["structured_messages"][3] == structured[3] + assert any(row.get("content") == "H" * 5000 for row in sent["messages"]) + + +def test_raw_retrieve_call_ids_covers_both_shapes_and_ignores_others(): + """Retrieve ids are read from OpenAI tool_calls and Anthropic tool_use blocks; + non-retrieve calls, non-tool_use blocks, string content, and non-list inputs + yield nothing.""" + from litellm.proxy.guardrails.guardrail_hooks.headroom.headroom import _raw_retrieve_call_ids + + messages = [ + { + "role": "assistant", + "tool_calls": [ + {"id": "oa1", "function": {"name": HEADROOM_RETRIEVE_TOOL_NAME}}, + {"id": "other", "function": {"name": "get_weather"}}, + {"id": "malformed", "function": {"name": 123}}, + {"id": "nofunc"}, + ], + }, + { + "role": "assistant", + "content": [ + {"type": "tool_use", "id": "an1", "name": "mcp__hr__headroom_retrieve", "input": {}}, + {"type": "tool_use", "id": "an2", "name": "jira_get_issue", "input": {}}, + {"type": "text", "text": "noise"}, + ], + }, + {"role": "user", "content": "plain string content, not a list"}, + ] + + assert _raw_retrieve_call_ids(messages) == frozenset({"oa1", "an1"}) + assert _raw_retrieve_call_ids("not a list") == frozenset() + assert _raw_retrieve_call_ids(None) == frozenset() + + @pytest.mark.asyncio async def test_nothing_compressible_returns_inputs_untouched(guardrail: HeadroomGuardrail): """A single-turn request is all protected, so there is nothing to send and From c83b4a1d1985e79253b7d32f1454cfc96f48e01c Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 2 Sep 2026 12:10:54 -0700 Subject: [PATCH 45/48] fix(proxy-extras): honor a raised command timeout for migrate deploy and name the right knob on db push timeouts --- .../litellm_proxy_extras/prisma_toolchain.py | 13 +++-- .../litellm_proxy_extras/utils.py | 5 +- .../test_prisma_toolchain.py | 51 +++++++++++++++++-- 3 files changed, 60 insertions(+), 9 deletions(-) diff --git a/litellm-proxy-extras/litellm_proxy_extras/prisma_toolchain.py b/litellm-proxy-extras/litellm_proxy_extras/prisma_toolchain.py index 5feb7a953b4..2283814ab35 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/prisma_toolchain.py +++ b/litellm-proxy-extras/litellm_proxy_extras/prisma_toolchain.py @@ -26,7 +26,10 @@ finish any sooner. Migrate deploy therefore runs under its own budget. All three budgets are overridable so an operator can widen them without a release: ``LITELLM_PRISMA_BOOTSTRAP_TIMEOUT`` for the toolchain install, ``LITELLM_PRISMA_MIGRATE_DEPLOY_TIMEOUT`` for ``prisma migrate deploy`` and -``LITELLM_PRISMA_COMMAND_TIMEOUT`` for every other Prisma command. +``LITELLM_PRISMA_COMMAND_TIMEOUT`` for every other Prisma command. The +per-command budget used to bound migrate deploy as well, so a deployment that +raised it above the deploy default keeps that larger budget for deploy unless +the deploy override says otherwise. """ import math @@ -102,9 +105,11 @@ def prisma_bootstrap_timeout() -> float: def prisma_migrate_deploy_timeout() -> float: """Seconds one ``prisma migrate deploy`` may run for, however many migrations are pending.""" - return _timeout_from_env( - PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR, DEFAULT_PRISMA_MIGRATE_DEPLOY_TIMEOUT - ) + if os.getenv(PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR) is not None: + return _timeout_from_env( + PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR, DEFAULT_PRISMA_MIGRATE_DEPLOY_TIMEOUT + ) + return max(DEFAULT_PRISMA_MIGRATE_DEPLOY_TIMEOUT, prisma_command_timeout()) def nodeenv_cache_dir() -> Optional[Path]: diff --git a/litellm-proxy-extras/litellm_proxy_extras/utils.py b/litellm-proxy-extras/litellm_proxy_extras/utils.py index ab9ec1e8a3a..f6647268624 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/utils.py +++ b/litellm-proxy-extras/litellm_proxy_extras/utils.py @@ -15,6 +15,7 @@ from litellm_proxy_extras.replica_identity import ( apply_replica_identity_full, ) from litellm_proxy_extras.prisma_toolchain import ( + PRISMA_COMMAND_TIMEOUT_ENV_VAR, PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR, ensure_prisma_toolchain, prisma_command_timeout, @@ -1135,9 +1136,9 @@ class ProxyExtrasDBManager: return True except subprocess.TimeoutExpired: logger.warning( - "Attempt %s timed out. Raise %s if this database needs longer to apply its pending migrations.", + "Attempt %s timed out. Raise %s if this database needs longer to apply its schema.", attempt + 1, - PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR, + PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR if use_migrate else PRISMA_COMMAND_TIMEOUT_ENV_VAR, ) time.sleep(random.randrange(5, 15)) except subprocess.CalledProcessError as e: diff --git a/tests/proxy_migration_tests/test_prisma_toolchain.py b/tests/proxy_migration_tests/test_prisma_toolchain.py index 1bd440d37b4..4c258c4d007 100644 --- a/tests/proxy_migration_tests/test_prisma_toolchain.py +++ b/tests/proxy_migration_tests/test_prisma_toolchain.py @@ -16,11 +16,13 @@ the proxy gave up after four identical timeouts. import ast import json +import logging import os import sys import time from collections.abc import Callable from pathlib import Path +from typing import Optional import pytest @@ -52,10 +54,10 @@ import time args = sys.argv[1:] cache_dir = os.environ["PRISMA_NODEENV_CACHE_DIR"] log_path = pathlib.Path(os.environ["FAKE_PRISMA_LOG"]) -earlier_deploys = sum( +earlier_same_command = sum( 1 for line in (log_path.read_text().splitlines() if log_path.exists() else []) - if json.loads(line)["args"][:2] == ["migrate", "deploy"] + if json.loads(line)["args"][:2] == args[:2] ) with log_path.open("a") as log: log.write( @@ -64,12 +66,14 @@ with log_path.open("a") as log: ) time.sleep(float(os.environ.get("FAKE_PRISMA_SLEEP", "0"))) if args[:2] == ["migrate", "deploy"]: - if earlier_deploys == 0: + if earlier_same_command == 0: time.sleep(float(os.environ.get("FAKE_PRISMA_FIRST_DEPLOY_SLEEP", "0"))) elif os.environ.get("FAKE_PRISMA_LATER_DEPLOY_STDERR"): print(os.environ["FAKE_PRISMA_LATER_DEPLOY_STDERR"], file=sys.stderr) sys.exit(1) print("No pending migrations to apply") +if args[:2] == ["db", "push"] and earlier_same_command == 0: + time.sleep(float(os.environ.get("FAKE_PRISMA_FIRST_PUSH_SLEEP", "0"))) sys.exit(0) """ @@ -269,6 +273,47 @@ def test_migrate_deploy_stops_at_its_own_timeout( assert elapsed < 30 +def test_db_push_timeout_hint_names_the_per_command_budget( + toolchain_env: tuple[Path, Path], monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture +) -> None: + """``db push`` keeps the per-command budget, so its timeout hint has to name that variable.""" + _, log_path = toolchain_env + monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@localhost:9/x") + monkeypatch.setenv(PRISMA_COMMAND_TIMEOUT_ENV_VAR, "1") + monkeypatch.setenv("FAKE_PRISMA_FIRST_PUSH_SLEEP", "3") + + with caplog.at_level(logging.WARNING, logger="litellm_proxy_extras"): + assert ProxyExtrasDBManager.setup_database(use_migrate=False, use_v2_resolver=False) is True + + assert [call["args"][:2] for call in _fake_prisma_calls(log_path)].count(["db", "push"]) == 2 + assert [record.getMessage() for record in caplog.records if "timed out" in record.getMessage()] == [ + f"Attempt 1 timed out. Raise {PRISMA_COMMAND_TIMEOUT_ENV_VAR} if this database needs longer to apply its schema." + ] + + +@pytest.mark.parametrize( + ("command_timeout", "deploy_timeout", "expected"), + [ + ("900", None, 900.0), + ("12", None, DEFAULT_PRISMA_MIGRATE_DEPLOY_TIMEOUT), + ("900", "1200", 1200.0), + ("900", "300", 300.0), + ], + ids=["raised_command_budget_carries_over", "lowered_command_budget_does_not", "override_wins_upward", "override_wins_downward"], +) +def test_migrate_deploy_budget_keeps_a_raised_command_budget( + command_timeout: str, deploy_timeout: Optional[str], expected: float, monkeypatch: pytest.MonkeyPatch +) -> None: + """Deployments that raised the per-command budget to survive a long deploy keep that budget for deploy.""" + monkeypatch.setenv(PRISMA_COMMAND_TIMEOUT_ENV_VAR, command_timeout) + if deploy_timeout is None: + monkeypatch.delenv(PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR, raising=False) + else: + monkeypatch.setenv(PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR, deploy_timeout) + + assert prisma_migrate_deploy_timeout() == expected + + @pytest.mark.parametrize( "raw", ["", "0", "-5", "not-a-number", "nan", "inf", "-inf", "1e400"], From 3af19cbf61706f1cf3d3360b8802b5cbc1991827 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 2 Sep 2026 12:36:29 -0700 Subject: [PATCH 46/48] test(proxy-extras): use the modern optional annotation in the deploy budget test --- tests/proxy_migration_tests/test_prisma_toolchain.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/tests/proxy_migration_tests/test_prisma_toolchain.py b/tests/proxy_migration_tests/test_prisma_toolchain.py index 4c258c4d007..733870f3239 100644 --- a/tests/proxy_migration_tests/test_prisma_toolchain.py +++ b/tests/proxy_migration_tests/test_prisma_toolchain.py @@ -22,7 +22,6 @@ import sys import time from collections.abc import Callable from pathlib import Path -from typing import Optional import pytest @@ -302,7 +301,7 @@ def test_db_push_timeout_hint_names_the_per_command_budget( ids=["raised_command_budget_carries_over", "lowered_command_budget_does_not", "override_wins_upward", "override_wins_downward"], ) def test_migrate_deploy_budget_keeps_a_raised_command_budget( - command_timeout: str, deploy_timeout: Optional[str], expected: float, monkeypatch: pytest.MonkeyPatch + command_timeout: str, deploy_timeout: str | None, expected: float, monkeypatch: pytest.MonkeyPatch ) -> None: """Deployments that raised the per-command budget to survive a long deploy keep that budget for deploy.""" monkeypatch.setenv(PRISMA_COMMAND_TIMEOUT_ENV_VAR, command_timeout) From cde9d94c3650dfbc0b704280973d6e4319a7530a Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 2 Sep 2026 12:40:24 -0700 Subject: [PATCH 47/48] feat(agentcore-a2a): derive runtime session id from A2A message.contextId (#39371) Native AgentCore A2A always sent either a fresh generated runtime session id or the single configured runtimeSessionId, so related turns lost context and unrelated callers shared one AgentCore microVM. The runtime session id is now params.message.contextId scoped to the calling key hash, then runtimeSessionId, then generated, and is length-validated (33-256) before the header is signed. Invalid ids surface as JSON-RPC -32602 / HTTP 400 instead of a 500. Co-authored-by: yassin Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../bedrock_agentcore/transformation.py | 51 ++++- litellm/a2a_protocol/utils.py | 25 +++ litellm/llms/langflow/a2a.py | 27 +-- .../proxy/agent_endpoints/a2a_endpoints.py | 2 + .../test_bedrock_agentcore_a2a.py | 191 ++++++++++++++++++ .../agent_endpoints/test_a2a_endpoints.py | 56 +++++ 6 files changed, 325 insertions(+), 27 deletions(-) diff --git a/litellm/a2a_protocol/providers/bedrock_agentcore/transformation.py b/litellm/a2a_protocol/providers/bedrock_agentcore/transformation.py index 32252711997..1e8cc4ff90e 100644 --- a/litellm/a2a_protocol/providers/bedrock_agentcore/transformation.py +++ b/litellm/a2a_protocol/providers/bedrock_agentcore/transformation.py @@ -10,8 +10,19 @@ from collections.abc import AsyncIterator, Mapping from typing import Any, Final from litellm._logging import verbose_logger +from litellm.a2a_protocol.litellm_completion_bridge.handler import ( + A2A_USER_API_KEY_HASH_PARAM, +) +from litellm.a2a_protocol.utils import ( + get_session_id_from_a2a_params, + scope_session_to_principal, +) +from litellm.exceptions import BadRequestError from litellm.llms.bedrock.chat.agentcore.transformation import AmazonAgentCoreConfig +RUNTIME_SESSION_ID_MIN_LENGTH: Final = 33 +RUNTIME_SESSION_ID_MAX_LENGTH: Final = 256 + # Reserved outbound header names that must never be sourced from per-request # ``agent_extra_headers`` for AgentCore requests. ``agent_extra_headers`` carries # values rewritten from the client-controlled ``x-a2a-{agent}-*`` convention, so @@ -19,8 +30,9 @@ from litellm.llms.bedrock.chat.agentcore.transformation import AmazonAgentCoreCo # request identity / SigV4 metadata by overwriting headers the proxy sets from # trusted server-side config. # -# The runtime headers (session / user id) are derived server-side from -# ``runtimeSessionId`` / ``runtimeUserId`` in the agent's ``litellm_params``; +# The runtime headers (session / user id) are derived server-side from the A2A +# ``message.contextId`` and ``runtimeSessionId`` / ``runtimeUserId`` in the +# agent's ``litellm_params``; # ``authorization`` is set by the AgentCore signer (JWT or SigV4); ``host`` and # the ``x-amz-*`` family are owned by SigV4 itself. _RESERVED_EXACT_HEADERS: Final = frozenset( @@ -66,6 +78,31 @@ def _filter_reserved_headers( return filtered or None +def _request_scoped_runtime_session_id( + params: Mapping[str, Any], + litellm_params: Mapping[str, Any], +) -> str | None: + context_id: Final = get_session_id_from_a2a_params(params) + if not isinstance(context_id, str) or not context_id: + return None + return scope_session_to_principal(context_id, litellm_params.get(A2A_USER_API_KEY_HASH_PARAM)) + + +def _validate_runtime_session_id(session_id: str, model: str) -> str: + if RUNTIME_SESSION_ID_MIN_LENGTH <= len(session_id) <= RUNTIME_SESSION_ID_MAX_LENGTH: + return session_id + raise BadRequestError( + message=( + f"Invalid AgentCore runtime session id {session_id!r}: AWS requires " + f"{RUNTIME_SESSION_ID_MIN_LENGTH}-{RUNTIME_SESSION_ID_MAX_LENGTH} characters. It is built from the A2A " + "message.contextId (prefixed with a 16-hex-char hash of the calling key and '-') when set, " + "otherwise from the agent's configured runtimeSessionId." + ), + model=model, + llm_provider="bedrock", + ) + + class BedrockAgentCoreA2ATransformation: """ Request/response transformation for Bedrock AgentCore A2A agents. @@ -100,7 +137,9 @@ class BedrockAgentCoreA2ATransformation: here to prevent a caller-controlled ``x-a2a-{agent}-*`` header from spoofing the AgentCore runtime user id or other SigV4 metadata. Use ``api_key`` / ``runtimeUserId`` / ``runtimeSessionId`` in litellm_params - (not ``agent_extra_headers``) to override those values. + (not ``agent_extra_headers``) to override those values. The runtime + session id is taken from ``params["message"]["contextId"]`` (scoped to + the calling key) when present, then ``runtimeSessionId``, else generated. Returns: Tuple of (url, signed_headers, signed_body_bytes) @@ -139,7 +178,11 @@ class BedrockAgentCoreA2ATransformation: # Set required AgentCore session headers (normally set by transform_request, # which we skip because it also builds {"prompt": "..."}) headers: Final[dict] = {} - session_id: Final = agentcore_config._get_runtime_session_id(optional_params) + session_id: Final = _validate_runtime_session_id( + _request_scoped_runtime_session_id(params, litellm_params) + or agentcore_config._get_runtime_session_id(optional_params), + model=model, + ) headers["X-Amzn-Bedrock-AgentCore-Runtime-Session-Id"] = session_id runtime_user_id: Final = agentcore_config._get_runtime_user_id(optional_params) if runtime_user_id: diff --git a/litellm/a2a_protocol/utils.py b/litellm/a2a_protocol/utils.py index f2e61f66105..7c459daf720 100644 --- a/litellm/a2a_protocol/utils.py +++ b/litellm/a2a_protocol/utils.py @@ -2,6 +2,8 @@ Utility functions for A2A protocol. """ +import hashlib +from collections.abc import Mapping from typing import TYPE_CHECKING, Any, Final import litellm @@ -140,6 +142,29 @@ class A2ARequestUtils: return prompt_tokens, completion_tokens, total_tokens +def get_session_id_from_a2a_params(params: Mapping[str, Any]) -> str | None: + message: Final = params.get("message", {}) + if isinstance(message, dict): + return message.get("contextId") + return getattr(message, "contextId", None) + + +def scope_session_to_principal(session_id: str, principal: str | None) -> str: + """ + Bind a client-supplied A2A contextId to the authenticated principal. + + Without this, two distinct keys authorized for the same agent could set the + same contextId and read/append to each other's backend memory. The + principal is hashed (it is already a hashed token) so the raw value is never + sent to the agent backend, while the original contextId is kept as a suffix + for operator-side correlation. + """ + if not principal: + return session_id + principal_prefix: Final = hashlib.sha256(principal.encode("utf-8")).hexdigest()[:16] + return f"{principal_prefix}-{session_id}" + + # Backwards compatibility aliases def extract_text_from_a2a_message(message: Any) -> str: return A2ARequestUtils.extract_text_from_message(message) diff --git a/litellm/llms/langflow/a2a.py b/litellm/llms/langflow/a2a.py index cae750d586e..060dc0a4d05 100644 --- a/litellm/llms/langflow/a2a.py +++ b/litellm/llms/langflow/a2a.py @@ -1,28 +1,9 @@ -import hashlib from typing import Any, Final - -def get_session_id_from_a2a_params(params: dict[str, Any]) -> str | None: - message: Final = params.get("message", {}) - if isinstance(message, dict): - return message.get("contextId") - return getattr(message, "contextId", None) - - -def scope_session_to_principal(session_id: str, principal: str | None) -> str: - """ - Bind a client-supplied A2A contextId to the authenticated principal. - - Without this, two distinct keys authorized for the same LangFlow agent could - set the same contextId and read/append to each other's LangFlow memory. The - principal is hashed (it is already a hashed token) so the raw value is never - sent to the LangFlow backend, while the original contextId is kept as a - suffix for operator-side correlation. - """ - if not principal: - return session_id - principal_prefix: Final = hashlib.sha256(principal.encode("utf-8")).hexdigest()[:16] - return f"{principal_prefix}-{session_id}" +from litellm.a2a_protocol.utils import ( + get_session_id_from_a2a_params, + scope_session_to_principal, +) def merge_a2a_session_into_litellm_params( diff --git a/litellm/proxy/agent_endpoints/a2a_endpoints.py b/litellm/proxy/agent_endpoints/a2a_endpoints.py index 31b05320cd3..28882484db4 100644 --- a/litellm/proxy/agent_endpoints/a2a_endpoints.py +++ b/litellm/proxy/agent_endpoints/a2a_endpoints.py @@ -1019,4 +1019,6 @@ async def invoke_agent_a2a( ) except Exception: pass + if isinstance(e, litellm.BadRequestError): + return _jsonrpc_error(body.get("id"), -32602, e.message, 400) return _jsonrpc_error(body.get("id"), -32603, f"Internal error: {e}", 500) diff --git a/tests/test_litellm/a2a_protocol/providers/bedrock_agentcore/test_bedrock_agentcore_a2a.py b/tests/test_litellm/a2a_protocol/providers/bedrock_agentcore/test_bedrock_agentcore_a2a.py index 5503a5668bf..a8fe464ec32 100644 --- a/tests/test_litellm/a2a_protocol/providers/bedrock_agentcore/test_bedrock_agentcore_a2a.py +++ b/tests/test_litellm/a2a_protocol/providers/bedrock_agentcore/test_bedrock_agentcore_a2a.py @@ -11,7 +11,9 @@ Verifies that: import json +import httpx import pytest +import respx from unittest.mock import AsyncMock, MagicMock, patch @@ -295,6 +297,195 @@ class TestTransformation: assert headers["Authorization"].startswith("AWS4-HMAC-SHA256") +SESSION_HEADER = "X-Amzn-Bedrock-AgentCore-Runtime-Session-Id" +CONTEXT_ID = "conversation-alpha-0001-0000000000000000" +KEY_HASH = "hashed-key-of-caller-one" + + +def _params_with_context(context_id: object) -> dict: + return {"message": {**SAMPLE_PARAMS["message"], "contextId": context_id}} + + +def _scoped(context_id: str, key_hash: str) -> str: + import hashlib + + return f"{hashlib.sha256(key_hash.encode()).hexdigest()[:16]}-{context_id}" + + +def _session_header(params: dict, litellm_params: dict) -> str: + from litellm.a2a_protocol.providers.bedrock_agentcore.transformation import ( + BedrockAgentCoreA2ATransformation, + ) + + _, headers, _ = BedrockAgentCoreA2ATransformation.get_url_and_signed_request( + request_id="req-001", + params=params, + litellm_params=litellm_params, + ) + return headers[SESSION_HEADER] + + +@pytest.fixture +def httpx_transport(monkeypatch): + import litellm + + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + litellm.in_memory_llm_clients_cache.flush_cache() + yield + litellm.in_memory_llm_clients_cache.flush_cache() + + +class TestRequestScopedRuntimeSession: + """message.contextId selects the AgentCore runtime session, scoped to the calling key.""" + + def test_context_id_scoped_to_calling_key(self): + from litellm.a2a_protocol.litellm_completion_bridge.handler import ( + A2A_USER_API_KEY_HASH_PARAM, + ) + + litellm_params = {**SAMPLE_LITELLM_PARAMS, A2A_USER_API_KEY_HASH_PARAM: KEY_HASH} + assert _session_header(_params_with_context(CONTEXT_ID), litellm_params) == _scoped(CONTEXT_ID, KEY_HASH) + + def test_context_id_used_verbatim_without_principal(self): + assert _session_header(_params_with_context(CONTEXT_ID), SAMPLE_LITELLM_PARAMS) == CONTEXT_ID + + def test_same_context_id_reuses_session_and_other_context_isolated(self): + first = _session_header(_params_with_context(CONTEXT_ID), SAMPLE_LITELLM_PARAMS) + second = _session_header(_params_with_context(CONTEXT_ID), SAMPLE_LITELLM_PARAMS) + other = _session_header( + _params_with_context("conversation-beta-00002-0000000000000000"), + SAMPLE_LITELLM_PARAMS, + ) + assert first == second + assert other != first + + def test_same_context_id_from_different_keys_is_isolated(self): + from litellm.a2a_protocol.litellm_completion_bridge.handler import ( + A2A_USER_API_KEY_HASH_PARAM, + ) + + params = _params_with_context(CONTEXT_ID) + caller_one = _session_header(params, {**SAMPLE_LITELLM_PARAMS, A2A_USER_API_KEY_HASH_PARAM: KEY_HASH}) + caller_two = _session_header( + params, {**SAMPLE_LITELLM_PARAMS, A2A_USER_API_KEY_HASH_PARAM: "hashed-key-of-caller-two"} + ) + assert caller_one != caller_two + assert caller_one.endswith(f"-{CONTEXT_ID}") + assert caller_two.endswith(f"-{CONTEXT_ID}") + + def test_context_id_takes_precedence_over_configured_session(self): + litellm_params = {**SAMPLE_LITELLM_PARAMS, "runtimeSessionId": "a" * 40} + assert _session_header(_params_with_context(CONTEXT_ID), litellm_params) == CONTEXT_ID + + def test_configured_session_is_fallback_without_context_id(self): + litellm_params = {**SAMPLE_LITELLM_PARAMS, "runtimeSessionId": "a" * 40} + assert _session_header(SAMPLE_PARAMS, litellm_params) == "a" * 40 + assert _session_header(_params_with_context(""), litellm_params) == "a" * 40 + + def test_no_context_id_and_no_config_generates_new_session_per_request(self): + first = _session_header(SAMPLE_PARAMS, SAMPLE_LITELLM_PARAMS) + second = _session_header(SAMPLE_PARAMS, SAMPLE_LITELLM_PARAMS) + assert first != second + assert 33 <= len(first) <= 256 + + @pytest.mark.parametrize( + "context_id", + [ + "short-context-id", + "x" * 257, + ], + ) + def test_invalid_context_id_rejected_with_clear_error(self, context_id): + import litellm + + with pytest.raises(litellm.BadRequestError, match="Invalid AgentCore runtime session id") as exc_info: + _session_header(_params_with_context(context_id), SAMPLE_LITELLM_PARAMS) + assert exc_info.value.status_code == 400 + assert "33-256" in str(exc_info.value) + + def test_scoped_context_id_shorter_than_33_rejected(self): + import litellm + from litellm.a2a_protocol.litellm_completion_bridge.handler import ( + A2A_USER_API_KEY_HASH_PARAM, + ) + + litellm_params = {**SAMPLE_LITELLM_PARAMS, A2A_USER_API_KEY_HASH_PARAM: KEY_HASH} + with pytest.raises(litellm.BadRequestError, match=_scoped("c" * 15, KEY_HASH)): + _session_header(_params_with_context("c" * 15), litellm_params) + assert _session_header(_params_with_context("c" * 16), litellm_params) == _scoped("c" * 16, KEY_HASH) + + def test_invalid_configured_session_rejected(self): + import litellm + + litellm_params = {**SAMPLE_LITELLM_PARAMS, "runtimeSessionId": "too-short"} + with pytest.raises(litellm.BadRequestError, match="Invalid AgentCore runtime session id"): + _session_header(SAMPLE_PARAMS, litellm_params) + + def test_non_string_context_id_falls_back(self): + litellm_params = {**SAMPLE_LITELLM_PARAMS, "runtimeSessionId": "a" * 40} + assert _session_header(_params_with_context(12345), litellm_params) == "a" * 40 + + def test_spoofed_session_header_does_not_override_context_id(self): + from litellm.a2a_protocol.providers.bedrock_agentcore.transformation import ( + BedrockAgentCoreA2ATransformation, + ) + + _, headers, _ = BedrockAgentCoreA2ATransformation.get_url_and_signed_request( + request_id="req-001", + params=_params_with_context(CONTEXT_ID), + litellm_params=SAMPLE_LITELLM_PARAMS, + agent_extra_headers={SESSION_HEADER: "s" * 40}, + ) + assert headers[SESSION_HEADER] == CONTEXT_ID + + @pytest.mark.asyncio + async def test_context_id_session_header_on_outbound_non_streaming_post(self, httpx_transport): + from litellm.a2a_protocol.litellm_completion_bridge.handler import ( + A2A_USER_API_KEY_HASH_PARAM, + ) + from litellm.a2a_protocol.providers.bedrock_agentcore.config import ( + BedrockAgentCoreA2AConfig, + ) + + with respx.mock(assert_all_called=True) as router: + route = router.post(url__regex=r".*/invocations.*").mock( + return_value=httpx.Response(200, json={"jsonrpc": "2.0", "id": "req-001", "result": {}}) + ) + await BedrockAgentCoreA2AConfig().handle_non_streaming( + request_id="req-001", + params=_params_with_context(CONTEXT_ID), + litellm_params={**SAMPLE_LITELLM_PARAMS, A2A_USER_API_KEY_HASH_PARAM: KEY_HASH}, + ) + + assert route.calls.last.request.headers[SESSION_HEADER] == _scoped(CONTEXT_ID, KEY_HASH) + + @pytest.mark.asyncio + async def test_context_id_session_header_on_outbound_streaming_post(self, httpx_transport): + from litellm.a2a_protocol.litellm_completion_bridge.handler import ( + A2A_USER_API_KEY_HASH_PARAM, + ) + from litellm.a2a_protocol.providers.bedrock_agentcore.config import ( + BedrockAgentCoreA2AConfig, + ) + + with respx.mock(assert_all_called=True) as router: + route = router.post(url__regex=r".*/invocations.*").mock( + return_value=httpx.Response(200, json={"jsonrpc": "2.0", "id": "req-001", "result": {}}) + ) + events = [ + event + async for event in BedrockAgentCoreA2AConfig().handle_streaming( + request_id="req-001", + params=_params_with_context(CONTEXT_ID), + litellm_params={**SAMPLE_LITELLM_PARAMS, A2A_USER_API_KEY_HASH_PARAM: KEY_HASH}, + ) + ] + + assert events == [{"jsonrpc": "2.0", "id": "req-001", "result": {}}] + assert route.calls.last.request.headers[SESSION_HEADER] == _scoped(CONTEXT_ID, KEY_HASH) + + class TestNonStreaming: """Test end-to-end non-streaming flow.""" diff --git a/tests/test_litellm/proxy/agent_endpoints/test_a2a_endpoints.py b/tests/test_litellm/proxy/agent_endpoints/test_a2a_endpoints.py index 2ff38af80b1..43034f889f6 100644 --- a/tests/test_litellm/proxy/agent_endpoints/test_a2a_endpoints.py +++ b/tests/test_litellm/proxy/agent_endpoints/test_a2a_endpoints.py @@ -918,6 +918,62 @@ async def test_task_method_failure_hook_uses_enriched_request_data(): assert failure_data.get("agent_id") == "test-agent" +@pytest.mark.asyncio +async def test_agentcore_invalid_context_id_returns_jsonrpc_invalid_params_400(): + from litellm.proxy._types import UserAPIKeyAuth + + agent = _make_agent_mock() + agent.litellm_params = { + "custom_llm_provider": "bedrock", + "model": "bedrock/agentcore/arn:aws:bedrock-agentcore:us-west-2:123456789012:runtime/demo", + "api_key": "test-jwt-token", + } + mock_request = _make_request_mock( + "message/send", + { + "message": { + "role": "user", + "parts": [{"kind": "text", "text": "Hello"}], + "messageId": "msg-1", + "contextId": "too-short", + } + }, + ) + user_api_key_dict = UserAPIKeyAuth(api_key="sk-test", user_id="u1", team_id="t1") + + mock_proxy_logging = MagicMock() + mock_proxy_logging.pre_call_hook = AsyncMock( + side_effect=lambda user_api_key_dict, data, call_type: data + ) + mock_proxy_logging.post_call_failure_hook = AsyncMock(return_value=None) + + with ExitStack() as stack: + for p in _base_patches(agent): + stack.enter_context(p) + stack.enter_context( + patch( # test-quality-ok: same proxy_logging_obj injection the sibling failure-hook test uses; no HTTP call is made because the request is rejected before signing + "litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging + ) + ) + + from litellm.proxy.agent_endpoints.a2a_endpoints import invoke_agent_a2a + + response = await invoke_agent_a2a( + agent_id="test-agent", + request=mock_request, + fastapi_response=MagicMock(), + user_api_key_dict=user_api_key_dict, + ) + + body = json.loads(response.body.decode()) + assert response.status_code == 400 + assert body["id"] == "req-1" + assert body["error"]["code"] == -32602 + assert "Invalid AgentCore runtime session id" in body["error"]["message"] + assert "Internal error" not in body["error"]["message"] + mock_proxy_logging.post_call_failure_hook.assert_awaited_once() + + @pytest.mark.asyncio async def test_get_extended_agent_card_rewrites_url(): from litellm.proxy._types import UserAPIKeyAuth From 987ab769213527d9bd1cfaeb3b038674dbec40f4 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 2 Sep 2026 12:40:59 -0700 Subject: [PATCH 48/48] fix(proxy): share per-model budget counters across replicas through the spend counter cache (#39375) * fix(proxy): share per-model budget counters across replicas through the spend counter cache Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(proxy): keep the shared fake Redis store immutable Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: yassin Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../proxy/hooks/model_max_budget_limiter.py | 35 +++++--- litellm/proxy/proxy_server.py | 2 +- ...test_unit_test_max_model_budget_limiter.py | 85 +++++++++++++++++++ .../proxy/test_redis_auth_cache_flag.py | 24 +++++- 4 files changed, 132 insertions(+), 14 deletions(-) diff --git a/litellm/proxy/hooks/model_max_budget_limiter.py b/litellm/proxy/hooks/model_max_budget_limiter.py index c5d10b2749b..efaaab277a9 100644 --- a/litellm/proxy/hooks/model_max_budget_limiter.py +++ b/litellm/proxy/hooks/model_max_budget_limiter.py @@ -1,6 +1,6 @@ import json import time -from collections.abc import Iterable, Mapping +from collections.abc import Iterable, Mapping, Sequence from dataclasses import dataclass from types import MappingProxyType from typing import Final @@ -199,18 +199,10 @@ async def build_model_max_budget_usage( ) for budget_model, budget_config in budgets ) - batched: Final = await cache.async_batch_get_cache( - keys=list(spend_keys) # mutable-ok: async_batch_get_cache annotates keys as list, so one must exist here - ) - # async_batch_get_cache returns None if it fails internally, and its result is - # index-aligned with `keys` otherwise. An unusable result reads as a miss, - # which is what a never-written counter already reads as. - current_spends: Final = ( - tuple(batched) if isinstance(batched, list) and len(batched) == len(budgets) else (None,) * len(budgets) - ) + current_spends: Final = await _current_window_spends(cache=cache, spend_keys=spend_keys) return { budget_model: { - "current_spend": round(_as_spend(current_spend), 4), + "current_spend": round(current_spend, 4), "budget_limit": budget_config.max_budget, "time_period": budget_config.budget_duration, } @@ -218,6 +210,22 @@ async def build_model_max_budget_usage( } +async def _current_window_spends(cache: DualCache, spend_keys: Sequence[str]) -> tuple[float, ...]: + """Redis holds the window total across replicas; the in-memory copy is one replica's share.""" + keys: Final = list(spend_keys) # mutable-ok: both batch readers annotate their key argument as list + redis_cache: Final = cache.redis_cache + if redis_cache is not None: + shared: Final = await redis_cache.async_batch_get_cache(key_list=keys) + return tuple(_as_spend(shared.get(key)) for key in keys) + # async_batch_get_cache returns None if it fails internally, and its result is + # index-aligned with `keys` otherwise. An unusable result reads as a miss, + # which is what a never-written counter already reads as. + batched: Final = await cache.async_batch_get_cache(keys=keys) + if not isinstance(batched, list) or len(batched) != len(keys): + return (0.0,) * len(keys) + return tuple(_as_spend(current_spend) for current_spend in batched) + + def _usable_budget_config(raw_budget_config: object) -> BudgetConfig | None: try: budget_config: Final = BudgetConfig.model_validate(raw_budget_config) @@ -404,7 +412,10 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): return current_spend + _as_spend(await self._cached_spend(legacy_spend_key)) async def _cached_spend(self, spend_key: str) -> float | None: - return await self.dual_cache.async_get_cache(key=spend_key) + redis_cache: Final = self.dual_cache.redis_cache + if redis_cache is None: + return await self.dual_cache.async_get_cache(key=spend_key) + return await redis_cache.async_get_cache(key=spend_key) async def async_filter_deployments( self, diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 85a57e5af2b..d52f05f6166 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -2280,7 +2280,7 @@ user_api_key_cache: UserApiKeyCache = UserApiKeyCache( ) spend_counter_cache: Final = DualCache(default_in_memory_ttl=UserAPIKeyCacheTTLEnum.in_memory_cache_ttl.value) cli_sso_session_cache: Final = DualCache(default_in_memory_ttl=CLI_SSO_SESSION_TTL_SECONDS) -model_max_budget_limiter: Final = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=user_api_key_cache) +model_max_budget_limiter: Final = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=spend_counter_cache) litellm.logging_callback_manager.add_litellm_callback(model_max_budget_limiter) redis_usage_cache: RedisCache | None = None # redis cache used for tracking spend, tpm/rpm limits polling_via_cache_enabled: Literal["all"] | list[str] | bool = False diff --git a/tests/proxy_unit_tests/test_unit_test_max_model_budget_limiter.py b/tests/proxy_unit_tests/test_unit_test_max_model_budget_limiter.py index 3785ccdcfba..096efc33aaf 100644 --- a/tests/proxy_unit_tests/test_unit_test_max_model_budget_limiter.py +++ b/tests/proxy_unit_tests/test_unit_test_max_model_budget_limiter.py @@ -1,3 +1,5 @@ +import asyncio +from types import MappingProxyType from unittest.mock import AsyncMock, patch @@ -5,6 +7,7 @@ import pytest import litellm from litellm.caching.caching import DualCache +from litellm.caching.redis_cache import RedisCache from datetime import datetime, timezone from litellm.litellm_core_utils.duration_parser import duration_in_seconds @@ -1332,3 +1335,85 @@ async def test_the_user_scope_has_no_pre_upgrade_counter_to_carry(): await limiter.is_user_within_model_budget( user_id="u1", user_model_max_budget=model_max_budget, model="openai/gpt-4" ) + + +class _SharedFakeRedis(RedisCache): + """Stand-in for the one Redis every replica's DualCache is attached to. + + Only the methods the limiter and DualCache call are implemented, and + ``super().__init__`` is skipped so no connection is opened. + """ + + def __init__(self): + self._store = MappingProxyType({}) + + async def async_set_cache(self, key, value, **kwargs): + self._store = MappingProxyType({**self._store, key: value}) + + async def async_get_cache(self, key, **kwargs): + return self._store.get(key) + + async def async_batch_get_cache(self, key_list, **kwargs): + return {key: self._store.get(key) for key in key_list} + + async def async_increment_pipeline(self, increment_list, **kwargs): + for op in increment_list: + total = self._store.get(op["key"], 0.0) + op["increment_value"] + self._store = MappingProxyType({**self._store, op["key"]: total}) + return [self._store[op["key"]] for op in increment_list] + + +async def _log_spend(limiter, *, key_hash, model_max_budget, response_cost): + await limiter.async_log_success_event( + _success_kwargs( + model_group="gpt-4", + response_cost=response_cost, + key_hash=key_hash, + key_model_max_budget=model_max_budget, + ), + response_obj=None, + start_time=None, + end_time=None, + ) + # The Redis push is scheduled as a task rather than awaited inline. + await asyncio.gather(*(t for t in asyncio.all_tasks() if t is not asyncio.current_task())) + + +@pytest.mark.asyncio +async def test_spend_logged_on_one_replica_is_enforced_and_reported_on_another(): + """ + Each replica increments its own in-memory copy of the per-model counter and + pushes the increment to the shared Redis, so only Redis holds the window's + total. A replica that has served part of the traffic must still enforce and + report the total, not its own share. + + Regression: reads went to the in-memory tier first, so a replica whose local + copy sat under the cap kept admitting requests and /key/info on it reported + that local share, while the shared counter was already over the cap. + """ + shared_redis = _SharedFakeRedis() + replica_a = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=DualCache(redis_cache=shared_redis)) + replica_b = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=DualCache(redis_cache=shared_redis)) + key_hash = "vk-shared" + model_max_budget = {"gpt-4": {"budget_limit": 1.0, "time_period": "30d"}} + user_api_key = UserAPIKeyAuth(token=key_hash, model_max_budget=model_max_budget) + + await _log_spend(replica_b, key_hash=key_hash, model_max_budget=model_max_budget, response_cost=0.25) + await _log_spend(replica_a, key_hash=key_hash, model_max_budget=model_max_budget, response_cost=0.5) + await _log_spend(replica_a, key_hash=key_hash, model_max_budget=model_max_budget, response_cost=0.5) + + with pytest.raises(litellm.BudgetExceededError): + await replica_b.is_key_within_model_budget(user_api_key, "gpt-4") + + usage_on_b = await build_model_max_budget_usage( + entity_type=Litellm_EntityType.KEY, + entity_id=key_hash, + model_max_budget=model_max_budget, + cache=replica_b.dual_cache, + ) + assert usage_on_b["gpt-4"]["current_spend"] == 1.25 + + # Control: a replica that never served this key reads the same total. + replica_c = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=DualCache(redis_cache=shared_redis)) + with pytest.raises(litellm.BudgetExceededError): + await replica_c.is_key_within_model_budget(user_api_key, "gpt-4") diff --git a/tests/test_litellm/proxy/test_redis_auth_cache_flag.py b/tests/test_litellm/proxy/test_redis_auth_cache_flag.py index 772b08bc9d0..573bfc40c96 100644 --- a/tests/test_litellm/proxy/test_redis_auth_cache_flag.py +++ b/tests/test_litellm/proxy/test_redis_auth_cache_flag.py @@ -5,7 +5,7 @@ Verifies that _init_cache attaches Redis to user_api_key_cache only when the flag is explicitly set to True, and leaves it in-memory-only otherwise. """ -from contextlib import contextmanager +from contextlib import ExitStack, contextmanager import json from unittest.mock import MagicMock, patch @@ -167,3 +167,25 @@ class TestRedisAuthCacheFlag: f"cli_sso_session_cache must always get Redis " f"(enable_redis_auth_cache={flag_value!r})" ) + + def test_flag_absent_still_shares_the_model_budget_counters_over_redis(self): + """ + Per-model budget counters are spend counters: the limiter must be able to + push and read them through Redis without the auth-cache opt-in, or every + worker enforces and reports its own share of a key's spend + """ + fake_redis = _FakeRedisCache() + limiter_cache = ps.model_max_budget_limiter.dual_cache + touched_caches = ( + limiter_cache, + ps.spend_counter_cache, + ps.cli_sso_session_cache, + ps.user_api_key_cache, + ps.litellm_config_cache, + ) + with ExitStack() as detached: + for cache in touched_caches: + detached.enter_context(patch.object(cache, "redis_cache", None)) + ps._attach_redis_usage_cache(fake_redis, enable_redis_auth_cache=False) + assert limiter_cache.redis_cache is fake_redis + assert ps.user_api_key_cache.redis_cache is None