From a1514efa210c60c00809b21d2906503b0c452cc8 Mon Sep 17 00:00:00 2001 From: michelligabriele Date: Mon, 27 Jul 2026 17:07:58 +0200 Subject: [PATCH 01/97] 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 c24927cf2ac363af4c55cf701181732f1849c8f7 Mon Sep 17 00:00:00 2001 From: tin Date: Tue, 4 Aug 2026 20:31:55 +0000 Subject: [PATCH 02/97] fix(proxy): report requested model on Anthropic streaming message_start Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../streaming_model_restamp.py | 79 +++++++++++++ litellm/proxy/common_request_processing.py | 25 +++- .../test_streaming_model_restamp.py | 109 ++++++++++++++++++ 3 files changed, 212 insertions(+), 1 deletion(-) create mode 100644 litellm/proxy/anthropic_endpoints/streaming_model_restamp.py create mode 100644 tests/test_litellm/proxy/anthropic_endpoints/test_streaming_model_restamp.py diff --git a/litellm/proxy/anthropic_endpoints/streaming_model_restamp.py b/litellm/proxy/anthropic_endpoints/streaming_model_restamp.py new file mode 100644 index 00000000000..857b6abd065 --- /dev/null +++ b/litellm/proxy/anthropic_endpoints/streaming_model_restamp.py @@ -0,0 +1,79 @@ +""" +Restamp the public ``model`` on the Anthropic Messages ``message_start`` event, the only +stream event carrying a model, so streamed responses report the requested model like +non-streaming ones do. + +Chunks reach the serializer either as already-encoded SSE frames (``bytes``/``str``, the +provider passthrough path) or as event dicts (fake-stream and agentic paths). +""" + +import json + +from pydantic import TypeAdapter, ValidationError + +_MESSAGE_START_EVENT = "message_start" +_SSE_DATA_FIELD = "data:" + +_EVENT_ADAPTER: TypeAdapter[dict[str, object]] = TypeAdapter(dict[str, object]) + + +def _restamped_event(event: dict[str, object], requested_model: str) -> dict[str, object] | None: + message = event.get("message") + if event.get("type") != _MESSAGE_START_EVENT or not isinstance(message, dict): + return None + if message.get("model") == requested_model: + return None + return {**event, "message": {**message, "model": requested_model}} # mutable-ok: SSE payload, re-serialized as is + + +def _restamped_data_line(line: str, requested_model: str) -> str | None: + stripped = line.strip() + if not stripped.startswith(_SSE_DATA_FIELD): + return None + payload = stripped[len(_SSE_DATA_FIELD) :].strip() + if not payload or payload == "[DONE]": + return None + try: + event = _EVENT_ADAPTER.validate_json(payload) + except ValidationError: + return None + restamped = _restamped_event(event, requested_model) + if restamped is None: + return None + return f"data: {json.dumps(restamped, separators=(',', ':'))}" + + +def _restamped_frame(frame: str, requested_model: str) -> str | None: + lines = frame.split("\n") + restamped = tuple(_restamped_data_line(line, requested_model) for line in lines) + if all(line is None for line in restamped): + return None + return "\n".join(new if new is not None else old for new, old in zip(restamped, lines)) + + +def restamp_anthropic_stream_chunk_model(chunk: object, requested_model: str) -> object: + """ + Return ``chunk`` with the ``message_start`` model replaced by ``requested_model``. + + Chunks that carry no model are returned unchanged. + """ + if isinstance(chunk, dict): + try: + event = _EVENT_ADAPTER.validate_python(chunk) + except ValidationError: + return chunk + return _restamped_event(event, requested_model) or chunk + + if isinstance(chunk, (bytes, bytearray)): + if _MESSAGE_START_EVENT.encode() not in chunk: + return chunk + restamped = _restamped_frame(chunk.decode("utf-8", errors="ignore"), requested_model) + return chunk if restamped is None else restamped.encode("utf-8") + + if isinstance(chunk, str): + if _MESSAGE_START_EVENT not in chunk: + return chunk + restamped = _restamped_frame(chunk, requested_model) + return chunk if restamped is None else restamped + + return chunk diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index f9cad283166..fd5de0debec 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -70,6 +70,9 @@ if TYPE_CHECKING: ProxyConfig = _ProxyConfig else: ProxyConfig = Any +from litellm.proxy.anthropic_endpoints.streaming_model_restamp import ( + restamp_anthropic_stream_chunk_model, +) from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request from litellm.types.utils import ( ModelResponse, @@ -1953,6 +1956,9 @@ class ProxyBaseLLMRequestProcessing: request_data=self.data, proxy_logging_obj=proxy_logging_obj, request=request, + restamp_model=( + None if _should_return_raw_model_name(self.data) else requested_model_from_client + ), ) return await create_response( generator=selected_data_generator, @@ -2801,6 +2807,18 @@ class ProxyBaseLLMRequestProcessing: else: return chunk + @staticmethod + def _sse_chunk_serializer(restamp_model: str | None) -> StreamChunkSerializer: + if not restamp_model: + return ProxyBaseLLMRequestProcessing.return_sse_chunk + + def serialize(chunk: object) -> str: + return ProxyBaseLLMRequestProcessing.return_sse_chunk( + restamp_anthropic_stream_chunk_model(chunk, restamp_model) + ) + + return serialize + @staticmethod async def _finalize_streaming_generator_cleanup( request: Request | None, @@ -2990,6 +3008,7 @@ class ProxyBaseLLMRequestProcessing: request_data: dict, proxy_logging_obj: ProxyLogging, request: Request | None = None, + restamp_model: str | None = None, ) -> AsyncGenerator[str, None]: """ Anthropic /messages and Google /generateContent streaming data generator require SSE events. @@ -2998,13 +3017,17 @@ class ProxyBaseLLMRequestProcessing: SSE serializers directly (rather than re-wrapping it in another ``async for: yield`` trampoline), so a streamed chunk traverses one fewer async-generator layer / coroutine resume on the hot path. + + ``restamp_model`` publishes that name on the Anthropic ``message_start`` + event in place of the provider's model, matching what the non-streaming + response reports. """ return ProxyBaseLLMRequestProcessing.async_streaming_data_generator( response=response, user_api_key_dict=user_api_key_dict, request_data=request_data, proxy_logging_obj=proxy_logging_obj, - serialize_chunk=ProxyBaseLLMRequestProcessing.return_sse_chunk, + serialize_chunk=ProxyBaseLLMRequestProcessing._sse_chunk_serializer(restamp_model), serialize_error=lambda proxy_exc: ( f"{STREAM_SSE_DATA_PREFIX}{json.dumps({'error': proxy_exc.to_dict()})}\n\n" ), diff --git a/tests/test_litellm/proxy/anthropic_endpoints/test_streaming_model_restamp.py b/tests/test_litellm/proxy/anthropic_endpoints/test_streaming_model_restamp.py new file mode 100644 index 00000000000..6e2c3f49445 --- /dev/null +++ b/tests/test_litellm/proxy/anthropic_endpoints/test_streaming_model_restamp.py @@ -0,0 +1,109 @@ +""" +Tests for restamping the public model on Anthropic Messages streaming chunks. +""" + +import json +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from litellm.proxy.anthropic_endpoints.streaming_model_restamp import ( + restamp_anthropic_stream_chunk_model, +) +from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + + +def _message_start_frame(model: str) -> bytes: + payload = { + "type": "message_start", + "message": {"id": "msg_1", "type": "message", "role": "assistant", "model": model, "content": []}, + } + return f"event: message_start\ndata: {json.dumps(payload)}\n\n".encode() + + +def _proxy_logging_obj_streaming(frames: list[bytes]) -> MagicMock: + async def _iterator_hook(**_kwargs): + for frame in frames: + yield frame + + proxy_logging_obj = MagicMock() + proxy_logging_obj.async_post_call_streaming_iterator_hook = _iterator_hook + proxy_logging_obj.async_post_call_streaming_hook = AsyncMock(side_effect=lambda **kwargs: kwargs["response"]) + return proxy_logging_obj + + +def _model_from_frame(frame: bytes | str) -> str: + text = frame.decode("utf-8") if isinstance(frame, bytes) else frame + data_line = next(line for line in text.split("\n") if line.startswith("data:")) + return json.loads(data_line[len("data:") :])["message"]["model"] + + +def test_restamps_sse_bytes_frame(): + restamped = restamp_anthropic_stream_chunk_model( + _message_start_frame("claude-haiku-4-5-20251001"), "claude-auto-1" + ) + + assert isinstance(restamped, bytes) + assert _model_from_frame(restamped) == "claude-auto-1" + assert b"event: message_start" in restamped + + +def test_restamps_event_dict(): + chunk = {"type": "message_start", "message": {"id": "msg_1", "model": "claude-sonnet-4-6"}} + + restamped = restamp_anthropic_stream_chunk_model(chunk, "claude-auto-2") + + assert restamped == {"type": "message_start", "message": {"id": "msg_1", "model": "claude-auto-2"}} + assert chunk["message"]["model"] == "claude-sonnet-4-6" + + +@pytest.mark.parametrize( + "chunk", + [ + b'event: content_block_delta\ndata: {"type":"content_block_delta","delta":{"text":"hi"}}\n\n', + {"type": "content_block_delta", "delta": {"text": "hi"}}, + {"type": "message_start", "message": "not-a-dict"}, + b"event: message_start\ndata: not-json\n\n", + b"data: [DONE]\n\n", + ], +) +def test_leaves_chunks_without_a_model_untouched(chunk): + assert restamp_anthropic_stream_chunk_model(chunk, "claude-auto-1") == chunk + + +@pytest.mark.asyncio +async def test_sse_generator_publishes_requested_model_on_message_start(): + """The message_start event reports the requested model, not the provider's.""" + delta_frame = b'event: content_block_delta\ndata: {"type":"content_block_delta","delta":{"text":"hi"}}\n\n' + proxy_logging_obj = _proxy_logging_obj_streaming([_message_start_frame("claude-haiku-4-5-20251001"), delta_frame]) + + chunks = [ + chunk + async for chunk in ProxyBaseLLMRequestProcessing.async_sse_data_generator( + response=MagicMock(), + user_api_key_dict=MagicMock(), + request_data={"model": "claude-auto-1"}, + proxy_logging_obj=proxy_logging_obj, + restamp_model="claude-auto-1", + ) + ] + + assert _model_from_frame(chunks[0]) == "claude-auto-1" + assert chunks[1] == delta_frame + + +@pytest.mark.asyncio +async def test_sse_generator_keeps_provider_model_when_restamping_is_off(): + proxy_logging_obj = _proxy_logging_obj_streaming([_message_start_frame("claude-haiku-4-5-20251001")]) + + chunks = [ + chunk + async for chunk in ProxyBaseLLMRequestProcessing.async_sse_data_generator( + response=MagicMock(), + user_api_key_dict=MagicMock(), + request_data={"model": "claude-auto-1"}, + proxy_logging_obj=proxy_logging_obj, + ) + ] + + assert _model_from_frame(chunks[0]) == "claude-haiku-4-5-20251001" From c02c81452c932eff275755a43ff9c686cd8055d4 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 31 Aug 2026 13:16:48 -0700 Subject: [PATCH 03/97] fix(proxy): reassemble split SSE frames before restamping anthropic message_start --- .../streaming_model_restamp.py | 78 +++++++++++++++++ litellm/proxy/common_request_processing.py | 8 +- .../test_streaming_model_restamp.py | 83 +++++++++++++++++++ 3 files changed, 165 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/anthropic_endpoints/streaming_model_restamp.py b/litellm/proxy/anthropic_endpoints/streaming_model_restamp.py index bbb7f2ceaa3..e8d54f03949 100644 --- a/litellm/proxy/anthropic_endpoints/streaming_model_restamp.py +++ b/litellm/proxy/anthropic_endpoints/streaming_model_restamp.py @@ -14,7 +14,11 @@ from typing import Final from pydantic import TypeAdapter, ValidationError _MESSAGE_START_EVENT: Final = "message_start" +_MESSAGE_START_MARKER: Final = b"message_start" _SSE_DATA_FIELD: Final = "data:" +_SSE_FRAME_END: Final = b"\n\n" +_MAX_HELD_BYTES: Final = 65536 +_PING_MARKERS: Final = (b"event: ping", b'"type": "ping"', b'"type":"ping"') _EVENT_ADAPTER: Final = TypeAdapter(Mapping[str, object]) @@ -79,3 +83,77 @@ def restamp_anthropic_stream_chunk_model(chunk: object, requested_model: str) -> return chunk if restamped_text is None else restamped_text return chunk + + +def _is_ping_frame(frame: bytes) -> bool: + return any(marker in frame for marker in _PING_MARKERS) + + +class AnthropicStreamModelRestamper: + """ + Per-stream restamper for the encoded passthrough path, where chunks are raw + transport reads: the ``message_start`` SSE frame can arrive split across + chunks or coalesced with later frames. Complete frames are emitted as their + terminator closes them and an incomplete tail is held until it completes, + so the restamp never misses a torn frame. Once ``message_start`` has been + handled, or the first real event proves the stream carries none, every + later chunk passes through untouched. + """ + + def __init__(self, requested_model: str) -> None: + self._requested_model: Final = requested_model + self._held = b"" + self._armed = True + + def process(self, chunk: object) -> object: + if not self._armed: + return chunk + if isinstance(chunk, (bytes, bytearray)): + return self._process_encoded(bytes(chunk)) + if isinstance(chunk, str): + return self._process_encoded(chunk.encode("utf-8")) + restamped: Final = restamp_anthropic_stream_chunk_model(chunk, self._requested_model) + if isinstance(chunk, dict) and chunk.get("type") not in (None, "ping"): + self._armed = False + return restamped + + def _process_encoded(self, data: bytes) -> bytes: + combined: Final = self._held + data + if _SSE_FRAME_END not in combined: + if len(combined) > _MAX_HELD_BYTES: + self._held = b"" + self._armed = False + return combined + self._held = combined + return b"" + closed, _, tail = combined.rpartition(_SSE_FRAME_END) + emitted: Final = self._restamped_closed_block(closed + _SSE_FRAME_END) + if not self._armed: + self._held = b"" + return emitted + tail + self._held = tail + return emitted + + def _restamped_closed_block(self, closed: bytes) -> bytes: + frames: Final = tuple(closed.split(_SSE_FRAME_END)[:-1]) + decider: Final = next( + ( + index + for index, frame in enumerate(frames) + if _MESSAGE_START_MARKER in frame or (b"data:" in frame and not _is_ping_frame(frame)) + ), + None, + ) + if decider is None: + return closed + self._armed = False + decider_frame: Final = frames[decider] + _SSE_FRAME_END + if _MESSAGE_START_MARKER not in decider_frame: + return closed + restamped_text: Final = _restamped_frame(decider_frame.decode("utf-8", errors="ignore"), self._requested_model) + if restamped_text is None: + return closed + return b"".join( + restamped_text.encode("utf-8") if index == decider else frame + _SSE_FRAME_END + for index, frame in enumerate(frames) + ) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 23e887f34a3..989cc7c18fb 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -177,7 +177,7 @@ if TYPE_CHECKING: else: ProxyConfig = Any from litellm.proxy.anthropic_endpoints.streaming_model_restamp import ( - restamp_anthropic_stream_chunk_model, + AnthropicStreamModelRestamper, ) from litellm.proxy.litellm_pre_call_utils import ( add_litellm_data_to_request, @@ -3414,10 +3414,10 @@ class ProxyBaseLLMRequestProcessing: if not restamp_model: return ProxyBaseLLMRequestProcessing.return_sse_chunk + restamper: Final = AnthropicStreamModelRestamper(restamp_model) + def serialize(chunk: object) -> str: - return ProxyBaseLLMRequestProcessing.return_sse_chunk( - restamp_anthropic_stream_chunk_model(chunk, restamp_model) - ) + return ProxyBaseLLMRequestProcessing.return_sse_chunk(restamper.process(chunk)) return serialize diff --git a/tests/test_litellm/proxy/anthropic_endpoints/test_streaming_model_restamp.py b/tests/test_litellm/proxy/anthropic_endpoints/test_streaming_model_restamp.py index 6e2c3f49445..385173b24a9 100644 --- a/tests/test_litellm/proxy/anthropic_endpoints/test_streaming_model_restamp.py +++ b/tests/test_litellm/proxy/anthropic_endpoints/test_streaming_model_restamp.py @@ -8,6 +8,7 @@ from unittest.mock import AsyncMock, MagicMock import pytest from litellm.proxy.anthropic_endpoints.streaming_model_restamp import ( + AnthropicStreamModelRestamper, restamp_anthropic_stream_chunk_model, ) from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing @@ -107,3 +108,85 @@ async def test_sse_generator_keeps_provider_model_when_restamping_is_off(): ] assert _model_from_frame(chunks[0]) == "claude-haiku-4-5-20251001" + + +def test_restamps_message_start_split_across_transport_chunks(): + frame = _message_start_frame("claude-haiku-4-5-20251001") + restamper = AnthropicStreamModelRestamper("claude-auto-1") + + held = restamper.process(frame[:25]) + emitted = restamper.process(frame[25:]) + + assert held == b"" + assert isinstance(emitted, bytes) + assert _model_from_frame(emitted) == "claude-auto-1" + + +def test_emits_coalesced_frames_with_only_message_start_rewritten(): + delta = b'event: content_block_delta\ndata: {"type":"content_block_delta","delta":{"text":"hi"}}\n\n' + combined = _message_start_frame("claude-haiku-4-5-20251001") + delta + + emitted = restamper_output = AnthropicStreamModelRestamper("claude-auto-1").process(combined) + + assert isinstance(restamper_output, bytes) + assert _model_from_frame(emitted) == "claude-auto-1" + assert emitted.endswith(delta) + + +def test_ping_frames_keep_the_restamper_armed(): + ping = b'event: ping\ndata: {"type": "ping"}\n\n' + frame = _message_start_frame("claude-haiku-4-5-20251001") + restamper = AnthropicStreamModelRestamper("claude-auto-1") + + assert restamper.process(ping) == ping + reassembled = restamper.process(frame[:10]) + reassembled += restamper.process(frame[10:]) + + assert _model_from_frame(reassembled) == "claude-auto-1" + + +def test_first_non_ping_event_disarms_the_restamper(): + delta = b'event: content_block_delta\ndata: {"type":"content_block_delta","delta":{"text":"hi"}}\n\n' + late_message_start = _message_start_frame("claude-haiku-4-5-20251001") + restamper = AnthropicStreamModelRestamper("claude-auto-1") + + assert restamper.process(delta) == delta + assert restamper.process(late_message_start) == late_message_start + + +def test_oversized_unterminated_chunk_flushes_unmodified(): + blob = b"data: " + b"x" * 70000 + restamper = AnthropicStreamModelRestamper("claude-auto-1") + + assert restamper.process(blob) == blob + frame = _message_start_frame("claude-haiku-4-5-20251001") + assert restamper.process(frame) == frame + + +def test_dict_message_start_disarms_after_restamp(): + restamper = AnthropicStreamModelRestamper("claude-auto-1") + first = restamper.process({"type": "message_start", "message": {"id": "msg_1", "model": "claude-sonnet-4-6"}}) + second = {"type": "message_start", "message": {"id": "msg_2", "model": "claude-sonnet-4-6"}} + + assert first == {"type": "message_start", "message": {"id": "msg_1", "model": "claude-auto-1"}} + assert restamper.process(second) == second + + +@pytest.mark.asyncio +async def test_sse_generator_restamps_message_start_split_across_chunks(): + frame = _message_start_frame("claude-haiku-4-5-20251001") + proxy_logging_obj = _proxy_logging_obj_streaming([frame[:30], frame[30:]]) + + chunks = [ + chunk + async for chunk in ProxyBaseLLMRequestProcessing.async_sse_data_generator( + response=MagicMock(), + user_api_key_dict=MagicMock(), + request_data={"model": "claude-auto-1"}, + proxy_logging_obj=proxy_logging_obj, + restamp_model="claude-auto-1", + ) + ] + + joined = b"".join(chunk if isinstance(chunk, bytes) else chunk.encode("utf-8") for chunk in chunks) + assert _model_from_frame(joined) == "claude-auto-1" From 7edf5b36cfd25666fab30f451b599b6fec524566 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 31 Aug 2026 16:01:39 -0700 Subject: [PATCH 04/97] fix(guardrails): deliver modify_response block as valid SSE on streaming chat and Responses A guardrail modify_response verdict on a streaming request only produced a proper replacement on /v1/messages: the chat completions and Responses API translations had no build_block_sse_chunks, so the ModifyResponseException re-raised and surfaced as an in-stream 500 error frame (or a whole-request 500 in buffered mode) instead of the documented 200 replacement. Implement build_block_sse_chunks for both OpenAI translations: chat emits a content delta plus a finish_reason content_filter chunk with real usage; Responses emits the typed event sequence (standalone via build_synthetic_response_events pre-stream, or an output-item continuation under the in-progress response id mid-stream) ending in response.completed. --- basedpyright-code-budget.json | 2 +- .../chat/guardrail_translation/handler.py | 10 +- .../guardrail_translation/base_translation.py | 5 +- .../base_llm/guardrail_translation/utils.py | 55 ++++ .../chat/guardrail_translation/handler.py | 117 +++++++- .../guardrail_translation/handler.py | 204 +++++++++++++- litellm/responses/streaming_iterator.py | 8 +- .../test_responses_hooks.py | 2 +- .../test_openai_guardrail_handler.py | 51 ++++ ...test_openai_responses_guardrail_handler.py | 63 +++++ .../test_openai_streaming_block.py | 265 ++++++++++++++++++ type-discipline-budget.json | 4 +- 12 files changed, 769 insertions(+), 17 deletions(-) create mode 100644 tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_openai_streaming_block.py diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index 83969d8dedf..88a612fde98 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -117,7 +117,7 @@ "limit": 117 }, "reportUnnecessaryComparison": { - "limit": 697 + "limit": 696 }, "reportUnnecessaryContains": { "limit": 5 diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index b9ca18c7843..21390e5ddb4 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -144,7 +144,7 @@ class AnthropicMessagesHandler(BaseTranslation): self, exc: "ModifyResponseException", stream_started: bool = False, - responses_so_far: list[object] | None = None, + responses_so_far: Sequence[object] | None = None, ) -> list[bytes]: """ Build an Anthropic SSE sequence delivering the guardrail block message @@ -162,7 +162,7 @@ class AnthropicMessagesHandler(BaseTranslation): would make Anthropic clients reject the stream. """ if stream_started: - return self._block_continuation_chunks(exc, responses_so_far or []) + return self._block_continuation_chunks(exc, responses_so_far or ()) return self._standalone_block_chunks(exc) def _standalone_block_chunks(self, exc: "ModifyResponseException") -> list[bytes]: @@ -187,7 +187,9 @@ class AnthropicMessagesHandler(BaseTranslation): ) return list(FakeAnthropicMessagesStreamIterator(response=block_response)) - def _block_continuation_chunks(self, exc: "ModifyResponseException", responses_so_far: list[object]) -> list[bytes]: + def _block_continuation_chunks( + self, exc: "ModifyResponseException", responses_so_far: Sequence[object] + ) -> list[bytes]: """Continue an already-started message: close the open content block, append the block message as a new text block, then end the message -- without a second message_start.""" @@ -237,7 +239,7 @@ class AnthropicMessagesHandler(BaseTranslation): @staticmethod def _content_block_state( - responses_so_far: list[object], + responses_so_far: Sequence[object], ) -> tuple[int | None, int | None]: """From the SSE chunks already sent to the client, return (open content-block index or None, highest content-block index seen or None). diff --git a/litellm/llms/base_llm/guardrail_translation/base_translation.py b/litellm/llms/base_llm/guardrail_translation/base_translation.py index ba96ab3dc99..247d9bc7c40 100644 --- a/litellm/llms/base_llm/guardrail_translation/base_translation.py +++ b/litellm/llms/base_llm/guardrail_translation/base_translation.py @@ -1,4 +1,5 @@ from abc import ABC, abstractmethod +from collections.abc import Sequence from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any, Final, Optional @@ -127,8 +128,8 @@ class BaseTranslation(ABC): self, exc: "ModifyResponseException", stream_started: bool = False, - responses_so_far: list[Any] | None = None, - ) -> list[bytes] | None: + responses_so_far: Sequence[Any] | None = None, + ) -> Sequence[bytes] | None: """ Build the streaming chunks that deliver a guardrail block message and cleanly terminate the stream in this provider's wire format. diff --git a/litellm/llms/base_llm/guardrail_translation/utils.py b/litellm/llms/base_llm/guardrail_translation/utils.py index f09ee210e6c..d30cdcecff5 100644 --- a/litellm/llms/base_llm/guardrail_translation/utils.py +++ b/litellm/llms/base_llm/guardrail_translation/utils.py @@ -124,6 +124,61 @@ def blocked_responses_api_usage(original_response: object) -> ResponseAPIUsage: ) +def stream_item_field(item: object, field: str) -> object | None: + if isinstance(item, dict): + return item.get(field) + return getattr(item, field, None) + + +def blocked_chat_stream_usage(original_response: object) -> tuple[int, int]: + """ + ``(prompt_tokens, completion_tokens)`` for a synthetic guardrail-blocked + chat completions stream. + + A mid-stream block carries the chunks received so far as a list; real usage + rides on the final chunk when the upstream sent one + (``stream_options.include_usage``). Non-list originals defer to + ``blocked_response_usage``. + """ + if not isinstance(original_response, list): + usage: Final = blocked_response_usage(original_response) + return usage.get("input_tokens", 0), usage.get("output_tokens", 0) + usage_obj: Final = next( + ( + chunk_usage + for item in reversed(original_response) + if (chunk_usage := stream_item_field(item, "usage")) is not None + ), + None, + ) + return ( + _usage_tokens(usage_obj, "prompt_tokens", "input_tokens"), + _usage_tokens(usage_obj, "completion_tokens", "output_tokens"), + ) + + +def blocked_responses_stream_usage(original_response: object) -> ResponseAPIUsage: + """ + ``ResponseAPIUsage`` for a synthetic guardrail-blocked /v1/responses stream. + + A mid-stream block carries the events received so far as a list; real usage + rides on the ``response.completed`` event's response when the upstream sent + one. Non-list originals defer to ``blocked_responses_api_usage``. + """ + if not isinstance(original_response, list): + return blocked_responses_api_usage(original_response) + completed: Final = next( + ( + response + for item in reversed(original_response) + if str(stream_item_field(item, "type") or "") == "response.completed" + and (response := stream_item_field(item, "response")) is not None + ), + None, + ) + return blocked_responses_api_usage(completed) + + def effective_skip_system_message_for_guardrail(guardrail_to_apply: Any) -> bool: per: Final = getattr(guardrail_to_apply, "skip_system_message_in_guardrail", None) if per is not None: diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index 54673c77f80..31a36822c89 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/litellm/llms/openai/chat/guardrail_translation/handler.py @@ -14,8 +14,14 @@ Pattern Overview: This pattern can be replicated for other message formats (e.g., Anthropic). """ +import json +import time +import uuid +from collections.abc import Sequence from typing import TYPE_CHECKING, Any, Final, Union, cast +from typing_extensions import NotRequired, ReadOnly, TypedDict + import litellm from litellm._logging import verbose_proxy_logger from litellm.llms.base_llm.guardrail_translation.base_translation import ( @@ -23,6 +29,7 @@ from litellm.llms.base_llm.guardrail_translation.base_translation import ( StreamTransformSink, ) from litellm.llms.base_llm.guardrail_translation.utils import ( + blocked_chat_stream_usage, effective_scan_only_tool_results_for_guardrail, effective_skip_system_message_for_guardrail, effective_skip_tool_message_for_guardrail, @@ -31,6 +38,7 @@ from litellm.llms.base_llm.guardrail_translation.utils import ( openai_tool_name, role_out_of_guardrail_scope, scoped_structured_message_indices, + stream_item_field, ) from litellm.main import stream_chunk_builder from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam @@ -46,7 +54,10 @@ from litellm.types.utils import ( ) if TYPE_CHECKING: - from litellm.integrations.custom_guardrail import CustomGuardrail + from litellm.integrations.custom_guardrail import ( + CustomGuardrail, + ModifyResponseException, + ) from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj @@ -1000,3 +1011,107 @@ class OpenAIChatCompletionsHandler(BaseTranslation): else: # Subsequent chunks - clear the text content_item["text"] = "" + + def build_block_sse_chunks( + self, + exc: "ModifyResponseException", + stream_started: bool = False, + responses_so_far: Sequence[object] | None = None, + ) -> Sequence[bytes]: + """ + Build OpenAI chat-completions SSE chunks that deliver the guardrail + block message and terminate the stream cleanly, mirroring the + non-streaming block response: ``finish_reason`` ``content_filter`` plus + the real usage the upstream call consumed. + + - ``stream_started`` False (buffered / pre-stream): nothing has been + sent, so open a standalone completion with a ``role`` delta. + - ``stream_started`` True (sampling / mid-stream): chunks already + reached the client, so continue the in-progress completion (reuse its + id/created/model, content-only delta). + + The proxy's data generator appends ``data: [DONE]`` itself. + """ + chunk_id, created, model = _blocked_stream_identity(exc, responses_so_far or ()) + prompt_tokens, completion_tokens = blocked_chat_stream_usage(exc.original_response) + continuation_delta: Final[_BlockedChunkDelta] = {"content": exc.message} + standalone_delta: Final[_BlockedChunkDelta] = {"role": "assistant", "content": exc.message} + message_chunk: Final[_BlockedChunk] = { + "id": chunk_id, + "object": "chat.completion.chunk", + "created": created, + "model": model, + "choices": ( + { + "index": 0, + "delta": continuation_delta if stream_started else standalone_delta, + "finish_reason": None, + }, + ), + } + final_chunk: Final[_BlockedChunk] = { + "id": chunk_id, + "object": "chat.completion.chunk", + "created": created, + "model": model, + "choices": ({"index": 0, "delta": {}, "finish_reason": "content_filter"},), + "usage": { + "prompt_tokens": prompt_tokens, + "completion_tokens": completion_tokens, + "total_tokens": prompt_tokens + completion_tokens, + }, + } + return _chat_sse_chunk(message_chunk), _chat_sse_chunk(final_chunk) + + +class _BlockedChunkDelta(TypedDict, total=False): + role: ReadOnly[str] + content: ReadOnly[str] + + +class _BlockedChunkChoice(TypedDict): + index: ReadOnly[int] + delta: ReadOnly[_BlockedChunkDelta] + finish_reason: ReadOnly[str | None] + + +class _BlockedChunkUsage(TypedDict): + prompt_tokens: ReadOnly[int] + completion_tokens: ReadOnly[int] + total_tokens: ReadOnly[int] + + +class _BlockedChunk(TypedDict): + id: ReadOnly[str] + object: ReadOnly[str] + created: ReadOnly[int] + model: ReadOnly[str] + choices: ReadOnly[tuple[_BlockedChunkChoice, ...]] + usage: NotRequired[ReadOnly[_BlockedChunkUsage]] + + +def _chat_sse_chunk(payload: _BlockedChunk) -> bytes: + return f"data: {json.dumps(payload)}\n\n".encode() + + +def _blocked_stream_identity( + exc: "ModifyResponseException", responses_so_far: Sequence[object] +) -> tuple[str, int, str]: + identified: Final = next( + ( + (chunk_id, item) + for item in responses_so_far + if isinstance(chunk_id := stream_item_field(item, "id"), str) and chunk_id + ), + None, + ) + if identified is None: + return f"chatcmpl-{uuid.uuid4()}", int(time.time()), exc.model + chunk_id, source = identified + created: Final = stream_item_field(source, "created") + model: Final = stream_item_field(source, "model") + return ( + chunk_id, + created if isinstance(created, int) else int(time.time()), + model if isinstance(model, str) and model else exc.model, + ) diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index 7c5d8ac99ad..9a58899eccd 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -28,6 +28,8 @@ Output: response.output is List[GenericResponseOutputItem] where each has: - text: str """ +import time +import uuid from collections.abc import Sequence from typing import TYPE_CHECKING, Any, Final, Union, cast @@ -41,15 +43,31 @@ from litellm.completion_extras.litellm_responses_transformation.transformation i OpenAiResponsesToChatCompletionStreamIterator, ) from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation +from litellm.llms.base_llm.guardrail_translation.utils import ( + blocked_responses_stream_usage, + stream_item_field, +) from litellm.responses.litellm_completion_transformation.transformation import ( LiteLLMCompletionResponsesConfig, ) from litellm.types.llms.openai import ( AllMessageValues, + BaseLiteLLMOpenAIResponseObject, ChatCompletionToolCallChunk, ChatCompletionToolParam, + ContentPartAddedEvent, + ContentPartDoneEvent, + ContentPartDonePartOutputText, OpenAIMcpServerTool, + OutputItemAddedEvent, + OutputItemDoneEvent, + OutputTextDeltaEvent, + OutputTextDoneEvent, + ResponseAPIUsage, + ResponseCompletedEvent, + ResponsesAPIResponse, ResponsesAPIStreamEvents, + ResponsesAPIStreamingResponse, ) from litellm.types.responses.main import ( GenericResponseOutputItem, @@ -59,11 +77,13 @@ from litellm.types.responses.main import ( from litellm.types.utils import GenericGuardrailAPIInputs if TYPE_CHECKING: - from litellm.integrations.custom_guardrail import CustomGuardrail + from litellm.integrations.custom_guardrail import ( + CustomGuardrail, + ModifyResponseException, + ) from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.proxy._types import UserAPIKeyAuth from litellm.types.llms.openai import ResponseInputParam - from litellm.types.utils import ResponsesAPIResponse class ResponseOutputEnvelope(TypedDict, total=False): @@ -802,3 +822,183 @@ class OpenAIResponsesHandler(BaseTranslation): content[content_idx]["text"] = guardrail_response elif hasattr(content[content_idx], "text"): content[content_idx].text = guardrail_response + + def build_block_sse_chunks( + self, + exc: "ModifyResponseException", + stream_started: bool = False, + responses_so_far: Sequence[object] | None = None, + ) -> Sequence[bytes]: + """ + Build Responses API SSE events that deliver the guardrail block message + and terminate the stream cleanly, mirroring the non-streaming block + response: a completed response whose only output is the violation text, + with the real usage the upstream call consumed. + + - ``stream_started`` False (buffered / pre-stream): nothing has been + sent, so emit the full synthetic sequence (``response.created`` + through ``response.completed``). + - ``stream_started`` True (sampling / mid-stream): events already + reached the client, so continue the in-progress response: deliver the + block message as a new output item under the same response id and + close with a ``response.completed`` carrying only the replacement + item. + + The proxy's data generator appends ``data: [DONE]`` itself. + """ + events: Final = ( + self._block_continuation_events(exc, responses_so_far or ()) + if stream_started + else self._standalone_block_events(exc) + ) + return tuple( + f"data: {event.model_dump_json(exclude_none=True, exclude_unset=True)}\n\n".encode() for event in events + ) + + @staticmethod + def _standalone_block_events(exc: "ModifyResponseException") -> Sequence[ResponsesAPIStreamingResponse]: + from litellm.responses.streaming_iterator import build_synthetic_response_events + + return build_synthetic_response_events( + transformed=_blocked_response(exc, response_id=f"resp_{uuid.uuid4()}", model=exc.model), + logging_obj=None, + chunk_size=max(len(exc.message), 1), + ) + + @staticmethod + def _block_continuation_events( + exc: "ModifyResponseException", responses_so_far: Sequence[object] + ) -> Sequence[ResponsesAPIStreamingResponse]: + response_id, model, output_index = _continuation_identity(exc, responses_so_far) + item: Final = _blocked_output_item(exc) + item_id: Final = item["id"] + item_model: Final = BaseLiteLLMOpenAIResponseObject.model_validate(item) + part: Final[_BlockedContentPart] = {"type": "output_text", "text": exc.message, "annotations": ()} + done_part: Final[_BlockedDoneContentPart] = { + "type": "output_text", + "text": exc.message, + "annotations": (), + "logprobs": None, + } + return ( + OutputItemAddedEvent( + type=ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, + output_index=output_index, + item=item_model, + ), + ContentPartAddedEvent( + type=ResponsesAPIStreamEvents.CONTENT_PART_ADDED, + item_id=item_id, + output_index=output_index, + content_index=0, + part=BaseLiteLLMOpenAIResponseObject.model_validate(part), + ), + OutputTextDeltaEvent( + type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA, + item_id=item_id, + output_index=output_index, + content_index=0, + delta=exc.message, + ), + OutputTextDoneEvent( + type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DONE, + item_id=item_id, + output_index=output_index, + content_index=0, + text=exc.message, + ), + ContentPartDoneEvent( + type=ResponsesAPIStreamEvents.CONTENT_PART_DONE, + item_id=item_id, + output_index=output_index, + content_index=0, + part=ContentPartDonePartOutputText.model_validate(done_part), + ), + OutputItemDoneEvent( + type=ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE, + output_index=output_index, + item=item_model, + ), + ResponseCompletedEvent( + type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED, + response=_blocked_response(exc, response_id=response_id, model=model, output_item=item), + ), + ) + + +class _BlockedContentPart(TypedDict): + type: ReadOnly[str] + text: ReadOnly[str] + annotations: ReadOnly[tuple[object, ...]] + + +class _BlockedDoneContentPart(TypedDict): + type: ReadOnly[str] + text: ReadOnly[str] + annotations: ReadOnly[tuple[object, ...]] + logprobs: ReadOnly[None] + + +class _BlockedOutputItem(TypedDict): + type: ReadOnly[str] + id: ReadOnly[str] + status: ReadOnly[str] + role: ReadOnly[str] + content: ReadOnly[tuple[_BlockedContentPart, ...]] + + +class _BlockedResponsePayload(TypedDict): + id: ReadOnly[str] + object: ReadOnly[str] + created_at: ReadOnly[int] + model: ReadOnly[str] + output: ReadOnly[tuple[_BlockedOutputItem, ...]] + status: ReadOnly[str] + usage: ReadOnly[ResponseAPIUsage] + + +def _blocked_output_item(exc: "ModifyResponseException") -> _BlockedOutputItem: + item: Final[_BlockedOutputItem] = { + "type": "message", + "id": f"msg_{uuid.uuid4()}", + "status": "completed", + "role": "assistant", + "content": ({"type": "output_text", "text": exc.message, "annotations": ()},), + } + return item + + +def _blocked_response( + exc: "ModifyResponseException", + response_id: str, + model: str, + output_item: _BlockedOutputItem | None = None, +) -> ResponsesAPIResponse: + payload: Final[_BlockedResponsePayload] = { + "id": response_id, + "object": "response", + "created_at": int(time.time()), + "model": model, + "output": (output_item if output_item is not None else _blocked_output_item(exc),), + "status": "completed", + "usage": blocked_responses_stream_usage(exc.original_response), + } + return ResponsesAPIResponse.model_validate(payload) + + +def _continuation_identity(exc: "ModifyResponseException", responses_so_far: Sequence[object]) -> tuple[str, str, int]: + responses: Final = tuple( + response for item in responses_so_far if (response := stream_item_field(item, "response")) is not None + ) + response_id: Final = next( + (rid for response in responses if isinstance(rid := stream_item_field(response, "id"), str) and rid), + f"resp_{uuid.uuid4()}", + ) + model: Final = next( + (m for response in responses if isinstance(m := stream_item_field(response, "model"), str) and m), + exc.model, + ) + indices: Final = tuple( + index for item in responses_so_far if isinstance(index := stream_item_field(item, "output_index"), int) + ) + return response_id, model, max(indices) + 1 if indices else 0 diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 63adf950142..4fd47c50b1e 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -969,7 +969,7 @@ class MockResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): transformed: ResponsesAPIResponse, logging_obj: LiteLLMLoggingObj, ) -> None: - self._events: list[ResponsesAPIStreamingResponse] = _build_synthetic_response_events( + self._events: Sequence[ResponsesAPIStreamingResponse] = build_synthetic_response_events( transformed=transformed, logging_obj=logging_obj, chunk_size=self.CHUNK_SIZE, @@ -1036,7 +1036,7 @@ class CachedResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): transformed: ResponsesAPIResponse, logging_obj: LiteLLMLoggingObj, ) -> None: - self._events = _build_synthetic_response_events( + self._events = build_synthetic_response_events( transformed=transformed, logging_obj=logging_obj, chunk_size=MockResponsesAPIStreamingIterator.CHUNK_SIZE, @@ -1218,10 +1218,10 @@ def _add_text_like_part_events( ) -def _build_synthetic_response_events( +def build_synthetic_response_events( *, transformed: ResponsesAPIResponse, - logging_obj: LiteLLMLoggingObj, + logging_obj: LiteLLMLoggingObj | None, chunk_size: int, ) -> list[ResponsesAPIStreamingResponse]: openai_types: Final = _get_openai_response_types() diff --git a/tests/llm_responses_api_testing/test_responses_hooks.py b/tests/llm_responses_api_testing/test_responses_hooks.py index 66dbb29dba5..a86752c0172 100644 --- a/tests/llm_responses_api_testing/test_responses_hooks.py +++ b/tests/llm_responses_api_testing/test_responses_hooks.py @@ -841,7 +841,7 @@ def test_build_synthetic_response_events_covers_annotations_function_calls_and_r ) try: - events = streaming_module._build_synthetic_response_events( + events = streaming_module.build_synthetic_response_events( transformed=transformed, logging_obj=logging_obj, chunk_size=5, diff --git a/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py b/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py index a29e0be4655..8ad0026d1ab 100644 --- a/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py +++ b/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py @@ -1559,3 +1559,54 @@ class TestScanOnlyToolResults: assert data["messages"][3]["content"] == "page says [BLOCKED] here" assert data["messages"][3]["tool_call_id"] == "call_1" assert data["messages"][4]["content"] == "and then?" + + +class TestBuildBlockSseChunks: + """build_block_sse_chunks turns a streaming ModifyResponseException into 200 SSE chunks""" + + def _exc(self, original_response=None): + from litellm.exceptions import ModifyResponseException + + return ModifyResponseException( + message="Blocked by policy.", + model="gpt-5.4-mini", + request_data={}, + guardrail_name="test", + original_response=original_response, + ) + + def _payloads(self, chunks): + return [json.loads(chunk.decode().removeprefix("data: ").strip()) for chunk in chunks] + + def test_standalone_block_uses_fresh_identity_and_zero_usage(self): + handler = OpenAIChatCompletionsHandler() + first, final = self._payloads(handler.build_block_sse_chunks(self._exc(), stream_started=False)) + assert first["id"].startswith("chatcmpl-") + assert first["model"] == "gpt-5.4-mini" + assert first["choices"][0]["delta"] == {"role": "assistant", "content": "Blocked by policy."} + assert first["choices"][0]["finish_reason"] is None + assert final["choices"][0]["delta"] == {} + assert final["choices"][0]["finish_reason"] == "content_filter" + assert final["usage"] == {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0} + + def test_continuation_reuses_stream_identity_and_real_usage(self): + handler = OpenAIChatCompletionsHandler() + yielded = [ + {"id": "chatcmpl-live", "created": 1724900000, "model": "gpt-5.4-mini-2026-01-01"}, + ] + original = yielded + [ + {"id": "chatcmpl-live", "usage": {"prompt_tokens": 11, "completion_tokens": 5}}, + ] + first, final = self._payloads( + handler.build_block_sse_chunks( + self._exc(original_response=original), stream_started=True, responses_so_far=yielded + ) + ) + assert (first["id"], first["created"], first["model"]) == ( + "chatcmpl-live", + 1724900000, + "gpt-5.4-mini-2026-01-01", + ) + assert first["choices"][0]["delta"] == {"content": "Blocked by policy."} + assert final["id"] == "chatcmpl-live" + assert final["usage"] == {"prompt_tokens": 11, "completion_tokens": 5, "total_tokens": 16} diff --git a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py index 447175b09a6..cd0e7f29933 100644 --- a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py +++ b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py @@ -1229,3 +1229,66 @@ class TestOpenAIResponsesHandlerToolInjection: names = [t.get("name") for t in result["tools"]] assert "get_weather" in names assert "injected_tool" in names + + +class TestBuildBlockSseChunks: + """build_block_sse_chunks turns a streaming ModifyResponseException into 200 SSE events""" + + def _exc(self, original_response=None): + from litellm.exceptions import ModifyResponseException + + return ModifyResponseException( + message="Blocked by policy.", + model="gpt-5.4-mini", + request_data={}, + guardrail_name="test", + original_response=original_response, + ) + + def _payloads(self, chunks): + import json + + return [json.loads(chunk.decode().removeprefix("data: ").strip()) for chunk in chunks] + + def test_standalone_block_emits_complete_synthetic_stream(self): + handler = OpenAIResponsesHandler() + payloads = self._payloads(handler.build_block_sse_chunks(self._exc(), stream_started=False)) + types = [payload["type"] for payload in payloads] + assert types[0] == "response.created" + assert types[-1] == "response.completed" + completed = payloads[-1]["response"] + assert completed["id"].startswith("resp_") + assert completed["model"] == "gpt-5.4-mini" + assert completed["output"][0]["content"][0]["text"] == "Blocked by policy." + + def test_continuation_appends_item_at_next_output_index_with_real_usage(self): + handler = OpenAIResponsesHandler() + yielded = [ + {"type": "response.created", "response": {"id": "resp_live", "model": "gpt-5.4-mini-2026-01-01"}}, + {"type": "response.output_item.added", "output_index": 2, "item": {"id": "msg_orig"}}, + ] + original = yielded + [ + { + "type": "response.completed", + "response": { + "id": "resp_live", + "model": "gpt-5.4-mini-2026-01-01", + "output": [], + "usage": {"input_tokens": 7, "output_tokens": 21, "total_tokens": 28}, + }, + } + ] + payloads = self._payloads( + handler.build_block_sse_chunks( + self._exc(original_response=original), stream_started=True, responses_so_far=yielded + ) + ) + types = [payload["type"] for payload in payloads] + assert "response.created" not in types + assert types[0] == "response.output_item.added" + assert payloads[0]["output_index"] == 3 + completed = payloads[-1]["response"] + assert completed["id"] == "resp_live" + assert completed["model"] == "gpt-5.4-mini-2026-01-01" + assert completed["output"][0]["content"][0]["text"] == "Blocked by policy." + assert completed["usage"] == {"input_tokens": 7, "output_tokens": 21, "total_tokens": 28} diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_openai_streaming_block.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_openai_streaming_block.py new file mode 100644 index 00000000000..4d5581290ba --- /dev/null +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_openai_streaming_block.py @@ -0,0 +1,265 @@ +""" +Regression tests for blocking an OpenAI-format streaming response from the +unified guardrail post-call streaming iterator hook. + +When a guardrail's ``apply_guardrail`` raises ``ModifyResponseException`` +while (or at the end of) a chat completions or Responses API stream is being +relayed, the hook must emit a well-formed SSE termination sequence carrying +the block message - NOT a bare ``data: {"error": ...}`` blob that surfaces as +an HTTP 500 error frame and truncates the stream. +""" + +import json +from typing import Any, AsyncGenerator, List, Literal, Optional + +import pytest + +from litellm.integrations.custom_guardrail import ( + CustomGuardrail, + ModifyResponseException, +) +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import ( + UnifiedLLMGuardrails, +) +from litellm.types.utils import ( + Delta, + GenericGuardrailAPIInputs, + ModelResponseStream, + StreamingChoices, +) + +BLOCK_MESSAGE = "This response was replaced by policy." + + +class _BlockingGuardrail(CustomGuardrail): + """Mock guardrail that always blocks response scans by raising ModifyResponseException.""" + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[Any] = None, + ) -> GenericGuardrailAPIInputs: + raise ModifyResponseException( + message=BLOCK_MESSAGE, + model="gpt-5.4-mini", + request_data=request_data, + guardrail_name=self.guardrail_name, + ) + + +def _chat_chunk(delta: Delta, finish_reason: Optional[str] = None) -> ModelResponseStream: + return ModelResponseStream( + id="chatcmpl-live", + created=1724900000, + model="gpt-5.4-mini", + choices=[StreamingChoices(index=0, delta=delta, finish_reason=finish_reason)], + ) + + +async def _chat_stream(end: bool) -> AsyncGenerator[ModelResponseStream, None]: + yield _chat_chunk(Delta(role="assistant", content="This ")) + for text in ["is ", "the ", "original ", "answer."]: + yield _chat_chunk(Delta(content=text)) + if end: + yield _chat_chunk(Delta(), finish_reason="stop") + + +async def _responses_stream(end: bool) -> AsyncGenerator[dict, None]: + original_text = "This is the original answer." + response_envelope = {"id": "resp_live", "model": "gpt-5.4-mini", "status": "in_progress", "output": []} + yield {"type": "response.created", "response": response_envelope} + yield {"type": "response.in_progress", "response": response_envelope} + yield { + "type": "response.output_item.added", + "output_index": 0, + "item": {"id": "msg_orig", "type": "message", "role": "assistant", "content": []}, + } + yield { + "type": "response.content_part.added", + "item_id": "msg_orig", + "output_index": 0, + "content_index": 0, + "part": {"type": "output_text", "text": "", "annotations": []}, + } + for delta in ["This ", "is ", "the ", "original ", "answer."]: + yield { + "type": "response.output_text.delta", + "item_id": "msg_orig", + "output_index": 0, + "content_index": 0, + "delta": delta, + } + yield { + "type": "response.output_text.done", + "item_id": "msg_orig", + "output_index": 0, + "content_index": 0, + "text": original_text, + } + if end: + yield { + "type": "response.completed", + "response": { + "id": "resp_live", + "model": "gpt-5.4-mini", + "status": "completed", + "output": [ + { + "id": "msg_orig", + "type": "message", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": original_text, "annotations": []}], + } + ], + "usage": {"input_tokens": 7, "output_tokens": 21, "total_tokens": 28}, + }, + } + + +async def _run_hook( + route: str, + stream: AsyncGenerator[Any, None], + sampling_rate: int = 1, + end_of_stream_only: bool = False, + buffer_until_moderated: bool = False, +) -> List[Any]: + guardrail = _BlockingGuardrail(guardrail_name="test-blocking-guardrail", event_hook="post_call") + guardrail.streaming_sampling_rate = sampling_rate + guardrail.streaming_end_of_stream_only = end_of_stream_only + guardrail.streaming_buffer_until_moderated = buffer_until_moderated + + unified_guardrail = UnifiedLLMGuardrails() + user_api_key_dict = UserAPIKeyAuth(api_key="test", request_route=route) + request_data = { + "messages": [{"role": "user", "content": "hi"}], + "guardrail_to_apply": guardrail, + "metadata": {"guardrails": ["test-blocking-guardrail"]}, + } + + collected: List[Any] = [] + async for chunk in unified_guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=user_api_key_dict, + response=stream, + request_data=request_data, + ): + collected.append(chunk) + return collected + + +def _sse_payloads(collected: List[Any]) -> List[dict]: + payloads = [] + for chunk in collected: + if not isinstance(chunk, bytes): + continue + for block in chunk.decode().split("\n\n"): + for line in block.strip().split("\n"): + if line.startswith("data:"): + payloads.append(json.loads(line[len("data:") :].strip())) + return payloads + + +def _assert_no_error_frame(collected: List[Any]) -> None: + raw = "".join(chunk.decode() for chunk in collected if isinstance(chunk, bytes)) + assert '"error"' not in raw, f"unexpected error blob in stream: {raw!r}" + + +@pytest.mark.asyncio +async def test_chat_pre_stream_block_emits_standalone_completion(): + """Block on the first chunk: a standalone completion opens with a role delta + and ends with finish_reason content_filter.""" + collected = await _run_hook("/v1/chat/completions", _chat_stream(end=False)) + _assert_no_error_frame(collected) + payloads = _sse_payloads(collected) + assert payloads, "no block SSE chunks were emitted" + assert payloads[0]["choices"][0]["delta"] == {"role": "assistant", "content": BLOCK_MESSAGE} + assert payloads[-1]["choices"][0]["finish_reason"] == "content_filter" + + +@pytest.mark.asyncio +async def test_chat_mid_stream_block_continues_the_completion(): + """Regression for the LIT-6496 500 error frame: after chunks were already + forwarded, the block continues the same completion id and terminates with + finish_reason content_filter instead of raising into an error blob.""" + collected = await _run_hook("/v1/chat/completions", _chat_stream(end=False), sampling_rate=5) + _assert_no_error_frame(collected) + forwarded = [chunk for chunk in collected if isinstance(chunk, ModelResponseStream)] + assert forwarded, "original chunks should have streamed before the block" + payloads = _sse_payloads(collected) + assert payloads, "no block SSE chunks were emitted" + assert all(payload["id"] == "chatcmpl-live" for payload in payloads), ( + "block chunks must continue the in-progress completion, not start a new one" + ) + assert payloads[0]["choices"][0]["delta"] == {"content": BLOCK_MESSAGE} + assert payloads[-1]["choices"][0]["finish_reason"] == "content_filter" + + +@pytest.mark.asyncio +async def test_chat_end_of_stream_block_terminates_cleanly(): + collected = await _run_hook("/v1/chat/completions", _chat_stream(end=True), end_of_stream_only=True) + _assert_no_error_frame(collected) + payloads = _sse_payloads(collected) + assert BLOCK_MESSAGE in json.dumps(payloads) + assert payloads[-1]["choices"][0]["finish_reason"] == "content_filter" + + +@pytest.mark.asyncio +async def test_responses_buffered_block_emits_full_event_sequence(): + """Buffered moderation blocks before anything streams: a complete synthetic + Responses stream from response.created through response.completed carrying + the block message, with the original content never released.""" + collected = await _run_hook("/v1/responses", _responses_stream(end=True), buffer_until_moderated=True) + _assert_no_error_frame(collected) + assert not [chunk for chunk in collected if isinstance(chunk, dict)], ( + "buffered original chunks must never be released after a block" + ) + payloads = _sse_payloads(collected) + event_types = [payload["type"] for payload in payloads] + assert event_types[0] == "response.created" + assert "response.output_text.delta" in event_types + assert event_types[-1] == "response.completed" + completed = payloads[-1]["response"] + assert completed["status"] == "completed" + assert completed["output"][0]["content"][0]["text"] == BLOCK_MESSAGE + assert completed["usage"] == {"input_tokens": 7, "output_tokens": 21, "total_tokens": 28} + assert "original answer" not in json.dumps(payloads) + + +@pytest.mark.asyncio +async def test_responses_mid_stream_block_continues_the_response(): + """Regression for the LIT-6496 500 error frame: after events were already + forwarded, the block appends a new output item under the same response id + and closes with response.completed - never a second response.created.""" + collected = await _run_hook("/v1/responses", _responses_stream(end=False)) + _assert_no_error_frame(collected) + forwarded_types = [chunk["type"] for chunk in collected if isinstance(chunk, dict)] + assert "response.created" in forwarded_types, "original events should have streamed before the block" + payloads = _sse_payloads(collected) + assert payloads, "no block SSE chunks were emitted" + block_types = [payload["type"] for payload in payloads] + assert "response.created" not in block_types, "a mid-stream block must not restart the response" + assert block_types[0] == "response.output_item.added" + assert block_types[-1] == "response.completed" + assert payloads[0]["output_index"] == 1, "the block item must continue after the original output item" + completed = payloads[-1]["response"] + assert completed["id"] == "resp_live" + assert completed["output"][0]["content"][0]["text"] == BLOCK_MESSAGE + + +@pytest.mark.asyncio +async def test_responses_end_of_stream_block_reports_original_usage(): + collected = await _run_hook("/v1/responses", _responses_stream(end=True), end_of_stream_only=True) + _assert_no_error_frame(collected) + forwarded_types = [chunk["type"] for chunk in collected if isinstance(chunk, dict)] + assert "response.completed" not in forwarded_types, ( + "the original terminal event must be withheld and replaced by the block sequence" + ) + payloads = _sse_payloads(collected) + completed = payloads[-1]["response"] + assert payloads[-1]["type"] == "response.completed" + assert completed["id"] == "resp_live" + assert completed["output"][0]["content"][0]["text"] == BLOCK_MESSAGE + assert completed["usage"] == {"input_tokens": 7, "output_tokens": 21, "total_tokens": 28} diff --git a/type-discipline-budget.json b/type-discipline-budget.json index 7365cec4fdd..f17d0b7af4a 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,9 +1,9 @@ { "LIT001": { - "limit": 22704 + "limit": 22698 }, "LIT002": { - "limit": 26854 + "limit": 26853 }, "LIT003": { "limit": 269 From 31a9f7e6ad932b0adbc24f7666178aa4c1071984 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 31 Aug 2026 16:34:21 -0700 Subject: [PATCH 05/97] test(guardrails): type streaming-block test helpers and drop mutable accumulators --- .../test_openai_streaming_block.py | 50 ++++++++++--------- 1 file changed, 27 insertions(+), 23 deletions(-) diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_openai_streaming_block.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_openai_streaming_block.py index 4d5581290ba..99172ec3be6 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_openai_streaming_block.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_openai_streaming_block.py @@ -10,7 +10,7 @@ an HTTP 500 error frame and truncates the stream. """ import json -from typing import Any, AsyncGenerator, List, Literal, Optional +from typing import Any, AsyncGenerator, Dict, Literal, Optional, Tuple, Union import pytest @@ -31,6 +31,9 @@ from litellm.types.utils import ( BLOCK_MESSAGE = "This response was replaced by policy." +JsonPayload = Dict[str, object] +StreamChunk = Union[ModelResponseStream, JsonPayload, bytes] + class _BlockingGuardrail(CustomGuardrail): """Mock guardrail that always blocks response scans by raising ModifyResponseException.""" @@ -67,7 +70,7 @@ async def _chat_stream(end: bool) -> AsyncGenerator[ModelResponseStream, None]: yield _chat_chunk(Delta(), finish_reason="stop") -async def _responses_stream(end: bool) -> AsyncGenerator[dict, None]: +async def _responses_stream(end: bool) -> AsyncGenerator[JsonPayload, None]: original_text = "This is the original answer." response_envelope = {"id": "resp_live", "model": "gpt-5.4-mini", "status": "in_progress", "output": []} yield {"type": "response.created", "response": response_envelope} @@ -122,11 +125,11 @@ async def _responses_stream(end: bool) -> AsyncGenerator[dict, None]: async def _run_hook( route: str, - stream: AsyncGenerator[Any, None], + stream: AsyncGenerator[Union[ModelResponseStream, JsonPayload], None], sampling_rate: int = 1, end_of_stream_only: bool = False, buffer_until_moderated: bool = False, -) -> List[Any]: +) -> Tuple[StreamChunk, ...]: guardrail = _BlockingGuardrail(guardrail_name="test-blocking-guardrail", event_hook="post_call") guardrail.streaming_sampling_rate = sampling_rate guardrail.streaming_end_of_stream_only = end_of_stream_only @@ -140,29 +143,30 @@ async def _run_hook( "metadata": {"guardrails": ["test-blocking-guardrail"]}, } - collected: List[Any] = [] - async for chunk in unified_guardrail.async_post_call_streaming_iterator_hook( - user_api_key_dict=user_api_key_dict, - response=stream, - request_data=request_data, - ): - collected.append(chunk) - return collected + return tuple( + [ + chunk + async for chunk in unified_guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=user_api_key_dict, + response=stream, + request_data=request_data, + ) + ] + ) -def _sse_payloads(collected: List[Any]) -> List[dict]: - payloads = [] - for chunk in collected: - if not isinstance(chunk, bytes): - continue - for block in chunk.decode().split("\n\n"): - for line in block.strip().split("\n"): - if line.startswith("data:"): - payloads.append(json.loads(line[len("data:") :].strip())) - return payloads +def _sse_payloads(collected: Tuple[StreamChunk, ...]) -> Tuple[JsonPayload, ...]: + return tuple( + json.loads(line[len("data:") :].strip()) + for chunk in collected + if isinstance(chunk, bytes) + for block in chunk.decode().split("\n\n") + for line in block.strip().split("\n") + if line.startswith("data:") + ) -def _assert_no_error_frame(collected: List[Any]) -> None: +def _assert_no_error_frame(collected: Tuple[StreamChunk, ...]) -> None: raw = "".join(chunk.decode() for chunk in collected if isinstance(chunk, bytes)) assert '"error"' not in raw, f"unexpected error blob in stream: {raw!r}" From 78b57fb427de52a3fc0f10ba5c20005f46725a6f Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 31 Aug 2026 17:05:45 -0700 Subject: [PATCH 06/97] fix(guardrails): withhold chat finish chunk in end_of_stream_only mode and close open Responses items before a mid-stream block --- .../chat/guardrail_translation/handler.py | 22 +++ .../guardrail_translation/handler.py | 133 +++++++++++++++++- .../test_openai_guardrail_handler.py | 33 +++++ ...test_openai_responses_guardrail_handler.py | 26 +++- .../test_openai_streaming_block.py | 74 ++++++++-- 5 files changed, 273 insertions(+), 15 deletions(-) diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index 05716e65137..55b27947faa 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/litellm/llms/openai/chat/guardrail_translation/handler.py @@ -1015,6 +1015,21 @@ class OpenAIChatCompletionsHandler(BaseTranslation): # Subsequent chunks - clear the text content_item["text"] = "" + def _check_streaming_has_ended(self, responses_so_far: Sequence[object]) -> bool: + """ + True once any relayed chunk carries a non-null ``finish_reason``. + + The unified guardrail's ``end_of_stream_only`` streaming path probes + this via ``hasattr`` to withhold the terminal chunks until + end-of-stream moderation runs, so a block can replace the finish + instead of trailing after a ``finish_reason`` the client already saw. + """ + return any( + stream_item_field(choice, "finish_reason") is not None + for item in responses_so_far + for choice in _stream_chunk_choices(item) + ) + def build_block_sse_chunks( self, exc: "ModifyResponseException", @@ -1097,6 +1112,13 @@ def _chat_sse_chunk(payload: _BlockedChunk) -> bytes: return f"data: {json.dumps(payload)}\n\n".encode() +def _stream_chunk_choices(item: object) -> Sequence[object]: + choices: Final = stream_item_field(item, "choices") + if isinstance(choices, Sequence) and not isinstance(choices, (str, bytes)): + return choices + return () + + def _blocked_stream_identity( exc: "ModifyResponseException", responses_so_far: Sequence[object] ) -> tuple[str, int, str]: diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index 811c4feeacc..cf01242f151 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -31,6 +31,7 @@ Output: response.output is List[GenericResponseOutputItem] where each has: import time import uuid from collections.abc import Sequence +from dataclasses import dataclass from typing import TYPE_CHECKING, Any, Final, Union, cast from openai.types.responses.response_function_tool_call import ResponseFunctionToolCall @@ -874,10 +875,10 @@ class OpenAIResponsesHandler(BaseTranslation): sent, so emit the full synthetic sequence (``response.created`` through ``response.completed``). - ``stream_started`` True (sampling / mid-stream): events already - reached the client, so continue the in-progress response: deliver the - block message as a new output item under the same response id and - close with a ``response.completed`` carrying only the replacement - item. + reached the client, so continue the in-progress response: close the + output item still open on the wire, deliver the block message as a + new output item under the same response id, and close with a + ``response.completed`` carrying only the replacement item. The proxy's data generator appends ``data: [DONE]`` itself. """ @@ -887,7 +888,8 @@ class OpenAIResponsesHandler(BaseTranslation): else self._standalone_block_events(exc) ) return tuple( - f"data: {event.model_dump_json(exclude_none=True, exclude_unset=True)}\n\n".encode() for event in events + f"data: {event.model_dump_json(exclude_none=True, exclude_unset=True, serialize_as_any=True)}\n\n".encode() + for event in events ) @staticmethod @@ -915,6 +917,7 @@ class OpenAIResponsesHandler(BaseTranslation): "logprobs": None, } return ( + *_open_item_closing_events(responses_so_far), OutputItemAddedEvent( type=ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, output_index=output_index, @@ -1036,3 +1039,123 @@ def _continuation_identity(exc: "ModifyResponseException", responses_so_far: Seq index for item in responses_so_far if isinstance(index := stream_item_field(item, "output_index"), int) ) return response_id, model, max(indices) + 1 if indices else 0 + + +@dataclass(frozen=True, slots=True) +class _OpenItemState: + item_id: str + item_type: str + role: str + output_index: int + content_index: int + text: str + part_open: bool + + +def _open_item_state(responses_so_far: Sequence[object]) -> _OpenItemState | None: + typed: Final = tuple((str(stream_item_field(event, "type") or ""), event) for event in responses_so_far) + added: Final = tuple( + (added_index, stream_item_field(event, "item")) + for event_type, event in typed + if event_type == "response.output_item.added" + and isinstance(added_index := stream_item_field(event, "output_index"), int) + ) + done_indices: Final = frozenset( + done_index + for event_type, event in typed + if event_type == "response.output_item.done" + and isinstance(done_index := stream_item_field(event, "output_index"), int) + ) + open_added: Final = tuple((index, payload) for index, payload in added if index not in done_indices) + if not open_added: + return None + output_index, item_payload = open_added[-1] + item_id: Final = stream_item_field(item_payload, "id") if item_payload is not None else None + if not isinstance(item_id, str) or not item_id: + return None + raw_type: Final = stream_item_field(item_payload, "type") + raw_role: Final = stream_item_field(item_payload, "role") + part_added: Final = tuple( + part_index + for event_type, event in typed + if event_type == "response.content_part.added" + and stream_item_field(event, "item_id") == item_id + and isinstance(part_index := stream_item_field(event, "content_index"), int) + ) + part_done: Final = frozenset( + part_done_index + for event_type, event in typed + if event_type == "response.content_part.done" + and stream_item_field(event, "item_id") == item_id + and isinstance(part_done_index := stream_item_field(event, "content_index"), int) + ) + open_parts: Final = tuple(index for index in part_added if index not in part_done) + text: Final = "".join( + delta + for event_type, event in typed + if event_type == "response.output_text.delta" + and stream_item_field(event, "item_id") == item_id + and isinstance(delta := stream_item_field(event, "delta"), str) + ) + return _OpenItemState( + item_id=item_id, + item_type=raw_type if isinstance(raw_type, str) and raw_type else "message", + role=raw_role if isinstance(raw_role, str) and raw_role else "assistant", + output_index=output_index, + content_index=open_parts[-1] if open_parts else 0, + text=text, + part_open=bool(open_parts), + ) + + +def _open_item_closing_events(responses_so_far: Sequence[object]) -> Sequence[ResponsesAPIStreamingResponse]: + """Close the output item still in progress on the relayed stream before the + block item is appended: strict Responses clients reject a + ``response.completed`` that arrives while an earlier ``output_item.added`` + was never closed. The closing text is exactly what the client has received + for that item so far.""" + open_item: Final = _open_item_state(responses_so_far) + if open_item is None: + return () + partial_part: Final[_BlockedContentPart] = { + "type": "output_text", + "text": open_item.text, + "annotations": (), + } + closed_payload: Final[_BlockedItemPayload] = { + "type": open_item.item_type, + "id": open_item.item_id, + "status": "completed", + "role": open_item.role, + "content": (partial_part,), + } + item_done: Final = OutputItemDoneEvent( + type=ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE, + output_index=open_item.output_index, + item=GenericResponseOutputItem.model_validate(closed_payload), + ) + if not open_item.part_open: + return (item_done,) + partial_done_part: Final[_BlockedDoneContentPart] = { + "type": "output_text", + "text": open_item.text, + "annotations": (), + "logprobs": None, + } + return ( + OutputTextDoneEvent( + type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DONE, + item_id=open_item.item_id, + output_index=open_item.output_index, + content_index=open_item.content_index, + text=open_item.text, + ), + ContentPartDoneEvent( + type=ResponsesAPIStreamEvents.CONTENT_PART_DONE, + item_id=open_item.item_id, + output_index=open_item.output_index, + content_index=open_item.content_index, + part=ContentPartDonePartOutputText.model_validate(partial_done_part), + ), + item_done, + ) diff --git a/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py b/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py index 8ad0026d1ab..7dd6065063a 100644 --- a/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py +++ b/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py @@ -1610,3 +1610,36 @@ class TestBuildBlockSseChunks: assert first["choices"][0]["delta"] == {"content": "Blocked by policy."} assert final["id"] == "chatcmpl-live" assert final["usage"] == {"prompt_tokens": 11, "completion_tokens": 5, "total_tokens": 16} + + +class TestCheckStreamingHasEnded: + """_check_streaming_has_ended lets end_of_stream_only withhold the finish chunk until moderation""" + + def test_empty_and_content_only_chunks_are_not_ended(self): + handler = OpenAIChatCompletionsHandler() + assert handler._check_streaming_has_ended([]) is False + content_only = [ + {"id": "chatcmpl-live", "choices": [{"index": 0, "delta": {"content": "hi"}, "finish_reason": None}]}, + {"id": "chatcmpl-live", "choices": []}, + {"id": "chatcmpl-live", "usage": {"prompt_tokens": 1, "completion_tokens": 1}}, + ] + assert handler._check_streaming_has_ended(content_only) is False + + def test_dict_finish_chunk_marks_stream_ended(self): + handler = OpenAIChatCompletionsHandler() + chunks = [ + {"id": "chatcmpl-live", "choices": [{"index": 0, "delta": {"content": "hi"}, "finish_reason": None}]}, + {"id": "chatcmpl-live", "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}]}, + ] + assert handler._check_streaming_has_ended(chunks) is True + + def test_object_finish_chunk_marks_stream_ended(self): + from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices + + handler = OpenAIChatCompletionsHandler() + chunks = [ + ModelResponseStream( + choices=[StreamingChoices(index=0, delta=Delta(content=None), finish_reason="stop")] + ) + ] + assert handler._check_streaming_has_ended(chunks) is True diff --git a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py index cd0e7f29933..1ae4e7a699b 100644 --- a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py +++ b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py @@ -1285,10 +1285,32 @@ class TestBuildBlockSseChunks: ) types = [payload["type"] for payload in payloads] assert "response.created" not in types - assert types[0] == "response.output_item.added" - assert payloads[0]["output_index"] == 3 + assert types[0] == "response.output_item.done" + assert payloads[0]["output_index"] == 2 + assert payloads[0]["item"]["id"] == "msg_orig" + assert payloads[0]["item"]["status"] == "completed" + assert types[1] == "response.output_item.added" + assert payloads[1]["output_index"] == 3 completed = payloads[-1]["response"] assert completed["id"] == "resp_live" assert completed["model"] == "gpt-5.4-mini-2026-01-01" assert completed["output"][0]["content"][0]["text"] == "Blocked by policy." assert completed["usage"] == {"input_tokens": 7, "output_tokens": 21, "total_tokens": 28} + + def test_continuation_without_open_item_emits_no_closing_events(self): + handler = OpenAIResponsesHandler() + yielded = [ + {"type": "response.created", "response": {"id": "resp_live", "model": "gpt-5.4-mini"}}, + {"type": "response.in_progress", "response": {"id": "resp_live"}}, + ] + payloads = self._payloads( + handler.build_block_sse_chunks( + self._exc(original_response=yielded), stream_started=True, responses_so_far=yielded + ) + ) + types = [payload["type"] for payload in payloads] + assert types[0] == "response.output_item.added" + assert types[-1] == "response.completed" + dones = [payload for payload in payloads if payload["type"] == "response.output_item.done"] + assert len(dones) == 1 + assert dones[0]["item"]["content"][0]["text"] == "Blocked by policy." diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_openai_streaming_block.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_openai_streaming_block.py index 99172ec3be6..42bdf41bc88 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_openai_streaming_block.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_openai_streaming_block.py @@ -53,6 +53,19 @@ class _BlockingGuardrail(CustomGuardrail): ) +class _PassingGuardrail(CustomGuardrail): + """Mock guardrail that always lets response scans through unchanged.""" + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[Any] = None, + ) -> GenericGuardrailAPIInputs: + return inputs + + def _chat_chunk(delta: Delta, finish_reason: Optional[str] = None) -> ModelResponseStream: return ModelResponseStream( id="chatcmpl-live", @@ -129,8 +142,13 @@ async def _run_hook( sampling_rate: int = 1, end_of_stream_only: bool = False, buffer_until_moderated: bool = False, + blocks: bool = True, ) -> Tuple[StreamChunk, ...]: - guardrail = _BlockingGuardrail(guardrail_name="test-blocking-guardrail", event_hook="post_call") + guardrail = ( + _BlockingGuardrail(guardrail_name="test-blocking-guardrail", event_hook="post_call") + if blocks + else _PassingGuardrail(guardrail_name="test-passing-guardrail", event_hook="post_call") + ) guardrail.streaming_sampling_rate = sampling_rate guardrail.streaming_end_of_stream_only = end_of_stream_only guardrail.streaming_buffer_until_moderated = buffer_until_moderated @@ -140,7 +158,7 @@ async def _run_hook( request_data = { "messages": [{"role": "user", "content": "hi"}], "guardrail_to_apply": guardrail, - "metadata": {"guardrails": ["test-blocking-guardrail"]}, + "metadata": {"guardrails": [guardrail.guardrail_name]}, } return tuple( @@ -203,13 +221,38 @@ async def test_chat_mid_stream_block_continues_the_completion(): @pytest.mark.asyncio async def test_chat_end_of_stream_block_terminates_cleanly(): + """Regression for bugbot's finish-ordering finding: in end_of_stream_only + mode the original finish chunk must be withheld until moderation decides, + so a block's content_filter finish is the only stream terminator a client + ever sees - never policy text trailing after finish_reason stop.""" collected = await _run_hook("/v1/chat/completions", _chat_stream(end=True), end_of_stream_only=True) _assert_no_error_frame(collected) + forwarded = [chunk for chunk in collected if isinstance(chunk, ModelResponseStream)] + assert forwarded, "content chunks still stream to the client before end-of-stream moderation" + assert all(choice.finish_reason is None for chunk in forwarded for choice in chunk.choices), ( + "the original finish chunk must be withheld until moderation decides" + ) payloads = _sse_payloads(collected) assert BLOCK_MESSAGE in json.dumps(payloads) assert payloads[-1]["choices"][0]["finish_reason"] == "content_filter" +@pytest.mark.asyncio +async def test_chat_end_of_stream_pass_releases_withheld_finish_chunk(): + """When end-of-stream moderation passes, the withheld finish chunk is + released so a clean stream still terminates normally.""" + collected = await _run_hook( + "/v1/chat/completions", _chat_stream(end=True), end_of_stream_only=True, blocks=False + ) + assert not [chunk for chunk in collected if isinstance(chunk, bytes)], ( + "a clean stream must carry no synthetic block frames" + ) + forwarded = [chunk for chunk in collected if isinstance(chunk, ModelResponseStream)] + finish_reasons = [choice.finish_reason for chunk in forwarded for choice in chunk.choices] + assert finish_reasons[-1] == "stop", "the withheld finish chunk must be released after moderation passes" + assert all(reason is None for reason in finish_reasons[:-1]) + + @pytest.mark.asyncio async def test_responses_buffered_block_emits_full_event_sequence(): """Buffered moderation blocks before anything streams: a complete synthetic @@ -234,20 +277,35 @@ async def test_responses_buffered_block_emits_full_event_sequence(): @pytest.mark.asyncio async def test_responses_mid_stream_block_continues_the_response(): - """Regression for the LIT-6496 500 error frame: after events were already - forwarded, the block appends a new output item under the same response id - and closes with response.completed - never a second response.created.""" + """Regression for the LIT-6496 500 error frame and bugbot's unclosed-item + finding: after events were already forwarded, the block first closes the + output item still open on the wire, then appends the replacement item under + the same response id, and closes with response.completed - never a second + response.created and never a completed response with an item left open.""" collected = await _run_hook("/v1/responses", _responses_stream(end=False)) _assert_no_error_frame(collected) - forwarded_types = [chunk["type"] for chunk in collected if isinstance(chunk, dict)] + forwarded = [chunk for chunk in collected if isinstance(chunk, dict)] + forwarded_types = [chunk["type"] for chunk in forwarded] assert "response.created" in forwarded_types, "original events should have streamed before the block" payloads = _sse_payloads(collected) assert payloads, "no block SSE chunks were emitted" block_types = [payload["type"] for payload in payloads] assert "response.created" not in block_types, "a mid-stream block must not restart the response" - assert block_types[0] == "response.output_item.added" assert block_types[-1] == "response.completed" - assert payloads[0]["output_index"] == 1, "the block item must continue after the original output item" + + all_events = forwarded + list(payloads) + opened = sorted(event["output_index"] for event in all_events if event["type"] == "response.output_item.added") + closed = sorted(event["output_index"] for event in all_events if event["type"] == "response.output_item.done") + assert opened == closed, "every output item opened on the stream must be closed before response.completed" + original_done_position = block_types.index("response.output_item.done") + block_item_position = block_types.index("response.output_item.added") + assert original_done_position < block_item_position, ( + "the in-progress original item must be closed before the block item is appended" + ) + assert payloads[original_done_position]["item"]["id"] == "msg_orig" + assert payloads[block_item_position]["output_index"] == 1, ( + "the block item must continue after the original output item" + ) completed = payloads[-1]["response"] assert completed["id"] == "resp_live" assert completed["output"][0]["content"][0]["text"] == BLOCK_MESSAGE From 38825cf9c62eeb3c374b87a4ea9b2c31cc9f19cc Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 31 Aug 2026 17:19:55 -0700 Subject: [PATCH 07/97] fix(guardrails): match Responses stream event types by value so enum-typed events close the open item --- .../guardrail_translation/handler.py | 2 +- ...test_openai_responses_guardrail_handler.py | 59 +++++++++++++++++++ 2 files changed, 60 insertions(+), 1 deletion(-) diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index cf01242f151..6295df1dbfa 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -1053,7 +1053,7 @@ class _OpenItemState: def _open_item_state(responses_so_far: Sequence[object]) -> _OpenItemState | None: - typed: Final = tuple((str(stream_item_field(event, "type") or ""), event) for event in responses_so_far) + typed: Final = tuple((stream_item_field(event, "type"), event) for event in responses_so_far) added: Final = tuple( (added_index, stream_item_field(event, "item")) for event_type, event in typed diff --git a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py index 1ae4e7a699b..d6e56f0faf1 100644 --- a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py +++ b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py @@ -1297,6 +1297,65 @@ class TestBuildBlockSseChunks: assert completed["output"][0]["content"][0]["text"] == "Blocked by policy." assert completed["usage"] == {"input_tokens": 7, "output_tokens": 21, "total_tokens": 28} + def test_continuation_closes_open_item_given_pydantic_events_with_enum_types(self): + from litellm.types.llms.openai import ( + BaseLiteLLMOpenAIResponseObject, + ContentPartAddedEvent, + OutputItemAddedEvent, + OutputTextDeltaEvent, + ResponsesAPIStreamEvents, + ) + + handler = OpenAIResponsesHandler() + open_item = GenericResponseOutputItem.model_validate( + {"type": "message", "id": "msg_live", "status": "in_progress", "role": "assistant", "content": []} + ) + yielded = [ + OutputItemAddedEvent( + type=ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, output_index=0, item=open_item + ), + ContentPartAddedEvent( + type=ResponsesAPIStreamEvents.CONTENT_PART_ADDED, + item_id="msg_live", + output_index=0, + content_index=0, + part=BaseLiteLLMOpenAIResponseObject.model_validate( + {"type": "output_text", "text": "", "annotations": []} + ), + ), + OutputTextDeltaEvent( + type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA, + item_id="msg_live", + output_index=0, + content_index=0, + delta="partial ", + ), + OutputTextDeltaEvent( + type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA, + item_id="msg_live", + output_index=0, + content_index=0, + delta="text", + ), + ] + payloads = self._payloads( + handler.build_block_sse_chunks( + self._exc(original_response=yielded), stream_started=True, responses_so_far=yielded + ) + ) + types = [payload["type"] for payload in payloads] + assert types[:3] == [ + "response.output_text.done", + "response.content_part.done", + "response.output_item.done", + ] + assert payloads[0]["text"] == "partial text" + assert payloads[2]["item"]["id"] == "msg_live" + assert payloads[2]["item"]["status"] == "completed" + assert payloads[2]["item"]["content"][0]["text"] == "partial text" + assert types[3] == "response.output_item.added" + assert payloads[3]["output_index"] == 1 + def test_continuation_without_open_item_emits_no_closing_events(self): handler = OpenAIResponsesHandler() yielded = [ From 329654765a9d07a3f8d15abc6a86040044668473 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 31 Aug 2026 17:31:59 -0700 Subject: [PATCH 08/97] test(guardrails): update chat eos block tests for finish-chunk withholding --- .../guardrails/guardrail_hooks/test_bedrock_guardrails.py | 7 +++++-- .../unified_guardrails/test_unified_guardrail.py | 3 ++- 2 files changed, 7 insertions(+), 3 deletions(-) diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py index 235a5c0c09b..bacaa1c3e22 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py @@ -5524,7 +5524,9 @@ async def test_streaming_end_of_stream_block_emits_error_frame_instead_of_trunca """Regression for PR #38722: a topicPolicy DENY caught by the end-of-stream scan used to raise after SSE headers were flushed, so the client saw a silently truncated stream. The unified hook must emit the chat in-stream - error frame instead.""" + error frame instead. The finish chunk is withheld while the end-of-stream + scan runs, so on a block it is dropped rather than relayed before the + frame.""" from litellm.llms import load_guardrail_translation_mappings from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail import ( unified_guardrail as unified_module, @@ -5582,8 +5584,9 @@ async def test_streaming_end_of_stream_block_emits_error_frame_instead_of_trunca finally: unified_module.endpoint_guardrail_translation_mappings = None - assert len(out) == 3 + assert len(out) == 2 assert isinstance(out[0], ModelResponseStream) + assert out[0].choices[0].finish_reason is None frame = out[-1] assert isinstance(frame, bytes) payload = json.loads(frame.decode()[len("data: ") :]) diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py index 0b32558a00a..8cad1c634a9 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py @@ -1844,7 +1844,8 @@ class TestStreamingHttpErrorFrames: out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks) - assert out[:2] == chunks + assert out[0] == chunks[0] + assert chunks[1] not in out frame = out[-1] assert isinstance(frame, bytes) text = frame.decode() 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 09/97] 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 10/97] 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 c21e895fe26b16cc97f0d4ed8e389f6ede9f2688 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 1 Sep 2026 12:16:27 -0700 Subject: [PATCH 11/97] fix(proxy): handle CRLF and CR SSE frame terminators and flush held tail in anthropic stream restamper --- .../streaming_model_restamp.py | 49 ++++++--- litellm/proxy/common_request_processing.py | 21 ++-- .../test_streaming_model_restamp.py | 100 +++++++++++++++++- 3 files changed, 144 insertions(+), 26 deletions(-) diff --git a/litellm/proxy/anthropic_endpoints/streaming_model_restamp.py b/litellm/proxy/anthropic_endpoints/streaming_model_restamp.py index e8d54f03949..7da5e5099fc 100644 --- a/litellm/proxy/anthropic_endpoints/streaming_model_restamp.py +++ b/litellm/proxy/anthropic_endpoints/streaming_model_restamp.py @@ -8,6 +8,7 @@ provider passthrough path) or as event dicts (fake-stream and agentic paths). """ import json +import re from collections.abc import Mapping from typing import Final @@ -16,7 +17,7 @@ from pydantic import TypeAdapter, ValidationError _MESSAGE_START_EVENT: Final = "message_start" _MESSAGE_START_MARKER: Final = b"message_start" _SSE_DATA_FIELD: Final = "data:" -_SSE_FRAME_END: Final = b"\n\n" +_SSE_FRAME_END_PATTERN: Final = re.compile(rb"\r\n\r\n|\r\r|\n\n") _MAX_HELD_BYTES: Final = 65536 _PING_MARKERS: Final = (b"event: ping", b'"type": "ping"', b'"type":"ping"') @@ -46,15 +47,16 @@ def _restamped_data_line(line: str, requested_model: str) -> str | None: restamped: Final = _restamped_event(event, requested_model) if restamped is None: return None - return f"data: {json.dumps(restamped, separators=(',', ':'))}" + terminator: Final = line[len(line.rstrip("\r\n")) :] + return f"data: {json.dumps(restamped, separators=(',', ':'))}{terminator}" def _restamped_frame(frame: str, requested_model: str) -> str | None: - lines: Final = frame.split("\n") + lines: Final = frame.splitlines(keepends=True) restamped: Final = tuple(_restamped_data_line(line, requested_model) for line in lines) if all(line is None for line in restamped): return None - return "\n".join(new if new is not None else old for new, old in zip(restamped, lines)) + return "".join(new if new is not None else old for new, old in zip(restamped, lines)) def restamp_anthropic_stream_chunk_model(chunk: object, requested_model: str) -> object: @@ -93,10 +95,12 @@ class AnthropicStreamModelRestamper: """ Per-stream restamper for the encoded passthrough path, where chunks are raw transport reads: the ``message_start`` SSE frame can arrive split across - chunks or coalesced with later frames. Complete frames are emitted as their - terminator closes them and an incomplete tail is held until it completes, - so the restamp never misses a torn frame. Once ``message_start`` has been - handled, or the first real event proves the stream carries none, every + chunks or coalesced with later frames. Complete frames (``\\n\\n``, + ``\\r\\n\\r\\n``, or ``\\r\\r`` terminated) are emitted as their terminator + closes them and an incomplete tail is held until it completes, so the + restamp never misses a torn frame; ``flush`` returns whatever is still held + when the stream ends so no bytes are swallowed. Once ``message_start`` has + been handled, or the first real event proves the stream carries none, every later chunk passes through untouched. """ @@ -117,17 +121,27 @@ class AnthropicStreamModelRestamper: self._armed = False return restamped + def flush(self) -> bytes: + held: Final = self._held + self._held = b"" + self._armed = False + if not held: + return b"" + restamped: Final = restamp_anthropic_stream_chunk_model(held, self._requested_model) + return restamped if isinstance(restamped, bytes) else held + def _process_encoded(self, data: bytes) -> bytes: combined: Final = self._held + data - if _SSE_FRAME_END not in combined: + boundaries: Final = tuple(match.end() for match in _SSE_FRAME_END_PATTERN.finditer(combined)) + if not boundaries: if len(combined) > _MAX_HELD_BYTES: self._held = b"" self._armed = False return combined self._held = combined return b"" - closed, _, tail = combined.rpartition(_SSE_FRAME_END) - emitted: Final = self._restamped_closed_block(closed + _SSE_FRAME_END) + emitted: Final = self._restamped_closed_block(combined[: boundaries[-1]]) + tail: Final = combined[boundaries[-1] :] if not self._armed: self._held = b"" return emitted + tail @@ -135,7 +149,8 @@ class AnthropicStreamModelRestamper: return emitted def _restamped_closed_block(self, closed: bytes) -> bytes: - frames: Final = tuple(closed.split(_SSE_FRAME_END)[:-1]) + boundaries: Final = tuple(match.end() for match in _SSE_FRAME_END_PATTERN.finditer(closed)) + frames: Final = tuple(closed[start:end] for start, end in zip((0, *boundaries[:-1]), boundaries)) decider: Final = next( ( index @@ -147,13 +162,13 @@ class AnthropicStreamModelRestamper: if decider is None: return closed self._armed = False - decider_frame: Final = frames[decider] + _SSE_FRAME_END - if _MESSAGE_START_MARKER not in decider_frame: + if _MESSAGE_START_MARKER not in frames[decider]: return closed - restamped_text: Final = _restamped_frame(decider_frame.decode("utf-8", errors="ignore"), self._requested_model) + restamped_text: Final = _restamped_frame( + frames[decider].decode("utf-8", errors="ignore"), self._requested_model + ) if restamped_text is None: return closed return b"".join( - restamped_text.encode("utf-8") if index == decider else frame + _SSE_FRAME_END - for index, frame in enumerate(frames) + restamped_text.encode("utf-8") if index == decider else frame for index, frame in enumerate(frames) ) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 989cc7c18fb..eda7ebfa624 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -3410,12 +3410,10 @@ class ProxyBaseLLMRequestProcessing: return chunk @staticmethod - def _sse_chunk_serializer(restamp_model: str | None) -> StreamChunkSerializer: - if not restamp_model: + def _sse_chunk_serializer(restamper: AnthropicStreamModelRestamper | None) -> StreamChunkSerializer: + if restamper is None: return ProxyBaseLLMRequestProcessing.return_sse_chunk - restamper: Final = AnthropicStreamModelRestamper(restamp_model) - def serialize(chunk: object) -> str: return ProxyBaseLLMRequestProcessing.return_sse_chunk(restamper.process(chunk)) @@ -3481,11 +3479,16 @@ class ProxyBaseLLMRequestProcessing: serialize_chunk: StreamChunkSerializer, serialize_error: StreamErrorSerializer, request: Request | None = None, + flush_tail: Callable[[], bytes] | None = None, ) -> AsyncGenerator[str, None]: """ Shared streaming data generator: runs proxy iterator hook, per-chunk hook, cost injection, then yields chunks via serialize_chunk; on exception runs failure hook and yields via serialize_error. Use for SSE or NDJSON. + + ``flush_tail`` runs once after the upstream iterator completes cleanly and + its non-empty result is yielded, so a serializer that buffers bytes across + chunks can emit anything still held at end of stream. """ verbose_proxy_logger.debug("inside generator") # Resolve per-stream (not per-chunk) whether the heavy per-chunk path @@ -3548,6 +3551,9 @@ class ProxyBaseLLMRequestProcessing: # so it must not suppress that refund. delivered_chunk = delivered_chunk or chunk != STREAM_SSE_KEEPALIVE_PING_BYTES yield serialize_chunk(chunk) + held_tail: Final = flush_tail() if flush_tail is not None else b"" + if held_tail: + yield serialize_chunk(held_tail) stream_completed = True except (asyncio.CancelledError, GeneratorExit): # Client disconnected mid-stream. CancelledError / GeneratorExit @@ -3558,8 +3564,7 @@ class ProxyBaseLLMRequestProcessing: # billing and release exactly once. This is the outermost generator # Starlette closes on disconnect, so the nested iterator hook (which # only sees GeneratorExit on GC) cannot own the refund. - if not stream_completed: - client_disconnected = True + client_disconnected = not stream_completed if not delivered_chunk and not _withheld_provider_output(response): from litellm.proxy.spend_tracking.budget_reservation import ( release_budget_reservation_on_cancel, @@ -3627,16 +3632,18 @@ class ProxyBaseLLMRequestProcessing: event in place of the provider's model, matching what the non-streaming response reports. """ + restamper: Final = AnthropicStreamModelRestamper(restamp_model) if restamp_model else None return ProxyBaseLLMRequestProcessing.async_streaming_data_generator( response=response, user_api_key_dict=user_api_key_dict, request_data=request_data, proxy_logging_obj=proxy_logging_obj, - serialize_chunk=ProxyBaseLLMRequestProcessing._sse_chunk_serializer(restamp_model), + serialize_chunk=ProxyBaseLLMRequestProcessing._sse_chunk_serializer(restamper), serialize_error=lambda proxy_exc: ( f"{STREAM_SSE_DATA_PREFIX}{json.dumps({'error': proxy_exc.to_dict()})}\n\n" ), request=request, + flush_tail=None if restamper is None else restamper.flush, ) @overload diff --git a/tests/test_litellm/proxy/anthropic_endpoints/test_streaming_model_restamp.py b/tests/test_litellm/proxy/anthropic_endpoints/test_streaming_model_restamp.py index 385173b24a9..b7bc670c7f8 100644 --- a/tests/test_litellm/proxy/anthropic_endpoints/test_streaming_model_restamp.py +++ b/tests/test_litellm/proxy/anthropic_endpoints/test_streaming_model_restamp.py @@ -14,12 +14,12 @@ from litellm.proxy.anthropic_endpoints.streaming_model_restamp import ( from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing -def _message_start_frame(model: str) -> bytes: +def _message_start_frame(model: str, line_end: str = "\n") -> bytes: payload = { "type": "message_start", "message": {"id": "msg_1", "type": "message", "role": "assistant", "model": model, "content": []}, } - return f"event: message_start\ndata: {json.dumps(payload)}\n\n".encode() + return f"event: message_start{line_end}data: {json.dumps(payload)}{line_end}{line_end}".encode() def _proxy_logging_obj_streaming(frames: list[bytes]) -> MagicMock: @@ -190,3 +190,99 @@ async def test_sse_generator_restamps_message_start_split_across_chunks(): joined = b"".join(chunk if isinstance(chunk, bytes) else chunk.encode("utf-8") for chunk in chunks) assert _model_from_frame(joined) == "claude-auto-1" + + +def test_restamps_crlf_terminated_message_start_frame(): + frame = _message_start_frame("claude-haiku-4-5-20251001", line_end="\r\n") + delta = b'event: content_block_delta\r\ndata: {"type":"content_block_delta","delta":{"text":"hi"}}\r\n\r\n' + restamper = AnthropicStreamModelRestamper("claude-auto-1") + + emitted = restamper.process(frame) + + assert isinstance(emitted, bytes) + assert _model_from_frame(emitted) == "claude-auto-1" + assert emitted.endswith(b"\r\n\r\n") + assert restamper.process(delta) == delta + + +def test_restamps_cr_terminated_message_start_frame(): + frame = _message_start_frame("claude-haiku-4-5-20251001", line_end="\r") + restamper = AnthropicStreamModelRestamper("claude-auto-1") + + emitted = restamper.process(frame) + + assert isinstance(emitted, bytes) + assert b'"model":"claude-auto-1"' in emitted + assert emitted.endswith(b"\r\r") + + +def test_restamps_crlf_message_start_split_across_transport_chunks(): + frame = _message_start_frame("claude-haiku-4-5-20251001", line_end="\r\n") + restamper = AnthropicStreamModelRestamper("claude-auto-1") + + held = restamper.process(frame[:25]) + emitted = restamper.process(frame[25:]) + + assert held == b"" + assert isinstance(emitted, bytes) + assert _model_from_frame(emitted) == "claude-auto-1" + + +def test_flush_returns_restamped_held_tail(): + unterminated = _message_start_frame("claude-haiku-4-5-20251001")[:-2] + restamper = AnthropicStreamModelRestamper("claude-auto-1") + + assert restamper.process(unterminated) == b"" + flushed = restamper.flush() + + assert b'"model":"claude-auto-1"' in flushed + assert restamper.flush() == b"" + + +def test_flush_disarms_the_restamper(): + restamper = AnthropicStreamModelRestamper("claude-auto-1") + frame = _message_start_frame("claude-haiku-4-5-20251001") + + assert restamper.flush() == b"" + assert restamper.process(frame) == frame + + +@pytest.mark.asyncio +async def test_sse_generator_flushes_held_tail_at_end_of_stream(): + unterminated = _message_start_frame("claude-haiku-4-5-20251001")[:-2] + proxy_logging_obj = _proxy_logging_obj_streaming([unterminated]) + + chunks = [ + chunk + async for chunk in ProxyBaseLLMRequestProcessing.async_sse_data_generator( + response=MagicMock(), + user_api_key_dict=MagicMock(), + request_data={"model": "claude-auto-1"}, + proxy_logging_obj=proxy_logging_obj, + restamp_model="claude-auto-1", + ) + ] + + joined = b"".join(chunk if isinstance(chunk, bytes) else chunk.encode("utf-8") for chunk in chunks) + assert b'"model":"claude-auto-1"' in joined + + +@pytest.mark.asyncio +async def test_sse_generator_restamps_crlf_stream(): + frame = _message_start_frame("claude-haiku-4-5-20251001", line_end="\r\n") + delta = b'event: content_block_delta\r\ndata: {"type":"content_block_delta","delta":{"text":"hi"}}\r\n\r\n' + proxy_logging_obj = _proxy_logging_obj_streaming([frame, delta]) + + chunks = [ + chunk + async for chunk in ProxyBaseLLMRequestProcessing.async_sse_data_generator( + response=MagicMock(), + user_api_key_dict=MagicMock(), + request_data={"model": "claude-auto-1"}, + proxy_logging_obj=proxy_logging_obj, + restamp_model="claude-auto-1", + ) + ] + + assert _model_from_frame(chunks[0]) == "claude-auto-1" + assert chunks[1] == delta 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 12/97] 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 13/97] 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 14/97] 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 4567fc784c6fe43acd41a6db32c834465b8d653d Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 1 Sep 2026 12:45:23 -0700 Subject: [PATCH 15/97] fix: close non-message open items as incomplete when a stream is blocked --- .../guardrail_translation/handler.py | 38 +++++++++++++++--- ...test_openai_responses_guardrail_handler.py | 40 +++++++++++++++++++ 2 files changed, 73 insertions(+), 5 deletions(-) diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index 6295df1dbfa..14759f6475e 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -30,13 +30,14 @@ Output: response.output is List[GenericResponseOutputItem] where each has: import time import uuid -from collections.abc import Sequence +from collections.abc import Mapping, Sequence from dataclasses import dataclass +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Union, cast from openai.types.responses.response_function_tool_call import ResponseFunctionToolCall from openai.types.responses.tool_param import FunctionToolParam -from pydantic import BaseModel +from pydantic import BaseModel, TypeAdapter from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_proxy_logger @@ -1050,6 +1051,7 @@ class _OpenItemState: content_index: int text: str part_open: bool + payload: object def _open_item_state(responses_so_far: Sequence[object]) -> _OpenItemState | None: @@ -1070,7 +1072,9 @@ def _open_item_state(responses_so_far: Sequence[object]) -> _OpenItemState | Non if not open_added: return None output_index, item_payload = open_added[-1] - item_id: Final = stream_item_field(item_payload, "id") if item_payload is not None else None + if item_payload is None: + return None + item_id: Final = stream_item_field(item_payload, "id") if not isinstance(item_id, str) or not item_id: return None raw_type: Final = stream_item_field(item_payload, "type") @@ -1105,18 +1109,42 @@ def _open_item_state(responses_so_far: Sequence[object]) -> _OpenItemState | Non content_index=open_parts[-1] if open_parts else 0, text=text, part_open=bool(open_parts), + payload=item_payload, ) +_item_fields_adapter: Final = TypeAdapter(Mapping[str, object]) +_no_item_fields: Final[Mapping[str, object]] = MappingProxyType({}) + + +def _incomplete_item_fields(payload: object) -> Mapping[str, object]: + raw: Final = payload.model_dump() if isinstance(payload, BaseModel) else payload + if not isinstance(raw, dict): + return _no_item_fields + return _item_fields_adapter.validate_python(raw) + + def _open_item_closing_events(responses_so_far: Sequence[object]) -> Sequence[ResponsesAPIStreamingResponse]: """Close the output item still in progress on the relayed stream before the block item is appended: strict Responses clients reject a ``response.completed`` that arrives while an earlier ``output_item.added`` - was never closed. The closing text is exactly what the client has received - for that item so far.""" + was never closed. A message item closes ``completed`` with exactly the text + the client has received so far; any other item type (a function call the + guardrail rejected, for instance) closes ``incomplete`` so the synthetic + done event can never authorize acting on it.""" open_item: Final = _open_item_state(responses_so_far) if open_item is None: return () + if open_item.item_type != "message": + return ( + OutputItemDoneEvent( + type=ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE, + output_index=open_item.output_index, + item=BaseLiteLLMOpenAIResponseObject.model_validate( + MappingProxyType({**_incomplete_item_fields(open_item.payload), "status": "incomplete"}) + ), + ), + ) partial_part: Final[_BlockedContentPart] = { "type": "output_text", "text": open_item.text, diff --git a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py index d6e56f0faf1..49f4e9df7c9 100644 --- a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py +++ b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py @@ -1356,6 +1356,46 @@ class TestBuildBlockSseChunks: assert types[3] == "response.output_item.added" assert payloads[3]["output_index"] == 1 + def test_continuation_closes_open_function_call_as_incomplete(self): + handler = OpenAIResponsesHandler() + yielded = [ + {"type": "response.created", "response": {"id": "resp_live", "model": "gpt-5.4-mini"}}, + { + "type": "response.output_item.added", + "output_index": 0, + "item": { + "id": "fc_live", + "type": "function_call", + "status": "in_progress", + "call_id": "call_1", + "name": "run_payment", + "arguments": "", + }, + }, + { + "type": "response.function_call_arguments.delta", + "item_id": "fc_live", + "output_index": 0, + "delta": '{"amount": 100}', + }, + ] + payloads = self._payloads( + handler.build_block_sse_chunks( + self._exc(original_response=yielded), stream_started=True, responses_so_far=yielded + ) + ) + types = [payload["type"] for payload in payloads] + assert types[0] == "response.output_item.done" + closed = payloads[0]["item"] + assert closed["id"] == "fc_live" + assert closed["type"] == "function_call" + assert closed["status"] == "incomplete" + assert closed["name"] == "run_payment" + assert "content" not in closed + assert types[1] == "response.output_item.added" + assert payloads[1]["output_index"] == 1 + assert types[-1] == "response.completed" + def test_continuation_without_open_item_emits_no_closing_events(self): handler = OpenAIResponsesHandler() yielded = [ 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 16/97] 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 17/97] 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 18/97] 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 19/97] 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 20/97] 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 21/97] 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 22/97] 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 23/97] 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 4da12795fc5f90cd4e5e87fc61eb080c792b4355 Mon Sep 17 00:00:00 2001 From: mateo Date: Tue, 1 Sep 2026 21:50:45 +0000 Subject: [PATCH 24/97] fix: filter deployment default API key limits Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/types/utils.py | 2 + .../chat/test_anthropic_chat_handler.py | 39 +++++++++++++++++++ 2 files changed, 41 insertions(+) diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 3a0883b6607..5783a39b30c 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -3636,6 +3636,8 @@ all_litellm_params = ( "client", "rpm", "tpm", + "default_api_key_rpm_limit", + "default_api_key_tpm_limit", "itpm", "otpm", "max_parallel_requests", diff --git a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_handler.py b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_handler.py index bd750a47f63..043537f8c1f 100644 --- a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_handler.py +++ b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_handler.py @@ -3,11 +3,13 @@ import threading from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from unittest.mock import AsyncMock, MagicMock, patch +import httpx import pytest import litellm from litellm.constants import RESPONSE_FORMAT_TOOL_NAME from litellm.llms.anthropic.chat.handler import ModelResponseIterator, make_call +from litellm.llms.custom_httpx.http_handler import HTTPHandler from litellm.types.llms.openai import ( ChatCompletionToolCallChunk, ChatCompletionToolCallFunctionChunk, @@ -46,6 +48,43 @@ async def test_make_call_passes_logging_obj_to_client_post(): assert call_kwargs.get("logging_obj") is logging_obj +def test_anthropic_completion_does_not_send_deployment_default_limits(): + captured_requests: list[httpx.Request] = [] + + def respond(request: httpx.Request) -> httpx.Response: + captured_requests.append(request) + return httpx.Response( + 200, + json={ + "id": "msg_default_limits", + "type": "message", + "role": "assistant", + "model": "claude-3-5-haiku-20241022", + "content": [{"type": "text", "text": "Hello"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 1, "output_tokens": 1}, + }, + ) + + client = HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(respond))) + try: + litellm.completion( + model="anthropic/claude-3-5-haiku-20241022", + messages=[{"role": "user", "content": "Hello"}], + api_key="test-key", + client=client, + default_api_key_rpm_limit=60, + default_api_key_tpm_limit=5000000, + ) + finally: + client.close() + + request_body = json.loads(captured_requests[0].content) + assert "default_api_key_rpm_limit" not in request_body + assert "default_api_key_tpm_limit" not in request_body + + def test_redacted_thinking_content_block_delta(): chunk = { "type": "content_block_start", From af11db9fe571c9d10e9175ef58cceb0d6301f8d2 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Tue, 1 Sep 2026 14:51:15 -0700 Subject: [PATCH 25/97] test(e2e): cover retry-on-timeout and the context-window fallback Two P0 rows in the reliability coverage registry had no test. reliability.retry.timeout.succeeds_within_retries gets a new file. The model group is a pair: an always-timing-out deployment holding all of the group's shuffle weight, and a healthy backup at weight 0. The weighted pick always opens on the timing-out one, its first Timeout benches it via an allowed_fails_policy of TimeoutErrorAllowedFails 0, and the retry falls through to the only deployment left, so the outcome is a completion plus a reported retry with no random first pick in the middle. reliability.fallback.context_window.routes_to_fallback joins the existing fallbacks spec. It registers a genuinely small-context OpenAI deployment, sends a prompt past its limit so the provider refuses it on length, and reroutes with context_window_fallbacks, which is the setting that handles that refusal rather than plain fallbacks. Both drive real provider calls through router_settings_override, so no config change and no second proxy is needed. Reliability & Performance goes 16/36 to 18/36. Claude-Session: https://claude.ai/code/session_01QvQzYztinxj8ZuD5YxbVdL --- tests/e2e/models.py | 2 + tests/e2e/router/reliability_support.py | 46 ++++++++++++ .../router/test_reliability_fallbacks_e2e.py | 20 +++++ .../router/test_reliability_retries_e2e.py | 73 +++++++++++++++++++ 4 files changed, 141 insertions(+) create mode 100644 tests/e2e/router/test_reliability_retries_e2e.py diff --git a/tests/e2e/models.py b/tests/e2e/models.py index 6d9ccad9a24..b48a37f16ae 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -805,6 +805,7 @@ class LiteLLMParamsBody(BaseModel): mock_response: str | None = None timeout: float | None = None tpm: int | None = None + weight: int | None = None ModelMode = Literal["batch", "realtime", "image_generation"] @@ -819,6 +820,7 @@ class ModelInfoBody(BaseModel): mode: ModelMode | None = None access_groups: list[str] | None = None team_id: str | None = None + allowed_fails_policy: dict[str, int] | None = None class ModelNewBody(BaseModel): diff --git a/tests/e2e/router/reliability_support.py b/tests/e2e/router/reliability_support.py index e342aa363ca..5822058003c 100644 --- a/tests/e2e/router/reliability_support.py +++ b/tests/e2e/router/reliability_support.py @@ -19,6 +19,8 @@ from models import ( ChatMessage, ChatResponse, LiteLLMParamsBody, + ModelInfoBody, + ModelNewBody, ReliabilityChatBody, RouterSettingsOverride, ) @@ -26,6 +28,18 @@ from models import ( REAL_MODEL = "openai/gpt-5.5" REAL_KEY = "os.environ/OPENAI_API_KEY" +# The smallest-context chat model OpenAI still serves (16385 tokens). A prompt +# past that limit comes back as a real `context_length_exceeded` 400, which is +# what litellm maps to ContextWindowExceededError. +SMALL_CONTEXT_MODEL = "openai/gpt-3.5-turbo" +SMALL_CONTEXT_LIMIT_TOKENS = 16385 + + +def oversized_prompt(marker: str) -> str: + """A prompt comfortably past SMALL_CONTEXT_MODEL's context limit, so the + provider refuses it on length rather than answering a truncated version.""" + return f"{marker} " + ("token " * (SMALL_CONTEXT_LIMIT_TOKENS + 4000)) + def create_bad_base_deployment(proxy: ProxyClient, name: str) -> str: """Register a deployment pointing at an unreachable base, so every call to it @@ -40,6 +54,38 @@ def create_timeout_deployment(proxy: ProxyClient, name: str) -> str: return proxy.create_model(name, LiteLLMParamsBody(model=REAL_MODEL, api_key=REAL_KEY, timeout=0.001)) +def create_small_context_deployment(proxy: ProxyClient, name: str) -> str: + """Register a deployment on the smallest-context model OpenAI still serves, so an + oversized prompt earns a real context-window refusal from the provider.""" + return proxy.create_model(name, LiteLLMParamsBody(model=SMALL_CONTEXT_MODEL, api_key=REAL_KEY)) + + +def create_always_timing_out_deployment(proxy: ProxyClient, name: str) -> str: + """The always-picked half of a retry pair: a 1ms deadline the backend always + exceeds, all of the model group's shuffle weight, and a cooldown policy that + benches it on its first Timeout so the retry cannot land on it again.""" + return proxy.register_model( + ModelNewBody( + model_name=name, + litellm_params=LiteLLMParamsBody(model=REAL_MODEL, api_key=REAL_KEY, timeout=0.001, weight=1), + model_info=ModelInfoBody(allowed_fails_policy={"TimeoutErrorAllowedFails": 0}), + ) + ) + + +def create_zero_weight_backup_deployment(proxy: ProxyClient, name: str) -> str: + """The other half of a retry pair: healthy, but weight 0, so the weighted shuffle + never opens on it. It is reachable only once its sibling is benched and the + weighted pick falls through to a uniform one over what is left.""" + return proxy.register_model( + ModelNewBody( + model_name=name, + litellm_params=LiteLLMParamsBody(model=REAL_MODEL, api_key=REAL_KEY, weight=0), + model_info=ModelInfoBody(), + ) + ) + + def chat_override( proxy: ProxyClient, key: str, diff --git a/tests/e2e/router/test_reliability_fallbacks_e2e.py b/tests/e2e/router/test_reliability_fallbacks_e2e.py index d3ce62f8f95..8cece41ce2d 100644 --- a/tests/e2e/router/test_reliability_fallbacks_e2e.py +++ b/tests/e2e/router/test_reliability_fallbacks_e2e.py @@ -9,6 +9,10 @@ in the x-litellm-attempted-fallbacks header. Empty content is accepted only when `finish_reason == "length"` and the response billed completion tokens, since gpt-5.5 counts reasoning against max_tokens and can consume the whole budget before emitting any text; a fallback that produced nothing at all still fails. + +The context-window case is a different reroute from a plain failure: the provider +refuses the prompt on length, and `context_window_fallbacks` is the setting that +reroutes it, not `fallbacks`. """ from __future__ import annotations @@ -25,8 +29,10 @@ from reliability_support import ( completion_tokens_of, content_of, create_bad_base_deployment, + create_small_context_deployment, create_timeout_deployment, finish_reason_of, + oversized_prompt, reasoning_tokens_of, ) @@ -82,3 +88,17 @@ class TestReliabilityFallbacks: override=RouterSettingsOverride(fallbacks=[{primary: ["gpt-5.5"]}]), ) _assert_served_by_fallback(resp) + + @pytest.mark.covers("reliability.fallback.context_window.routes_to_fallback") + def test_context_window_routes_to_fallback( + self, client: ComplexityRouterClient, resources: ResourceManager, scoped_key: str + ) -> None: + primary = f"reliability-ctxfail-{unique_marker()}" + model_id = create_small_context_deployment(client.proxy, primary) + resources.defer(lambda: client.proxy.delete_model(model_id)) + + resp = chat_override( + client.proxy, scoped_key, primary, oversized_prompt(unique_marker()), + override=RouterSettingsOverride(context_window_fallbacks=[{primary: ["gpt-5.5"]}]), + ) + _assert_served_by_fallback(resp) diff --git a/tests/e2e/router/test_reliability_retries_e2e.py b/tests/e2e/router/test_reliability_retries_e2e.py new file mode 100644 index 00000000000..5441412935c --- /dev/null +++ b/tests/e2e/router/test_reliability_retries_e2e.py @@ -0,0 +1,73 @@ +"""Live e2e: a request that fails on its first deployment is retried inside its own +model group and still comes back a completion. + +The model group is a pair: an always-timing-out deployment that holds all of the +group's shuffle weight, and a healthy backup at weight 0. The weighted pick always +opens on the timing-out one, its first Timeout benches it (an +`allowed_fails_policy` of `TimeoutErrorAllowedFails: 0`), and the retry falls +through to the only deployment left. So the customer sees a completion and the +proxy reports that it took a retry to get there, with no random first pick in the +middle of it. +""" + +from __future__ import annotations + +import pytest + +from complexity_router_client import ComplexityRouterClient +from e2e_config import unique_marker +from lifecycle import ResourceManager +from models import RouterSettingsOverride +from reliability_support import ( + chat_override, + completion_tokens_of, + content_of, + create_always_timing_out_deployment, + create_zero_weight_backup_deployment, + finish_reason_of, +) + +pytestmark = pytest.mark.e2e + + +class TestReliabilityRetries: + @pytest.mark.covers("reliability.retry.timeout.succeeds_within_retries") + def test_timeout_on_first_deployment_succeeds_on_retry( + self, client: ComplexityRouterClient, resources: ResourceManager, scoped_key: str + ) -> None: + group = f"reliability-retry-{unique_marker()}" + timing_out = create_always_timing_out_deployment(client.proxy, group) + resources.defer(lambda: client.proxy.delete_model(timing_out)) + backup = create_zero_weight_backup_deployment(client.proxy, group) + resources.defer(lambda: client.proxy.delete_model(backup)) + + resp = chat_override( + client.proxy, + scoped_key, + group, + f"say hi {unique_marker()}", + override=RouterSettingsOverride(num_retries=2), + ) + + assert resp.status_code == 200, ( + f"the retry should have landed on the healthy backup, got {resp.status_code}: {resp.body[:300]}" + ) + + attempted = resp.headers.get("x-litellm-attempted-retries") + assert attempted is not None, "response is missing the x-litellm-attempted-retries header" + assert int(attempted) >= 1, ( + f"x-litellm-attempted-retries is {attempted!r}; a 200 with no retry means the request never " + "opened on the timing-out deployment, so this proves nothing about retries" + ) + + content = content_of(resp) + finish_reason = finish_reason_of(resp) + completion_tokens = completion_tokens_of(resp) or 0 + assert isinstance(content, str), ( + f"the retry should have returned a completion body, got content {content!r} (body={resp.body[:300]})" + ) + assert content or (finish_reason == "length" and completion_tokens > 0), ( + f"the retry returned empty content with finish_reason={finish_reason!r}, " + f"completion_tokens={completion_tokens}; empty content is only acceptable when the budget " + f"was spent on non-visible reasoning (body={resp.body[:300]})" + ) From 55d638412bb1f1b9c8eb9bb255ffbb8b4aa0c1c4 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Tue, 1 Sep 2026 15:17:00 -0700 Subject: [PATCH 26/97] fix: stop a cleared Team field from blocking personal key creation Clearing the Team combobox in the Create Key modal left team_id set to an empty string, so /key/generate treated the request as team key generation and failed with a team-not-found error for non-admin members. TeamDropdown now emits null on clear, and GenerateKeyRequest normalizes an empty team_id to None so the request runs the personal key path. --- litellm/proxy/_types.py | 7 +++ .../test_key_management_endpoints.py | 46 ++++++++++++++++++ .../common_components/team_dropdown.test.tsx | 48 +++++++++++++++++++ .../common_components/team_dropdown.tsx | 4 +- 4 files changed, 103 insertions(+), 2 deletions(-) create mode 100644 ui/litellm-dashboard/src/components/common_components/team_dropdown.test.tsx diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index e0a2097919b..18714256a8f 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -1216,6 +1216,13 @@ class GenerateKeyRequest(KeyRequestBase): organization_id: str | None = None project_id: str | None = None + @field_validator("team_id", mode="before") + @classmethod + def treat_cleared_team_id_as_unset(cls, v: object) -> object: + if v == "": + return None + return v + class GenerateKeyResponse(KeyRequestBase): key: str diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index a81c6b4c656..7e2e680743f 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -17489,3 +17489,49 @@ async def test_check_project_key_limits_still_rejects_real_model_outside_project assert exc_info.value.status_code == 400 assert "Model 'gpt-5.4-mini' not in project's allowed models" in exc_info.value.detail["error"] + + +def test_generate_key_request_blank_team_id_is_personal(): + """The UI Team-field clear submits team_id=""; it must count as no team (LIT-3925).""" + from litellm.proxy._types import RegenerateKeyRequest + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _is_team_key, + ) + + cleared = GenerateKeyRequest(team_id="") + assert cleared.team_id is None + assert _is_team_key(data=cleared) is False + assert RegenerateKeyRequest(team_id="").team_id is None + assert GenerateKeyRequest(team_id="team-1").team_id == "team-1" + + +def test_key_generation_check_blank_team_id_uses_personal_permissions(monkeypatch): + """key_generation_check with team_id="" must take the personal-key path instead + of failing the team lookup with "Unable to find team object" (LIT-3925).""" + from litellm.proxy._types import KeyManagementRoutes + from litellm.proxy.management_endpoints.key_management_endpoints import ( + key_generation_check, + ) + + monkeypatch.setattr( + litellm, + "key_generation_settings", + { + "team_key_generation": {"allowed_team_member_roles": ["admin"]}, + "personal_key_generation": {"allowed_user_roles": ["proxy_admin", "internal_user"]}, + }, + ) + + assert ( + key_generation_check( + team_table=None, + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + api_key="sk-alice", + user_id="alice", + ), + data=GenerateKeyRequest(key_alias="personal", team_id=""), + route=KeyManagementRoutes.KEY_GENERATE, + ) + is True + ) diff --git a/ui/litellm-dashboard/src/components/common_components/team_dropdown.test.tsx b/ui/litellm-dashboard/src/components/common_components/team_dropdown.test.tsx new file mode 100644 index 00000000000..90f7be1477c --- /dev/null +++ b/ui/litellm-dashboard/src/components/common_components/team_dropdown.test.tsx @@ -0,0 +1,48 @@ +import { render, screen } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; +import { describe, expect, it, vi } from "vitest"; + +import { chooseSelectOption } from "../../../tests/test-utils"; +import type { Team } from "../key_team_helpers/key_list"; +import TeamDropdown from "./team_dropdown"; + +const TEAMS = [ + { team_id: "team-1", team_alias: "Alpha Team" }, + { team_id: "team-2", team_alias: "Beta Team" }, +] as unknown as Team[]; + +vi.mock("@/app/(dashboard)/hooks/teams/useTeams", () => ({ + useInfiniteTeams: () => ({ + data: { pages: [{ teams: TEAMS }] }, + fetchNextPage: vi.fn(), + hasNextPage: false, + isFetchingNextPage: false, + isLoading: false, + }), +})); + +describe("TeamDropdown", () => { + it("emits the picked team's id and full object", async () => { + const user = userEvent.setup(); + const onChange = vi.fn(); + const onTeamSelect = vi.fn(); + render(); + + await chooseSelectOption(user, screen.getByRole("combobox"), /^Beta Team/); + + expect(onChange).toHaveBeenCalledWith("team-2"); + expect(onTeamSelect).toHaveBeenCalledWith(TEAMS[1]); + }); + + it("emits null, never the empty string, when the selection is cleared", async () => { + const user = userEvent.setup(); + const onChange = vi.fn(); + const onTeamSelect = vi.fn(); + render(); + + await user.click(document.querySelector('[data-slot="combobox-clear"]') as HTMLElement); + + expect(onChange).toHaveBeenCalledWith(null); + expect(onTeamSelect).toHaveBeenCalledWith(null); + }); +}); diff --git a/ui/litellm-dashboard/src/components/common_components/team_dropdown.tsx b/ui/litellm-dashboard/src/components/common_components/team_dropdown.tsx index 35121f41598..7d385c2a3f7 100644 --- a/ui/litellm-dashboard/src/components/common_components/team_dropdown.tsx +++ b/ui/litellm-dashboard/src/components/common_components/team_dropdown.tsx @@ -5,7 +5,7 @@ import { Team } from "../key_team_helpers/key_list"; interface TeamDropdownProps { value?: string; - onChange?: (value: string) => void; + onChange?: (value: string | null) => void; /** Callback with the full Team object (or null on clear). */ onTeamSelect?: (team: Team | null) => void; disabled?: boolean; @@ -47,7 +47,7 @@ const TeamDropdown: React.FC = ({ }, [data]); const handleChange = (teamId: string) => { - onChange?.(teamId); + onChange?.(teamId || null); if (onTeamSelect) { onTeamSelect(teamId ? teams.find((t) => t.team_id === teamId) ?? null : null); } From 4acc1d15fb7f9172d39417967445cf50176c5012 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Tue, 1 Sep 2026 15:50:19 -0700 Subject: [PATCH 27/97] fix(ui): map a cleared Team dropdown back to an empty string in the auto-router form --- .../add_model/add_auto_router_tab.test.tsx | 47 ++++++++++++++----- .../add_model/add_auto_router_tab.tsx | 4 +- 2 files changed, 39 insertions(+), 12 deletions(-) diff --git a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx index d8a955b719c..ba380662403 100644 --- a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx +++ b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx @@ -101,18 +101,24 @@ vi.mock("./build_complexity_router_config", async (importOriginal) => { }); // A real TeamDropdown fetches teams and renders an antd Select; the wiring under test is -// whether team_id is registered, validated and forwarded, so a plain control stands in. +// whether team_id is registered, validated and forwarded, so a plain control stands in. The +// clear button mirrors the real dropdown's x, which emits null rather than a string. vi.mock("../common_components/team_dropdown", () => ({ - default: ({ value, onChange }: { value?: string; onChange?: (next: string) => void }) => ( - + default: ({ value, onChange }: { value?: string; onChange?: (next: string | null) => void }) => ( + <> + + + ), })); @@ -354,6 +360,25 @@ describe("AddAutoRouterTab", () => { expect(handleAddAutoRouterSubmit).not.toHaveBeenCalled(); }); + // The shared dropdown emits null on clear while this form's schema wants a string, so the + // form maps null back to "": the user sees the pick-a-team message, not a zod type error. + it("treats a team picked and then cleared like no team at all", async () => { + const user = userEvent.setup(); + vi.mocked(getMissingTiersError).mockReturnValue(null); + + renderWithProviders( + , + ); + + await user.type(screen.getByPlaceholderText(/smart_router/i), "team-scoped-router"); + await user.selectOptions(screen.getByTestId("team-dropdown"), "team-1"); + await user.click(screen.getByTestId("team-dropdown-clear")); + await user.click(screen.getByRole("button", { name: /add auto router/i })); + + expect(await screen.findByText("Please select a team to continue")).toBeInTheDocument(); + expect(handleAddAutoRouterSubmit).not.toHaveBeenCalled(); + }); + it("defaults a new router to session affinity off, matching the backend field default", async () => { const user = userEvent.setup(); vi.mocked(getMissingTiersError).mockReturnValue(null); diff --git a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx index 318adcce369..03b7432e4ec 100644 --- a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx +++ b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx @@ -548,7 +548,9 @@ const AddAutoRouterTab: React.FC = ({ "Select the team this auto router belongs to. Only keys for this team will be able to call it.", )} > - {({ id, value, onChange }) => } + {({ id, value, onChange }) => ( + onChange(next ?? "")} /> + )} )} From 00e40c0afe558cf18ef1f446f94e52f8594b535a Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Tue, 1 Sep 2026 15:28:12 -0700 Subject: [PATCH 28/97] Record each e2e test's source location in the JUnit report The JUnit report is the only thing that leaves the e2e run, and it says where a test's results came from but never where its code lives. A reader looking at `test_cell_claimed_only_by_a_skipped_test_is_uncovered` on the status page has a name and nothing else -- no file, no line, no way to reach the source short of grepping the repo by hand. Pytest knows the location; the report format loses it. The `xunit1` family wrote `file=` and `line=` onto every ``, and the `xunit2` default this suite runs on drops both. Switching families back would change the document for every consumer of the same XML -- the Buildkite Test Engine upload and the Loki pipeline included -- so add the location the way this suite already adds `package` and `covers`: as a ``, which is purely additive. `source` is repo-relative and one-based (`tests/e2e/a2a/test_x.py:41`), so a consumer can build a link without knowing how pytest was started. That takes normalizing the two launch shapes -- the runner image runs from its own copy at /app/e2e, a developer runs from the repo root -- which is the same normalization `package_from_nodeid` was already doing in reverse, now factored into `suite_parts` so the two cannot drift apart. Paths that escape the suite, and tests pytest reports no line for, emit an empty string: a test with no link beats a link that 404s. Claude-Session: https://claude.ai/code/session_017dTKXwJkzhtVLzDhePHsKG --- tests/e2e/junit_properties.py | 84 ++++++++++++++--- tests/e2e/test_junit_properties.py | 145 +++++++++++++++++++++++++++++ 2 files changed, 215 insertions(+), 14 deletions(-) create mode 100644 tests/e2e/test_junit_properties.py diff --git a/tests/e2e/junit_properties.py b/tests/e2e/junit_properties.py index e4f59f5c4d2..5b5e239bf9a 100644 --- a/tests/e2e/junit_properties.py +++ b/tests/e2e/junit_properties.py @@ -2,10 +2,20 @@ The e2e suite ships results to Loki/Grafana from a standard pytest JUnit report (`--junitxml=e2e-report.xml`), not a bespoke log line. JUnit already records -outcome, duration, and node id for every ``; the only signals it cannot -derive on its own are the normalized suite package and the coverage-registry cell -ids a test covers. Those ride along as JUnit `` entries via each item's -`user_properties`, attached in `conftest.py::pytest_collection_modifyitems`. +outcome, duration, and node id for every ``; the signals it cannot +derive on its own are the normalized suite package, the coverage-registry cell +ids a test covers, and where the test's source lives. Those ride along as JUnit +`` entries via each item's `user_properties`, attached in +`conftest.py::pytest_collection_modifyitems`. + +`source` is here because of a reporter limitation rather than a missing pytest +fact. Pytest knows every test's file and line, and its `xunit1` report family +wrote them as `file=` / `line=` attributes on ``. The default `xunit2` +family -- pytest's since 6.0, and this suite's, since pytest.ini names no family +-- drops both. Switching families to get them back would change the document +shape for every consumer of the same XML, the Buildkite Test Engine upload and +the Loki pipeline included; a property is additive, so nothing that reads the +report today sees a difference. """ from __future__ import annotations @@ -14,22 +24,66 @@ from collections.abc import Iterable import pytest +# This module's own directory, relative to the repo root. Hardcoded because it +# cannot be discovered at runtime: the e2e runner image copies tests/e2e/ to +# /app/e2e and runs pytest from there, so no ancestor of this file names the +# suite's place in the litellm tree. Moving tests/e2e/ means editing this line, +# and test_junit_properties.py fails from a checkout until you do. +SUITE_ROOT = "tests/e2e" + + +def suite_parts(path_part: str) -> tuple[str, ...]: + """Path components of a suite file, relative to tests/e2e, either way it ran. + + Pytest reports paths relative to its rootdir, which moves with the + invocation: a repo-root run gives `tests/e2e/logging/test_x.py`, a suite-cwd + run (what the runner image does) gives `logging/test_x.py`. Strip the + `tests/e2e` prefix when present so both collapse to the same components. + """ + raw = tuple(p for p in path_part.replace("\\", "/").split("/") if p and p != ".") + return raw[2:] if len(raw) >= 3 and raw[0] == "tests" and raw[1] == "e2e" else raw + def package_from_nodeid(nodeid: str) -> str: - """Top-level suite package under tests/e2e/, or 'root' for top-level files. - - Pytest nodeids are relative to the invocation cwd. Repo-root runs look like - `tests/e2e/logging/...`; suite-cwd runs look like `logging/...`. Strip the - `tests/e2e` prefix so package is the suite dir either way. - """ - path_part = nodeid.split("::", 1)[0].replace("\\", "/") - raw = tuple(p for p in path_part.split("/") if p and p != ".") - parts = raw[2:] if len(raw) >= 3 and raw[0] == "tests" and raw[1] == "e2e" else raw + """Top-level suite package under tests/e2e/, or 'root' for top-level files.""" + parts = suite_parts(nodeid.split("::", 1)[0]) if len(parts) <= 1: return "root" return parts[0] +def source_from_location(path: str, lineno: int | None) -> str: + """Repo-relative `path:line` for a test, or '' when nothing is linkable. + + `pytest.Item.location` supplies a rootdir-relative path and a ZERO-based + line number, and neither travels as-is. The path is re-rooted at SUITE_ROOT + so consumers never have to know how pytest was started, and the line is + emitted ONE-based, matching editors, tracebacks, and code hosts (GitHub's + `#L41` is the file's 41st line). A decorated test anchors at its first + decorator, which is where pytest reports it and which puts the marks and the + `def` on screen together. + + Returns '' rather than a guess when pytest reports no line, or when the path + escapes the suite root (absolute, or reaching upward): a test that renders + without a link is a smaller failure than one that links somewhere wrong. + """ + if lineno is None: + return "" + normalized = path.replace("\\", "/") + if normalized.startswith("/") or ".." in normalized.split("/"): + return "" + parts = suite_parts(normalized) + if not parts: + return "" + return f"{'/'.join((SUITE_ROOT, *parts))}:{lineno + 1}" + + +def source_from_item(item: pytest.Item) -> str: + """Read the repo-relative `path:line` off a pytest Item's reported location.""" + path, lineno, _ = item.location + return source_from_location(path, lineno) + + def dedupe_covers(marker_args: Iterable[tuple[object, ...]]) -> tuple[str, ...]: """Flatten @pytest.mark.covers arg lists into unique, order-preserving cell ids, dropping anything that is not a non-empty string.""" @@ -43,10 +97,12 @@ def covers_from_item(item: pytest.Item) -> tuple[str, ...]: def result_properties(item: pytest.Item) -> tuple[tuple[str, str], ...]: """The custom signals a standard reporter cannot derive: the normalized suite - package and the comma-joined coverage-registry cell ids this test covers.""" + package, the comma-joined coverage-registry cell ids this test covers, and the + repo-relative `path:line` its source sits at.""" return ( ("package", package_from_nodeid(item.nodeid)), ("covers", ",".join(covers_from_item(item))), + ("source", source_from_item(item)), ) diff --git a/tests/e2e/test_junit_properties.py b/tests/e2e/test_junit_properties.py new file mode 100644 index 00000000000..8aa267673cb --- /dev/null +++ b/tests/e2e/test_junit_properties.py @@ -0,0 +1,145 @@ +"""Harness coverage for the custom JUnit properties. + +No proxy and no ``e2e`` marker. Pins the two normalizations that have to agree +about where a suite file lives -- ``package_from_nodeid`` (strip the suite root) +and ``source_from_location`` (re-root at it) -- across both ways the suite is +launched, plus the one-based line offset and the refusal to emit a path that +escapes the suite. The consumers of these properties are the Loki/Grafana +rollups and, for ``source``, the status page's per-test links to GitHub. +""" + +from __future__ import annotations + +from pathlib import Path + +import pytest +from junit_properties import ( + SUITE_ROOT, + attach_result_properties, + dedupe_covers, + package_from_nodeid, + result_properties, + source_from_location, + suite_parts, +) + + +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 repo_root() -> Path | None: + """The litellm checkout above this file, or None when there isn't one.""" + return next((p for p in Path(__file__).resolve().parents if (p / ".git").exists()), None) + + +class TestSuiteParts: + @pytest.mark.parametrize( + "path", + ["logging/test_x.py", "tests/e2e/logging/test_x.py", "./logging/test_x.py", "tests\\e2e\\logging\\test_x.py"], + ) + def test_both_invocation_shapes_collapse_to_the_same_components(self, path: str) -> None: + """A repo-root run and a suite-cwd run report the same file differently; + every downstream signal has to see one spelling.""" + assert suite_parts(path) == ("logging", "test_x.py") + + def test_top_level_suite_file_keeps_its_single_component(self) -> None: + assert suite_parts("tests/e2e/test_fixture_mode.py") == ("test_fixture_mode.py",) + + +class TestPackageFromNodeid: + @pytest.mark.parametrize( + ("nodeid", "expected"), + [ + ("logging/test_x.py::TestFoo::test_bar", "logging"), + ("tests/e2e/logging/test_x.py::TestFoo::test_bar", "logging"), + ("quota_management/spend_tracking/test_x.py::test_bar", "quota_management"), + ("test_fixture_mode.py::TestParseFixtureMode::test_known_values_normalize", "root"), + ("tests/e2e/test_fixture_mode.py::test_bar", "root"), + ], + ) + def test_package_is_the_first_dir_under_the_suite_root(self, nodeid: str, expected: str) -> None: + assert package_from_nodeid(nodeid) == expected + + +class TestSourceFromLocation: + @pytest.mark.parametrize("path", ["a2a/test_a2a_agent_e2e.py", "tests/e2e/a2a/test_a2a_agent_e2e.py"]) + def test_path_is_repo_relative_however_pytest_was_started(self, path: str) -> None: + assert source_from_location(path, 40) == "tests/e2e/a2a/test_a2a_agent_e2e.py:41" + + def test_line_is_emitted_one_based(self) -> None: + """pytest.Item.location counts from 0; editors, tracebacks and GitHub's + #L anchor all count from 1, and an off-by-one lands on the decorator.""" + assert source_from_location("a2a/test_x.py", 0) == "tests/e2e/a2a/test_x.py:1" + + def test_top_level_suite_file_sits_directly_under_the_suite_root(self) -> None: + assert source_from_location("test_fixture_mode.py", 39) == "tests/e2e/test_fixture_mode.py:40" + + @pytest.mark.parametrize( + ("path", "lineno"), + [ + ("a2a/test_x.py", None), + ("/app/e2e/a2a/test_x.py", 40), + ("../conftest.py", 40), + ("", 40), + ], + ) + def test_nothing_linkable_yields_empty_rather_than_a_guess(self, path: str, lineno: int | None) -> None: + """A test that renders without a link is a smaller failure than one whose + link 404s or points into another repo's file.""" + assert source_from_location(path, lineno) == "" + + +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"), + ("covers", "LOG-1,LOG-2"), + ("source", "tests/e2e/logging/test_x.py:41"), + ) + + def test_attach_is_idempotent(self) -> 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")) + attach_result_properties(item) + attach_result_properties(item) + assert [name for name, _ in item.user_properties] == ["package", "covers", "source"] + + +class TestSuiteRoot: + def test_suite_root_names_this_file_s_real_home(self) -> None: + """SUITE_ROOT is hardcoded because the runner image has no repo to read it + from. Where there IS a checkout, prove the constant still points at us -- + otherwise a moved tests/e2e/ ships links that 404.""" + root = repo_root() + if root is None: + pytest.skip("no checkout above this file (the runner image copies tests/e2e/ to /app/e2e)") + assert (root / SUITE_ROOT / Path(__file__).name).resolve() == Path(__file__).resolve() + + +class TestDedupeCovers: + def test_ids_are_unique_order_preserving_and_non_empty_strings(self) -> None: + assert dedupe_covers([("A", "B"), ("B", ""), ("C", 7)]) == ("A", "B", "C") From 0c2d4c5773a7269f026831d359ca86bea20b5159 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Tue, 1 Sep 2026 16:02:54 -0700 Subject: [PATCH 29/97] Refuse a source path carrying a colon `path:line` cannot represent a path that itself contains a colon, and the one way pytest produces one is a Windows absolute location: separator normalization turns `C:\app\e2e\a2a\test_x.py` into `C:/app/...`, which slipped past the leading-slash check and composed the nonsense repo path `tests/e2e/C:/app/e2e/a2a/test_x.py`. Reject the colon itself rather than special-casing a drive letter: it is the character the format reserves, so no path containing one was ever linkable. Claude-Session: https://claude.ai/code/session_017dTKXwJkzhtVLzDhePHsKG --- tests/e2e/junit_properties.py | 49 ++++++++++++------------------ tests/e2e/test_junit_properties.py | 5 +-- 2 files changed, 23 insertions(+), 31 deletions(-) diff --git a/tests/e2e/junit_properties.py b/tests/e2e/junit_properties.py index 5b5e239bf9a..c5971c5362c 100644 --- a/tests/e2e/junit_properties.py +++ b/tests/e2e/junit_properties.py @@ -8,14 +8,10 @@ ids a test covers, and where the test's source lives. Those ride along as JUnit `` entries via each item's `user_properties`, attached in `conftest.py::pytest_collection_modifyitems`. -`source` is here because of a reporter limitation rather than a missing pytest -fact. Pytest knows every test's file and line, and its `xunit1` report family -wrote them as `file=` / `line=` attributes on ``. The default `xunit2` -family -- pytest's since 6.0, and this suite's, since pytest.ini names no family --- drops both. Switching families to get them back would change the document -shape for every consumer of the same XML, the Buildkite Test Engine upload and -the Loki pipeline included; a property is additive, so nothing that reads the -report today sees a difference. +`source` is a property rather than the `file=` / `line=` attributes pytest used +to write, because the `xunit2` family this suite runs on drops those, and +switching families would change the XML for every consumer of it -- the +Buildkite Test Engine upload and the Loki pipeline included. """ from __future__ import annotations @@ -24,21 +20,18 @@ from collections.abc import Iterable import pytest -# This module's own directory, relative to the repo root. Hardcoded because it -# cannot be discovered at runtime: the e2e runner image copies tests/e2e/ to -# /app/e2e and runs pytest from there, so no ancestor of this file names the -# suite's place in the litellm tree. Moving tests/e2e/ means editing this line, -# and test_junit_properties.py fails from a checkout until you do. +# Hardcoded because the runner image copies tests/e2e/ to /app/e2e, so nothing +# at runtime names this suite's place in the repo. test_junit_properties.py +# fails from a checkout if it moves. SUITE_ROOT = "tests/e2e" def suite_parts(path_part: str) -> tuple[str, ...]: - """Path components of a suite file, relative to tests/e2e, either way it ran. + """Path components of a suite file relative to tests/e2e, however it ran. - Pytest reports paths relative to its rootdir, which moves with the - invocation: a repo-root run gives `tests/e2e/logging/test_x.py`, a suite-cwd - run (what the runner image does) gives `logging/test_x.py`. Strip the - `tests/e2e` prefix when present so both collapse to the same components. + Pytest paths are rootdir-relative, and rootdir moves with the invocation: a + repo-root run gives `tests/e2e/logging/test_x.py`, a suite-cwd run (the + runner image) gives `logging/test_x.py`. Both collapse to the same tuple. """ raw = tuple(p for p in path_part.replace("\\", "/").split("/") if p and p != ".") return raw[2:] if len(raw) >= 3 and raw[0] == "tests" and raw[1] == "e2e" else raw @@ -55,22 +48,20 @@ def package_from_nodeid(nodeid: str) -> str: def source_from_location(path: str, lineno: int | None) -> str: """Repo-relative `path:line` for a test, or '' when nothing is linkable. - `pytest.Item.location` supplies a rootdir-relative path and a ZERO-based - line number, and neither travels as-is. The path is re-rooted at SUITE_ROOT - so consumers never have to know how pytest was started, and the line is - emitted ONE-based, matching editors, tracebacks, and code hosts (GitHub's - `#L41` is the file's 41st line). A decorated test anchors at its first - decorator, which is where pytest reports it and which puts the marks and the - `def` on screen together. + `pytest.Item.location` gives a rootdir-relative path and a ZERO-based line. + The path is re-rooted at SUITE_ROOT so consumers need not know how pytest was + started, and the line is emitted ONE-based to match editors, tracebacks and + code hosts. A decorated test anchors at its first decorator, which is where + pytest reports it. - Returns '' rather than a guess when pytest reports no line, or when the path - escapes the suite root (absolute, or reaching upward): a test that renders - without a link is a smaller failure than one that links somewhere wrong. + Empty rather than a guess for anything unlinkable: no line, a path reaching + upward, or a path carrying a colon, which is both how an absolute Windows + path arrives and a character `path:line` has no way to represent. """ if lineno is None: return "" normalized = path.replace("\\", "/") - if normalized.startswith("/") or ".." in normalized.split("/"): + if normalized.startswith("/") or ":" in normalized or ".." in normalized.split("/"): return "" parts = suite_parts(normalized) if not parts: diff --git a/tests/e2e/test_junit_properties.py b/tests/e2e/test_junit_properties.py index 8aa267673cb..c0596177cc1 100644 --- a/tests/e2e/test_junit_properties.py +++ b/tests/e2e/test_junit_properties.py @@ -97,13 +97,14 @@ class TestSourceFromLocation: [ ("a2a/test_x.py", None), ("/app/e2e/a2a/test_x.py", 40), + ("C:\\app\\e2e\\a2a\\test_x.py", 40), ("../conftest.py", 40), ("", 40), ], ) def test_nothing_linkable_yields_empty_rather_than_a_guess(self, path: str, lineno: int | None) -> None: - """A test that renders without a link is a smaller failure than one whose - link 404s or points into another repo's file.""" + """A colon is rejected on two counts: it is how a Windows absolute path + arrives, and `path:line` cannot represent one in the path half.""" assert source_from_location(path, lineno) == "" From 1964d92fc65733c13e78f4cbb9ae62357325c454 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Tue, 1 Sep 2026 16:04:45 -0700 Subject: [PATCH 30/97] test(ui): query the clear button and the models page tabs through accessible screen queries --- .../app/(dashboard)/models-and-endpoints/page.test.tsx | 10 +++++----- .../common_components/team_dropdown.test.tsx | 2 +- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.test.tsx index 1199b66621f..84a05113177 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.test.tsx @@ -111,17 +111,17 @@ describe("ModelsAndEndpointsPage", () => { // POST /model/new 403s a proxy_admin_viewer, so the form's tab must not render for one. it("hides the Add Model tab for a view-only admin session", () => { mockUseAuthorized.mockReturnValue(VIEW_ONLY_ADMIN); - const { getByRole, queryByRole } = renderPage(); - expect(queryByRole("tab", { name: "Add Model" })).not.toBeInTheDocument(); - expect(getByRole("tab", { name: "All Models" })).toBeInTheDocument(); + renderPage(); + expect(screen.queryByRole("tab", { name: "Add Model" })).not.toBeInTheDocument(); + expect(screen.getByRole("tab", { name: "All Models" })).toBeInTheDocument(); }); // Read parity: the Auto-Routers list stays reachable for a view-only admin; only the // create affordance inside it is withheld, which AutoRoutersTabPanel decides. it("keeps the Auto-Routers tab for a view-only admin session", () => { mockUseAuthorized.mockReturnValue(VIEW_ONLY_ADMIN); - const { getByRole } = renderPage(); - expect(getByRole("tab", { name: /Auto-Routers/ })).toBeInTheDocument(); + renderPage(); + expect(screen.getByRole("tab", { name: /Auto-Routers/ })).toBeInTheDocument(); }); // Auto-routers are excluded from the All Models table, so this tab is their home: the only diff --git a/ui/litellm-dashboard/src/components/common_components/team_dropdown.test.tsx b/ui/litellm-dashboard/src/components/common_components/team_dropdown.test.tsx index 90f7be1477c..01b8a12d99c 100644 --- a/ui/litellm-dashboard/src/components/common_components/team_dropdown.test.tsx +++ b/ui/litellm-dashboard/src/components/common_components/team_dropdown.test.tsx @@ -40,7 +40,7 @@ describe("TeamDropdown", () => { const onTeamSelect = vi.fn(); render(); - await user.click(document.querySelector('[data-slot="combobox-clear"]') as HTMLElement); + await user.click(screen.getByRole("button", { name: "Clear" })); expect(onChange).toHaveBeenCalledWith(null); expect(onTeamSelect).toHaveBeenCalledWith(null); From 3888a85045f05007baa41cb6d2b1a8677ef81fbd Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 1 Sep 2026 16:13:26 -0700 Subject: [PATCH 31/97] fix(budget): reject known estimates over remaining budget under fail_closed_budget_enforcement (#39214) Co-authored-by: yassin Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../spend_tracking/budget_reservation.py | 36 ++++++-- .../proxy/test_budget_reservation.py | 88 +++++++++++++++++++ 2 files changed, 115 insertions(+), 9 deletions(-) diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index f4d8fd7d906..91d2ece7a51 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -120,13 +120,15 @@ async def _apply_over_budget_reservation_policy( applied_entries: list[dict[str, float | str]], reservation_cost: float, current_spend: float, + fail_closed_budget_enforcement: bool = False, ) -> float: """ Decide what to do when a counter is over budget, and return the reservation cost to carry into the next counter. Three outcomes: an over-budget key that opted into throttling releases its own reservation (the rate limiter slows it) and keeps the cost; a partially-remaining budget resizes the reservation - down to what is left; anything else hard-blocks by raising. + down to what is left, unless strict enforcement is on, because the known + estimate already does not fit; anything else hard-blocks by raising. """ if _key_reservation_should_release_for_throttle(counter.counter_key, valid_token): await _release_applied_entries_best_effort(entries=[entry], default_reserved_cost=reservation_cost) @@ -134,21 +136,36 @@ async def _apply_over_budget_reservation_policy( return reservation_cost remaining_before_reservation: Final = counter.max_budget - (current_spend - reservation_cost) - if remaining_before_reservation > 1e-12: - await _resize_applied_reservation( - entries=applied_entries, - current_reserved_cost=reservation_cost, - new_reserved_cost=remaining_before_reservation, + if remaining_before_reservation <= 1e-12: + _raise_counter_budget_exceeded(counter=counter, current_cost=current_spend) + if fail_closed_budget_enforcement and current_spend - counter.max_budget > 1e-12: + _raise_counter_budget_exceeded( + counter=counter, + current_cost=current_spend - reservation_cost, + estimated_cost=reservation_cost, ) - return remaining_before_reservation + await _resize_applied_reservation( + entries=applied_entries, + current_reserved_cost=reservation_cost, + new_reserved_cost=remaining_before_reservation, + ) + return remaining_before_reservation + +def _raise_counter_budget_exceeded( + counter: _BudgetCounter, + current_cost: float, + estimated_cost: float | None = None, +) -> NoReturn: + estimate_detail: Final = "" if estimated_cost is None else f"Estimated request cost: {estimated_cost}, " raise litellm.BudgetExceededError( - current_cost=current_spend, + current_cost=current_cost, max_budget=counter.max_budget, message=( "Budget has been exceeded! " f"{counter.entity_type}={counter.entity_id} " - f"Current cost: {current_spend}, " + f"Current cost: {current_cost}, " + f"{estimate_detail}" f"Max budget: {counter.max_budget}" ), entity_type=_COUNTER_ENTITY_TYPES.get(counter.entity_type), @@ -258,6 +275,7 @@ async def reserve_budget_for_request( applied_entries=applied_entries, reservation_cost=reservation_cost, current_spend=current_spend, + fail_closed_budget_enforcement=fail_closed_budget_enforcement, ) continue except Exception: diff --git a/tests/test_litellm/proxy/test_budget_reservation.py b/tests/test_litellm/proxy/test_budget_reservation.py index 95067929ac1..b8fb6170d34 100644 --- a/tests/test_litellm/proxy/test_budget_reservation.py +++ b/tests/test_litellm/proxy/test_budget_reservation.py @@ -819,6 +819,94 @@ async def test_should_cap_known_estimate_to_remaining_budget( ) == pytest.approx(0.9) +@pytest.mark.asyncio +async def test_fail_closed_rejects_known_estimate_exceeding_remaining_budget( + spend_counter_state, +): + """LIT-5922: with strict enforcement on, a request whose known estimate does + not fit the remaining budget must be rejected before dispatch instead of + having its reservation shrunk to the headroom and admitted, and the counter + must be restored to the pre-request spend.""" + counter_cache, key_cache = spend_counter_state + proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache) + valid_token = UserAPIKeyAuth( + token="key-budget-known-estimate-fail-closed", + spend=0.9, + max_budget=1.0, + ) + counter_cache.in_memory_cache.set_cache( + key="spend:key:key-budget-known-estimate-fail-closed", + value=0.9, + ) + + with patch( # test-quality-ok: reserve_budget_for_request takes no estimator, so pinning the estimate needs this attribute + "litellm.proxy.spend_tracking.budget_reservation.estimate_request_max_cost", + return_value=0.6, + ): + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await reserve_budget_for_request( + request_body=_request_body(), + route="/chat/completions", + llm_router=None, + valid_token=valid_token, + team_object=None, + user_object=None, + prisma_client=None, + user_api_key_cache=key_cache, + proxy_logging_obj=proxy_logging_obj, + fail_closed_budget_enforcement=True, + ) + + assert exc_info.value.current_cost == pytest.approx(0.9) + assert exc_info.value.max_budget == pytest.approx(1.0) + assert "Current cost: 0.9, Estimated request cost: 0.6, Max budget: 1.0" in str(exc_info.value) + assert counter_cache.in_memory_cache.get_cache( + key="spend:key:key-budget-known-estimate-fail-closed" + ) == pytest.approx(0.9) + + +@pytest.mark.asyncio +async def test_fail_closed_tolerates_float_noise_when_estimate_exactly_fits( + spend_counter_state, +): + """0.1 + 0.2 lands a hair above 0.3 in floating point. Strict enforcement + must treat that as fitting the budget, not reject it.""" + counter_cache, key_cache = spend_counter_state + proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache) + valid_token = UserAPIKeyAuth( + token="key-budget-fail-closed-float-noise", + spend=0.1, + max_budget=0.3, + ) + counter_cache.in_memory_cache.set_cache( + key="spend:key:key-budget-fail-closed-float-noise", + value=0.1, + ) + + with patch( # test-quality-ok: reserve_budget_for_request takes no estimator, so pinning the estimate needs this attribute + "litellm.proxy.spend_tracking.budget_reservation.estimate_request_max_cost", + return_value=0.2, + ): + reservation = await reserve_budget_for_request( + request_body=_request_body(), + route="/chat/completions", + llm_router=None, + valid_token=valid_token, + team_object=None, + user_object=None, + prisma_client=None, + user_api_key_cache=key_cache, + proxy_logging_obj=proxy_logging_obj, + fail_closed_budget_enforcement=True, + ) + + assert reservation is not None + assert reservation["reserved_cost"] == pytest.approx(0.2) + assert counter_cache.in_memory_cache.get_cache( + key="spend:key:key-budget-fail-closed-float-noise" + ) == pytest.approx(0.3) + + @pytest.mark.asyncio async def test_should_clamp_reservation_to_default_when_output_cap_missing( spend_counter_state, From e11a2ec0f671f7b01696f543715910b270681157 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Tue, 1 Sep 2026 16:19:17 -0700 Subject: [PATCH 32/97] Re-run checks after retargeting to litellm_internal_staging The Guard main branch job ran while this PR still pointed at main and recorded a failure that cannot clear: re-running it replays the original event payload, base included. Its trigger is scoped to PRs against main, so it does not apply now and a fresh head SHA is what drops the stale run. Claude-Session: https://claude.ai/code/session_017dTKXwJkzhtVLzDhePHsKG From 4a68abfd49eea3b7d4fdfeaadd55598369848ef2 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 1 Sep 2026 23:26:23 +0000 Subject: [PATCH 33/97] fix(proxy): route container create and list through model_list deployments Container create and list requests had no container ID to decode, so the router called the provider handler directly and the OpenAI transformation fell back to the global OPENAI_API_KEY. Proxies configured only with model_list credentials sent Authorization: Bearer None. Route through _ageneric_api_call_with_fallbacks when the caller passes a model, expose the list endpoint's model query param to the router, and encode the managed container ID on the async create path so follow-up calls route to the same deployment. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/containers/main.py | 25 +++++--- .../proxy/container_endpoints/endpoints.py | 2 +- litellm/router.py | 11 +++- .../test_router_endpoints.py | 59 +++++++++++++++++++ .../containers/test_container_api.py | 36 +++++++++++ 5 files changed, 122 insertions(+), 11 deletions(-) diff --git a/litellm/containers/main.py b/litellm/containers/main.py index 97ca11872c1..90d59af009f 100644 --- a/litellm/containers/main.py +++ b/litellm/containers/main.py @@ -1,7 +1,7 @@ import asyncio import contextvars import json -from collections.abc import Coroutine, Mapping +from collections.abc import Callable, Coroutine, Mapping from functools import partial from typing import Final, Literal, overload @@ -47,6 +47,13 @@ __all__ = [ ##### Container Create ####################### +async def _encode_created_container_id( + pending: Coroutine[object, object, ContainerObject], + encode: Callable[[ContainerObject], ContainerObject], +) -> ContainerObject: + return encode(await pending) + + @client async def acreate_container( name: str, @@ -256,16 +263,16 @@ def create_container( _is_async=_is_async, ) - # Encode container_id with provider/model metadata for routing + encode: Final = partial( + ContainerRequestUtils.encode_container_id_in_response, + custom_llm_provider=custom_llm_provider, + litellm_metadata=kwargs.get("litellm_metadata"), + extra_body=extra_body, + ) if isinstance(container_obj, ContainerObject): - container_obj = ContainerRequestUtils.encode_container_id_in_response( - response_obj=container_obj, - custom_llm_provider=custom_llm_provider, - litellm_metadata=kwargs.get("litellm_metadata"), - extra_body=extra_body, - ) + return encode(container_obj) - return container_obj + return _encode_created_container_id(pending=container_obj, encode=encode) except Exception as e: raise litellm.exception_type( diff --git a/litellm/proxy/container_endpoints/endpoints.py b/litellm/proxy/container_endpoints/endpoints.py index 4a088140725..eaa3db336a9 100644 --- a/litellm/proxy/container_endpoints/endpoints.py +++ b/litellm/proxy/container_endpoints/endpoints.py @@ -208,7 +208,7 @@ async def list_containers( # Read query parameters query_params: Final = dict(request.query_params) - data: Final[dict[str, Any]] = {"query_params": query_params} + data: Final[dict[str, Any]] = {"query_params": query_params, "model": query_params.get("model")} # Extract custom_llm_provider using priority chain custom_llm_provider: Final = ( diff --git a/litellm/router.py b/litellm/router.py index 471a1116f44..52e858a0b6c 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -6589,7 +6589,9 @@ class Router: metadata. When present, decode the ID, replace ``container_id`` with the upstream value, and route through ``_ageneric_api_call_with_fallbacks`` so deployment credentials (e.g. regional ``api_base`` for Azure) match - :meth:`_init_responses_api_endpoints`. Otherwise call the handler directly. + :meth:`_init_responses_api_endpoints`. Create/list calls carry no container ID, so + they route through the deployment named by ``model`` when the caller passes one. + Otherwise call the handler directly with global provider credentials. """ if custom_llm_provider and "custom_llm_provider" not in kwargs: kwargs["custom_llm_provider"] = custom_llm_provider @@ -6621,6 +6623,13 @@ class Router: **kwargs, ) + requested_model: Final = kwargs.get("model") + if isinstance(requested_model, str) and requested_model.strip(): + return await self._ageneric_api_call_with_fallbacks( + original_function=original_function, + **kwargs, + ) + return await original_function(**kwargs) async def _init_responses_api_endpoints( diff --git a/tests/router_unit_tests/test_router_endpoints.py b/tests/router_unit_tests/test_router_endpoints.py index d37af5b456a..6fcbd77e054 100644 --- a/tests/router_unit_tests/test_router_endpoints.py +++ b/tests/router_unit_tests/test_router_endpoints.py @@ -1324,6 +1324,65 @@ async def test_init_containers_api_endpoints_managed_id_without_model_id_applies assert call_kw["custom_llm_provider"] == "azure" +@pytest.mark.asyncio +async def test_init_containers_api_endpoints_create_with_model_uses_deployment_credentials(monkeypatch): + """ + ``POST /v1/containers`` carries no container ID, so a ``model`` in the request + body is the only way to pick a deployment. The upstream call must receive that + deployment's ``api_key``/``api_base`` instead of falling back to the global + ``OPENAI_API_KEY`` (which may be unset on the proxy). + """ + monkeypatch.delenv("OPENAI_API_KEY", raising=False) + router = Router( + model_list=[ + { + "model_name": "gpt-5.4", + "litellm_params": { + "model": "openai/gpt-5.4", + "api_key": "sk-model-list-key", + "api_base": "https://custom.openai.example/v1", + }, + } + ] + ) + mock_original_function = AsyncMock(return_value={"id": "cntr_test", "name": "Test Container"}) + + await router._init_containers_api_endpoints( + original_function=mock_original_function, + custom_llm_provider="openai", + name="Test Container", + model="gpt-5.4", + ) + + mock_original_function.assert_called_once() + call_kw = mock_original_function.call_args.kwargs + assert call_kw["api_key"] == "sk-model-list-key" + assert call_kw["api_base"] == "https://custom.openai.example/v1" + assert call_kw["model"] == "openai/gpt-5.4" + assert call_kw["name"] == "Test Container" + + +@pytest.mark.asyncio +async def test_init_containers_api_endpoints_create_without_model_calls_directly(): + """ + Without ``model`` (or with ``model=None`` as the proxy forwards it), create/list + must keep calling the handler directly with global provider credentials. + """ + router = Router(model_list=[]) + router._ageneric_api_call_with_fallbacks = AsyncMock() + mock_original_function = AsyncMock(return_value={"id": "cntr_test"}) + + await router._init_containers_api_endpoints( + original_function=mock_original_function, + custom_llm_provider="openai", + name="Test Container", + model=None, + ) + + router._ageneric_api_call_with_fallbacks.assert_not_called() + mock_original_function.assert_called_once_with(custom_llm_provider="openai", name="Test Container", model=None) + + def test_router_model_group_encrypted_content_affinity_callback_registration(): from litellm.router_utils.pre_call_checks.deployment_affinity_check import ( DeploymentAffinityCheck, diff --git a/tests/test_litellm/containers/test_container_api.py b/tests/test_litellm/containers/test_container_api.py index 885c4cd294a..8f1d87f9603 100644 --- a/tests/test_litellm/containers/test_container_api.py +++ b/tests/test_litellm/containers/test_container_api.py @@ -152,6 +152,42 @@ class TestContainerAPI: assert response.id == "cntr_async_123" assert response.name == "Async Test Container" + @pytest.mark.asyncio + async def test_acreate_container_encodes_router_model_id(self): + """ + The async handler returns a coroutine, so the managed-ID encoding must run + after it resolves. Otherwise follow-up calls (retrieve/delete/files) lose the + deployment and fall back to global provider credentials. + """ + upstream_response = ContainerObject( + id="cntr_upstream_123", + object="container", + created_at=1747857508, + status="running", + expires_after={"anchor": "last_active_at", "minutes": 20}, + last_active_at=1747857508, + name="Routed Container", + ) + + async def _resolve_upstream(): + return upstream_response + + with patch.object( # test-quality-ok: create_container does not forward a client, so the handler is the only seam + base_llm_http_handler, + "container_create_handler", + side_effect=lambda **kwargs: _resolve_upstream() if kwargs["_is_async"] else upstream_response, + ): + response = await acreate_container( + name="Routed Container", + custom_llm_provider="openai", + litellm_metadata={"model_info": {"id": "deployment-abc"}}, + ) + + decoded = ResponsesAPIRequestUtils._decode_container_id(response.id) + assert decoded["model_id"] == "deployment-abc" + assert decoded["custom_llm_provider"] == "openai" + assert decoded["response_id"] == "cntr_upstream_123" + @pytest.mark.asyncio async def test_alist_containers_basic(self): """Test basic async container listing functionality.""" From acc65b27d210158a18db74d6d0cd48d12aeb3a40 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 1 Sep 2026 23:39:42 +0000 Subject: [PATCH 34/97] fix(router): pass container create/list through when model names no deployment Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/router.py | 6 ++-- .../test_router_endpoints.py | 33 +++++++++++++++++++ .../containers/test_container_api.py | 2 +- 3 files changed, 38 insertions(+), 3 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index 52e858a0b6c..7b471a69ecf 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -6590,8 +6590,9 @@ class Router: upstream value, and route through ``_ageneric_api_call_with_fallbacks`` so deployment credentials (e.g. regional ``api_base`` for Azure) match :meth:`_init_responses_api_endpoints`. Create/list calls carry no container ID, so - they route through the deployment named by ``model`` when the caller passes one. - Otherwise call the handler directly with global provider credentials. + they route through the deployment named by ``model`` when the caller passes one, + falling back to the direct call when no deployment matches. Otherwise call the + handler directly with global provider credentials. """ if custom_llm_provider and "custom_llm_provider" not in kwargs: kwargs["custom_llm_provider"] = custom_llm_provider @@ -6627,6 +6628,7 @@ class Router: if isinstance(requested_model, str) and requested_model.strip(): return await self._ageneric_api_call_with_fallbacks( original_function=original_function, + passthrough_on_no_deployment=True, **kwargs, ) diff --git a/tests/router_unit_tests/test_router_endpoints.py b/tests/router_unit_tests/test_router_endpoints.py index 6fcbd77e054..a06aaa363ad 100644 --- a/tests/router_unit_tests/test_router_endpoints.py +++ b/tests/router_unit_tests/test_router_endpoints.py @@ -1383,6 +1383,39 @@ async def test_init_containers_api_endpoints_create_without_model_calls_directly mock_original_function.assert_called_once_with(custom_llm_provider="openai", name="Test Container", model=None) +@pytest.mark.asyncio +async def test_init_containers_api_endpoints_create_with_unknown_model_passes_through(monkeypatch): + """ + A ``model`` that names no configured deployment must not turn into a 400. The call + falls through to the handler with the caller's model and no injected deployment + credentials, matching the behaviour before model-based routing existed. + """ + monkeypatch.delenv("OPENAI_API_KEY", raising=False) + router = Router( + model_list=[ + { + "model_name": "gpt-5.4", + "litellm_params": {"model": "openai/gpt-5.4", "api_key": "sk-model-list-key"}, + } + ] + ) + mock_original_function = AsyncMock(return_value={"id": "cntr_test"}) + + await router._init_containers_api_endpoints( + original_function=mock_original_function, + custom_llm_provider="openai", + name="Test Container", + model="does-not-exist", + ) + + mock_original_function.assert_called_once() + call_kw = mock_original_function.call_args.kwargs + assert call_kw["model"] == "does-not-exist" + assert call_kw["name"] == "Test Container" + assert "api_key" not in call_kw + assert "api_base" not in call_kw + + def test_router_model_group_encrypted_content_affinity_callback_registration(): from litellm.router_utils.pre_call_checks.deployment_affinity_check import ( DeploymentAffinityCheck, diff --git a/tests/test_litellm/containers/test_container_api.py b/tests/test_litellm/containers/test_container_api.py index 8f1d87f9603..16a3844431e 100644 --- a/tests/test_litellm/containers/test_container_api.py +++ b/tests/test_litellm/containers/test_container_api.py @@ -172,7 +172,7 @@ class TestContainerAPI: async def _resolve_upstream(): return upstream_response - with patch.object( # test-quality-ok: create_container does not forward a client, so the handler is the only seam + with patch.object( # test-quality-ok: create_container exposes no client seam, only the handler base_llm_http_handler, "container_create_handler", side_effect=lambda **kwargs: _resolve_upstream() if kwargs["_is_async"] else upstream_response, From 59da6e75a50024dcca1af5efa90e4eec3409b89b Mon Sep 17 00:00:00 2001 From: tin-berri Date: Tue, 1 Sep 2026 16:50:12 -0700 Subject: [PATCH 35/97] feat(router): fall back on anthropic safeguard refusals on /v1/messages (#39157) * fix(router): resolve fallbacks against the tier a pre-routing hook selected A complexity or auto router picks a tier behind the router group name, but fallback lookup kept using kwargs["model"], which is still the router name. The tier's configured chain never ran, so a provider failure on its first hop went straight back to the client with "No fallback model group found for original model_group=smart-router". The hook assigns the selected model to a local only, and fallback resolution runs on an outer kwargs dict that **kwargs already copied, so writing it there is not visible. Record the selection in the metadata bucket instead, which is a nested dict shared by reference across those copies and is how the router already carries values back up, then key fallback lookup off it when present. Applies to the generic, context-window, content-policy and weighted-failover lookups. Reporting keeps using the router name, since that is what the caller asked for. Fixes #38832 * fix(router): annotate the recorded-selection helper with a read-only mapping record_pre_routing_selection only reads the request kwargs, writing into the nested metadata bucket it finds there, so Mapping states what it actually needs and clears the LIT001 mutable-annotation budget without a suppression. * test(router): assert the no-kwargs path leaks nothing The tolerated-None case called the helper without checking anything, which the test-quality gate counts as a test with no assertion. Assert that a fresh mapping still reads back empty, so the case proves the call is a no-op rather than only that it does not raise. * fix(router): stop declaring loop-assigned locals Final in the selection helpers Both helpers annotated a loop-assigned local as Final, which reassigns a Final on every iteration and cost three basedpyright errors. Read the buckets through a generator instead, so the write path iterates a for-target and the read path resolves in one shot with next(), which also matches the functional style the type-discipline rules ask for. * style(router): apply ruff format to the selection helpers * fix(router): derive the pre-routing tier fresh on every fallback hop The metadata buckets also carry whatever the caller sent, so an inbound pre_routing_selected_model let a client pick which fallback chain its request fell into. A fallback hop also inherited the previous hop's tier, so the second hop keyed its own failure off the tier that already failed and never ran its own chain. Clear the key at the top of async_function_with_fallbacks. Every hop re-enters there, so only the hook that routed that hop can set it. * fix(router): drop the cast at the fallback-hop clear call site * feat(router): fall back on anthropic safeguard refusals on /v1/messages --------- Co-authored-by: Priyansh Nandwana --- .../messages/streaming_iterator.py | 25 +- .../messages/utils.py | 49 ++- litellm/router.py | 192 +++++++-- .../router_utils/fallback_event_handlers.py | 85 ++++ .../anthropic_messages/anthropic_response.py | 13 +- ...test_router_anthropic_messages_fallback.py | 402 ++++++++++++++++++ .../test_fallback_event_handlers.py | 119 ++++++ tests/test_litellm/test_router.py | 135 ++++++ 8 files changed, 970 insertions(+), 50 deletions(-) create mode 100644 tests/router_unit_tests/test_router_anthropic_messages_fallback.py diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py index 45c7825344b..66e36dab2ba 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py @@ -86,22 +86,41 @@ def _decoded_sse_data_line(line: bytes) -> object | None: return None -def _anthropic_error_event_payload(chunk: object) -> Mapping[str, object] | None: +def _anthropic_event_payload(chunk: object, event_type: str) -> Mapping[str, object] | None: if isinstance(chunk, dict): - return chunk if chunk.get("type") == "error" else None + return chunk if chunk.get("type") == event_type else None if isinstance(chunk, (bytes, bytearray)): decoded_lines: Final = (_decoded_sse_data_line(line) for line in chunk.splitlines()) return next( ( candidate for candidate in decoded_lines - if isinstance(candidate, dict) and candidate.get("type") == "error" + if isinstance(candidate, dict) and candidate.get("type") == event_type ), None, ) return None +def _anthropic_error_event_payload(chunk: object) -> Mapping[str, object] | None: + return _anthropic_event_payload(chunk, "error") + + +def parse_anthropic_refusal_stop_details(chunk: object) -> Mapping[str, object] | None: + """ + Return the ``stop_details`` object of an Anthropic SSE ``message_delta`` + chunk whose delta carries ``stop_reason: "refusal"`` (a safeguard refusal: + https://platform.claude.com/docs/en/build-with-claude/refusals-and-fallback), + or None for any other chunk, a plain refusal without ``stop_details`` included. + """ + payload: Final = _anthropic_event_payload(chunk, "message_delta") + delta: Final = payload.get("delta") if payload is not None else None + if not isinstance(delta, dict) or delta.get("stop_reason") != "refusal": + return None + stop_details: Final = delta.get("stop_details") + return stop_details if isinstance(stop_details, dict) else None + + def _anthropic_error_body(chunk: object) -> Mapping[str, object] | None: """Return the ``error`` object of an Anthropic SSE ``event: error`` chunk, or None.""" payload: Final = _anthropic_error_event_payload(chunk) diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/utils.py b/litellm/llms/anthropic/experimental_pass_through/messages/utils.py index 02d82887dde..9deff950724 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/utils.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/utils.py @@ -1,11 +1,40 @@ +from collections.abc import Mapping from functools import lru_cache -from typing import Any, Final, cast, get_type_hints +from typing import TYPE_CHECKING, Any, Final, cast, get_type_hints from litellm.types.llms.anthropic import AnthropicMessagesRequestOptionalParams from litellm.types.llms.anthropic_messages.anthropic_response import ( AnthropicMessagesResponse, ) +if TYPE_CHECKING: + from litellm.exceptions import ContentPolicyViolationError + + +def get_safeguard_refusal_stop_details(response: object) -> Mapping[str, Any] | None: + """ + Return the ``stop_details`` of an Anthropic Messages response refused by a + safeguard (``stop_reason: "refusal"`` carrying ``stop_details``: + https://platform.claude.com/docs/en/build-with-claude/refusals-and-fallback), + or None for any other response, a plain refusal without ``stop_details`` included. + """ + if not isinstance(response, dict) or response.get("stop_reason") != "refusal": + return None + stop_details: Final = response.get("stop_details") + return stop_details if isinstance(stop_details, dict) else None + + +def safeguard_refusal_error(model: str, stop_details: Mapping[str, object]) -> "ContentPolicyViolationError": + """The exception a safeguard-refused Anthropic response converts into so the + content-policy fallback chain can re-dispatch it.""" + from litellm.exceptions import ContentPolicyViolationError + + return ContentPolicyViolationError( + message=f"Anthropic safeguard refusal (category: {stop_details.get('category')}).", + model=model, + llm_provider="anthropic", + ) + @lru_cache(maxsize=1) def _anthropic_messages_optional_param_keys() -> frozenset[str]: @@ -100,14 +129,12 @@ def mock_response( model=model, ) return AnthropicMessagesResponse( - **{ - "content": [{"text": mock_response, "type": "text"}], - "id": "msg_013Zva2CMHLNnXjNJJKqJ2EF", - "model": "claude-sonnet-4-20250514", - "role": "assistant", - "stop_reason": "end_turn", - "stop_sequence": None, - "type": "message", - "usage": {"input_tokens": 2095, "output_tokens": 503}, - } + content=[{"text": mock_response, "type": "text"}], + id="msg_013Zva2CMHLNnXjNJJKqJ2EF", + model="claude-sonnet-4-20250514", + role="assistant", + stop_reason="end_turn", + stop_sequence=None, + type="message", + usage={"input_tokens": 2095, "output_tokens": 503}, ) diff --git a/litellm/router.py b/litellm/router.py index 471a1116f44..6e4405ebfef 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -143,7 +143,11 @@ from litellm.router_utils.cooldown_handlers import ( from litellm.router_utils.fallback_event_handlers import ( AttemptedFallbackTargets, _check_non_standard_fallback_format, - get_fallback_model_group, + clear_pre_routing_selection, + fallback_lookup_groups, + get_fallback_model_group_for_lookup_groups, + get_pre_routing_selection, + record_pre_routing_selection, run_async_fallback, ) from litellm.router_utils.get_retry_from_policy import ( @@ -4918,6 +4922,19 @@ class Router: ) response = await response + if self._should_raise_anthropic_refusal_error( + model=model, + original_generic_function=original_generic_function, + response=response, + kwargs=kwargs, + ): + from litellm.llms.anthropic.experimental_pass_through.messages.utils import ( + safeguard_refusal_error, + ) + + refusal_details: Final = cast(dict, response["stop_details"]) # cast-ok: gate verified the shape + raise safeguard_refusal_error(model=model, stop_details=refusal_details) + self.success_calls[model_name] += 1 verbose_router_logger.info("ageneric_api_call_with_fallbacks(model=%s)\x1b[32m 200 OK\x1b[0m", model_name) @@ -4964,6 +4981,11 @@ class Router: # fallback to the original reference for any non-picklable value. # The original_generic_function is preserved so the per-attempt # helper knows which underlying API to call on fallback. + # The pre-routing hook stamps its tier selection into this bucket during the primary + # attempt; seeding it before the snapshot gives both the live kwargs and the copy a + # bucket, so the post-call carry-over below always has somewhere to read and write. + kwargs.setdefault("litellm_metadata", {}) # mutable-ok: shared bucket # rebind-ok: stamp must be readable here + fallback_kwargs: Final[dict[str, object]] = kwargs.copy() if isinstance(fallback_kwargs.get("litellm_metadata"), dict): fallback_kwargs["litellm_metadata"] = safe_deep_copy(fallback_kwargs["litellm_metadata"]) @@ -4973,6 +4995,14 @@ class Router: response: Final = await self._ageneric_api_call_with_fallbacks(original_function=original_function, **kwargs) + # The snapshot predates the pre-routing hook, so the tier it stamped into the live kwargs + # is carried over write-or-clear: a stale or caller-supplied selection left in the copy + # would key the mid-stream fallback lookup off a tier this attempt never routed to. + clear_pre_routing_selection(fallback_kwargs) + live_pre_routing_selection: Final = get_pre_routing_selection(kwargs) + if live_pre_routing_selection is not None: + record_pre_routing_selection(fallback_kwargs, live_pre_routing_selection) + if kwargs.get("stream") and isinstance(response, BaseResponsesAPIStreamingIterator): return await self._aresponses_streaming_iterator( response=response, @@ -5030,6 +5060,10 @@ class Router: from litellm.llms.anthropic.experimental_pass_through.messages.streaming_iterator import ( aclose_if_supported, parse_anthropic_error_event, + parse_anthropic_refusal_stop_details, + ) + from litellm.llms.anthropic.experimental_pass_through.messages.utils import ( + safeguard_refusal_error, ) source_iterator: Final = response @@ -5068,13 +5102,35 @@ class Router: continue if _anthropic_stream_commits_now(chunk, has_generated_content, len(buffered_lifecycle_chunks)): has_generated_content = True # rebind-ok: real content seen, or the buffer cap was hit - error_event = parse_anthropic_error_event(chunk) + # A transport can split one SSE data line across byte chunks, so pre-content + # detection parses the accumulated buffer plus the current chunk, never the + # chunk alone; the buffer is already capped, which bounds this window too. + parse_window = ( # rebind-ok: freshly computed each iteration, never carried over + b"".join(c for c in (*buffered_lifecycle_chunks, chunk) if isinstance(c, (bytes, bytearray))) # pyright: ignore[reportUnnecessaryIsInstance] # bridge-path chunks are not always bytes at runtime + if not has_generated_content and isinstance(chunk, (bytes, bytearray)) # pyright: ignore[reportUnnecessaryIsInstance] # bridge-path chunks are not always bytes at runtime + else chunk + ) + error_event = parse_anthropic_error_event(parse_window) retriable_pending_error = ( # rebind-ok: freshly computed each iteration, never carried over not has_generated_content and error_event is not None and _is_retriable_anthropic_status(error_event[2]) and not _anthropic_stream_error_is_gateway_verdict(chunk) ) + refusal_stop_details = ( # rebind-ok: freshly computed each iteration, never carried over + parse_anthropic_refusal_stop_details(parse_window) + if not has_generated_content and error_event is None + else None + ) + if refusal_stop_details is not None and self._has_content_policy_fallback(model, initial_kwargs): + refusal_error = safeguard_refusal_error(model=model, stop_details=refusal_stop_details) + raise MidStreamFallbackError( + message=refusal_error.message, + model=model, + llm_provider="anthropic", + original_exception=refusal_error, + is_pre_first_chunk=True, + ) if not has_generated_content and not retriable_pending_error and error_event is None: buffered_lifecycle_chunks = (*buffered_lifecycle_chunks, chunk) continue @@ -5186,8 +5242,13 @@ class Router: kwargs=initial_kwargs, metadata_variable_name="litellm_metadata", ) + # The content-policy dispatch branch matches on the trigger's own type, so a refusal's + # MidStreamFallbackError envelope is unwrapped here or the wrong fallback list is consulted. + fallback_trigger: Final[Exception] = ( + e.original_exception if isinstance(e.original_exception, litellm.ContentPolicyViolationError) else e + ) fallback_response = await self.async_function_with_fallbacks_common_utils( # rebind-ok: set on success - e=e, + e=fallback_trigger, disable_fallbacks=False, fallbacks=fallbacks, context_window_fallbacks=context_window_fallbacks, @@ -5243,6 +5304,11 @@ class Router: # share, leaking primary-deployment metadata into the mid-stream # fallback request. safe_deep_copy avoids deep-copying the full # kwargs (which can hold non-deepcopyable logging handles/clients). + # The pre-routing hook stamps its tier selection into this bucket during the primary + # attempt; seeding it before the snapshot gives both the live kwargs and the copy a + # bucket, so the post-call carry-over below always has somewhere to read and write. + kwargs.setdefault("litellm_metadata", {}) # mutable-ok: shared bucket # rebind-ok: stamp must be readable here + fallback_kwargs: Final[dict[str, object]] = kwargs.copy() # mutable-ok: mutated below before re-entry if isinstance(fallback_kwargs.get("litellm_metadata"), dict): fallback_kwargs["litellm_metadata"] = safe_deep_copy(fallback_kwargs["litellm_metadata"]) @@ -5252,6 +5318,14 @@ class Router: response: Final = await self._ageneric_api_call_with_fallbacks(original_function=original_function, **kwargs) + # The snapshot predates the pre-routing hook, so the tier it stamped into the live kwargs + # is carried over write-or-clear: a stale or caller-supplied selection left in the copy + # would key the mid-stream fallback lookup off a tier this attempt never routed to. + clear_pre_routing_selection(fallback_kwargs) + live_pre_routing_selection: Final = get_pre_routing_selection(kwargs) + if live_pre_routing_selection is not None: + record_pre_routing_selection(fallback_kwargs, live_pre_routing_selection) + if kwargs.get("stream") and hasattr(response, "__aiter__"): return await self._aanthropic_messages_streaming_iterator( response=cast("AsyncIterator[bytes]", response), # cast-ok: stream=True always returns a byte iterator @@ -6807,6 +6881,9 @@ class Router: original_exception: Final = e fallback_model_group = None original_model_group: Final[str | None] = kwargs.get("model") + # A pre-routing hook (complexity / auto / adaptive / quality routers) picks a tier + # behind the router name, and fallbacks are configured per tier, not per router. + lookup_groups: Final[tuple[str, ...]] = fallback_lookup_groups(kwargs, model_group) fallback_failure_exception_str = "" if disable_fallbacks is True or original_model_group is None: @@ -6851,15 +6928,15 @@ class Router: ] # Get external fallbacks — handle both standard and non-standard formats external_fallback_group: list | None = None - if fallbacks is not None and model_group is not None: + if fallbacks is not None and lookup_groups: if _check_non_standard_fallback_format(fallbacks=fallbacks): # Non-standard formats (e.g. ["claude-3-haiku"] or # [{"model": "...", "messages": [...]}]) are passed through directly external_fallback_group = fallbacks else: - external_fallback_group, generic_idx = get_fallback_model_group( + external_fallback_group, generic_idx = get_fallback_model_group_for_lookup_groups( fallbacks=fallbacks, - model_group=cast(str, model_group), + lookup_groups=lookup_groups, ) if external_fallback_group is None and generic_idx is not None: external_fallback_group = fallbacks[generic_idx]["*"] @@ -6917,9 +6994,9 @@ class Router: if isinstance(e, litellm.ContextWindowExceededError): if context_window_fallbacks is not None: context_window_fallback_model_group: Final[list[str] | None] = ( - self._get_fallback_model_group_from_fallbacks( + self._get_fallback_model_group_for_lookup_groups( fallbacks=context_window_fallbacks, - model_group=model_group, + lookup_groups=lookup_groups, ) ) if context_window_fallback_model_group is None: @@ -6950,9 +7027,9 @@ class Router: elif isinstance(e, litellm.ContentPolicyViolationError): if content_policy_fallbacks is not None: content_policy_fallback_model_group: Final[list[str] | None] = ( - self._get_fallback_model_group_from_fallbacks( + self._get_fallback_model_group_for_lookup_groups( fallbacks=content_policy_fallbacks, - model_group=model_group, + lookup_groups=lookup_groups, ) ) if content_policy_fallback_model_group is None: @@ -6979,14 +7056,14 @@ class Router: if litellm.expose_router_debug_in_errors: e.message += f"\n{error_message}" - if fallbacks is not None and model_group is not None: + if fallbacks is not None and lookup_groups: verbose_router_logger.debug("inside model fallbacks: %s", mask_sensitive_structure(fallbacks)) ( fallback_model_group, generic_fallback_idx, - ) = get_fallback_model_group( + ) = get_fallback_model_group_for_lookup_groups( fallbacks=fallbacks, # if fallbacks = [{"gpt-3.5-turbo": ["claude-3-haiku"]}] - model_group=cast(str, model_group), + lookup_groups=lookup_groups, ) ## if none, check for generic fallback if fallback_model_group is None and generic_fallback_idx is not None: @@ -6995,12 +7072,12 @@ class Router: if fallback_model_group is None: masked_fallbacks: Final = mask_sensitive_structure(fallbacks) verbose_router_logger.info( - "No fallback model group found for original model_group=%s. Fallbacks=%s", - model_group, + "No fallback model group found for lookup_groups=%s. Fallbacks=%s", + " -> ".join(lookup_groups), masked_fallbacks, ) if hasattr(original_exception, "message") and litellm.expose_router_debug_in_errors: - original_exception.message += f"No fallback model group found for original model_group={model_group}. Fallbacks={masked_fallbacks}" + original_exception.message += f"No fallback model group found for lookup_groups={' -> '.join(lookup_groups)}. Fallbacks={masked_fallbacks}" raise original_exception input_kwargs.update( @@ -7046,6 +7123,7 @@ class Router: If it fails after num_retries, fall back to another model group """ model_group: Final[str | None] = kwargs.get("model") + clear_pre_routing_selection(kwargs) # pyright: ignore[reportUnknownArgumentType] # **kwargs is untyped at this boundary if not isinstance(kwargs.get("attempted_targets"), AttemptedFallbackTargets): _fallback_metadata_key: Final = _get_router_metadata_variable_name( function_name=getattr(kwargs.get("original_function"), "__name__", None) @@ -7471,6 +7549,24 @@ class Router: break return fallback_model_group + def _get_fallback_model_group_for_lookup_groups( + self, + fallbacks: list[dict[str, list[str]]], # mutable-ok: mirrors the sibling resolver's contract + lookup_groups: tuple[str, ...], + ) -> list[str] | None: # mutable-ok: mirrors the sibling resolver's contract + """First lookup group whose exact-key chain resolves (tier first, then requested group).""" + return next( + ( + resolved + for resolved in ( + self._get_fallback_model_group_from_fallbacks(fallbacks=fallbacks, model_group=group) + for group in lookup_groups + ) + if resolved is not None + ), + None, + ) + def _get_first_default_fallback(self) -> str | None: """ Returns the first model from the default_fallbacks list, if it exists. @@ -7886,6 +7982,31 @@ class Router: return True return False + def _has_content_policy_fallback(self, model_group: str, kwargs: Mapping[str, Any]) -> bool: + """ + Whether a content-policy fallback would resolve for this request, keyed the same way + async_function_with_fallbacks_common_utils resolves it: the tier a pre-routing hook + selected wins over the requested group. Raising without this returning True would turn + a deliverable response into an error the fallback chain cannot recover from. + """ + content_policy_fallbacks: Final = kwargs.get("content_policy_fallbacks", self.content_policy_fallbacks) + if content_policy_fallbacks is not None: + return ( + self._get_fallback_model_group_for_lookup_groups( + fallbacks=content_policy_fallbacks, + lookup_groups=fallback_lookup_groups(kwargs, model_group), + ) + is not None + ) + if self._has_default_fallbacks(): + return True + verbose_router_logger.debug( + "No content-policy fallback available. Returning original response. model=%s, content_policy_fallbacks=%s", + model_group, + content_policy_fallbacks, + ) + return False + def _should_raise_content_policy_error(self, model: str, response: ModelResponse, kwargs: dict) -> bool: """ Determines if a content policy error should be raised. @@ -7898,27 +8019,26 @@ class Router: if response.choices[0].finish_reason != "content_filter": return False - content_policy_fallbacks: Final = kwargs.get("content_policy_fallbacks", self.content_policy_fallbacks) + return self._has_content_policy_fallback(model, kwargs) - ### ONLY RAISE ERROR IF CP FALLBACK AVAILABLE ### - if content_policy_fallbacks is not None: - fallback_model_group = None - for item in content_policy_fallbacks: # [{"gpt-3.5-turbo": ["gpt-4"]}] - if list(item.keys())[0] == model: - fallback_model_group = item[model] - break - - if fallback_model_group is not None: - return True - elif self._has_default_fallbacks(): # default fallbacks set - return True - - verbose_router_logger.debug( - "Content Policy Error occurred. No available fallbacks. Returning original response. model=%s, content_policy_fallbacks=%s", - model, - content_policy_fallbacks, + def _should_raise_anthropic_refusal_error( + self, model: str, original_generic_function: Callable, response: object, kwargs: Mapping[str, Any] + ) -> bool: + """ + The /v1/messages twin of _should_raise_content_policy_error: an Anthropic safeguard + refusal (stop_reason "refusal" carrying stop_details) re-enters the fallback chain only + when a content-policy fallback is configured; a plain refusal without stop_details, or + any response with nothing configured, is returned to the client unchanged. + """ + from litellm.llms.anthropic.experimental_pass_through.messages.utils import ( + get_safeguard_refusal_stop_details, ) - return False + + if getattr(original_generic_function, "__name__", "") != "anthropic_messages": + return False + if get_safeguard_refusal_stop_details(response) is None: + return False + return self._has_content_policy_fallback(model, kwargs) def _get_healthy_deployments(self, model: str, parent_otel_span: Span | None): _all_deployments: list = [] @@ -12087,6 +12207,7 @@ class Router: if pre_routing_hook_response is not None: model = pre_routing_hook_response.model messages = pre_routing_hook_response.messages + record_pre_routing_selection(request_kwargs, model) if pre_routing_hook_response.litellm_params: accepted_tier_params: Final = self._tier_params_the_target_accepts( model, pre_routing_hook_response.litellm_params, request_kwargs @@ -12202,6 +12323,7 @@ class Router: if pre_routing_hook_response is not None: model = pre_routing_hook_response.model messages = pre_routing_hook_response.messages + record_pre_routing_selection(request_kwargs, model) if pre_routing_hook_response.litellm_params: accepted_tier_params: Final = self._tier_params_the_target_accepts( model, pre_routing_hook_response.litellm_params, request_kwargs diff --git a/litellm/router_utils/fallback_event_handlers.py b/litellm/router_utils/fallback_event_handlers.py index d2842294a08..3d37ca216a7 100644 --- a/litellm/router_utils/fallback_event_handlers.py +++ b/litellm/router_utils/fallback_event_handlers.py @@ -214,6 +214,91 @@ def _check_stripped_model_group(model_group: str, fallback_key: str) -> bool: return False +PRE_ROUTING_SELECTED_MODEL_KEY: Final = "pre_routing_selected_model" +_ROUTER_METADATA_BUCKETS: Final = ("metadata", "litellm_metadata") + + +def record_pre_routing_selection(request_kwargs: Mapping[str, Any] | None, selected_model: str) -> None: + """ + Remember which model a pre-routing hook picked, so fallback lookup can key off it. + + Fallback resolution runs on an outer kwargs dict that ``**kwargs`` already copied, so + writing the model there is invisible by the time routing picks a tier. The metadata + buckets are nested dicts shared by reference across those copies, which is how the + router already carries values back up. + + The write goes through the proxy-internal bucket resolver, never into both buckets: + on /v1/messages the top-level ``metadata`` dict is the provider's own request field, + so a blanket write would forward the tier stamp upstream. + """ + from litellm.litellm_core_utils.core_helpers import get_metadata_variable_name_from_kwargs + + if request_kwargs is None: + return + bucket: Final = request_kwargs.get(get_metadata_variable_name_from_kwargs(request_kwargs)) + if isinstance(bucket, dict): + bucket[PRE_ROUTING_SELECTED_MODEL_KEY] = selected_model + + +def clear_pre_routing_selection(request_kwargs: Mapping[str, object] | None) -> None: + """ + Drop any selection the router did not make itself on this hop. + + The buckets carry whatever the caller sent, so an inbound value is the caller + choosing a fallback chain rather than the router choosing a tier. A fallback hop + also inherits the previous hop's selection, which would key its own failure off + the tier that already failed. Clearing at the start of every hop leaves only a + value the pre-routing hook wrote while routing that hop. + """ + if request_kwargs is None: + return + for bucket in (request_kwargs.get(name) for name in _ROUTER_METADATA_BUCKETS): + if isinstance(bucket, dict) and PRE_ROUTING_SELECTED_MODEL_KEY in bucket: + del bucket[PRE_ROUTING_SELECTED_MODEL_KEY] + + +def get_pre_routing_selection(kwargs: Mapping[str, Any]) -> str | None: + """The model a pre-routing hook selected for this request, if one did.""" + buckets: Final = (kwargs.get(name) for name in _ROUTER_METADATA_BUCKETS) + selections: Final = (bucket.get(PRE_ROUTING_SELECTED_MODEL_KEY) for bucket in buckets if isinstance(bucket, dict)) + return next((selected for selected in selections if isinstance(selected, str) and selected), None) + + +def fallback_lookup_groups(kwargs: Mapping[str, Any], model_group: str | None) -> tuple[str, ...]: + """ + Ordered keys for resolving a fallback chain: the tier a pre-routing hook selected wins, + and the requested group still resolves when no tier-keyed chain exists, so configs keyed + on the router name (the documented contract) keep working behind auto-routers. + """ + ordered: Final = (get_pre_routing_selection(kwargs), model_group) + return tuple(dict.fromkeys(group for group in ordered if group)) + + +def _resolved_a_specific_chain( + fallbacks: list[Any], # mutable-ok: mirrors get_fallback_model_group's contract + result: tuple[list[str] | None, int | None], # mutable-ok: mirrors get_fallback_model_group's contract +) -> bool: + resolved, generic_idx = result + if resolved is None: + return False + return generic_idx is None or resolved is not fallbacks[generic_idx]["*"] + + +def get_fallback_model_group_for_lookup_groups( + fallbacks: list[Any], # mutable-ok: mirrors get_fallback_model_group's contract + lookup_groups: tuple[str, ...], +) -> tuple[list[str] | None, int | None]: # mutable-ok: mirrors get_fallback_model_group's contract + """ + First lookup group with a specifically-keyed chain wins; the generic "*" chain applies + only after every group missed, so a catch-all cannot shadow a later group's own chain. + """ + results: Final = tuple(get_fallback_model_group(fallbacks=fallbacks, model_group=group) for group in lookup_groups) + specific: Final = next((result for result in results if _resolved_a_specific_chain(fallbacks, result)), None) + if specific is not None: + return specific + return next((result for result in results if result[0] is not None), (None, None)) + + def get_fallback_model_group(fallbacks: list[Any], model_group: str) -> tuple[list[str] | None, int | None]: """ Returns: diff --git a/litellm/types/llms/anthropic_messages/anthropic_response.py b/litellm/types/llms/anthropic_messages/anthropic_response.py index 42ca3fd6d4b..4fe1dafc73b 100644 --- a/litellm/types/llms/anthropic_messages/anthropic_response.py +++ b/litellm/types/llms/anthropic_messages/anthropic_response.py @@ -78,6 +78,16 @@ class AnthropicUsage(TypedDict, total=False): server_tool_use: NotRequired[ReadOnly[ServerToolUsage]] +class AnthropicStopDetails(TypedDict, total=False): + """ + Safeguard verdict accompanying a `stop_reason: "refusal"` response: + https://platform.claude.com/docs/en/build-with-claude/refusals-and-fallback + """ + + category: ReadOnly[str | None] + explanation: ReadOnly[str | None] + + class AnthropicMessagesResponse(TypedDict, total=False): """ Anthropic Messages API Response: https://docs.anthropic.com/en/api/messages @@ -90,7 +100,8 @@ class AnthropicMessagesResponse(TypedDict, total=False): id: str model: str | None # This represents the Model type from Anthropic role: Literal["assistant"] | None - stop_reason: Literal["end_turn", "max_tokens", "stop_sequence", "tool_use"] | None + stop_reason: Literal["end_turn", "max_tokens", "stop_sequence", "tool_use", "refusal"] | None + stop_details: NotRequired[ReadOnly[AnthropicStopDetails | None]] stop_sequence: str | None type: Literal["message"] | None usage: AnthropicUsage | None diff --git a/tests/router_unit_tests/test_router_anthropic_messages_fallback.py b/tests/router_unit_tests/test_router_anthropic_messages_fallback.py new file mode 100644 index 00000000000..0c4d1dfc21e --- /dev/null +++ b/tests/router_unit_tests/test_router_anthropic_messages_fallback.py @@ -0,0 +1,402 @@ +""" +Unit tests for safeguard-refusal fallback on the /v1/messages router surface. + +An Anthropic safeguard refusal is an HTTP 200 whose body carries +stop_reason "refusal" plus a stop_details object; the router converts it +into a ContentPolicyViolationError so the content-policy fallback chain +runs, but only when a matching fallback is configured. A plain refusal +without stop_details, or any refusal with nothing configured, must reach +the client byte-identical. + +The upstream is faked at the HTTP boundary by intercepting the third-party +transport (httpx.AsyncClient.send), so requests run litellm's real +transformation, allowlist, and streaming pipeline end to end. +""" + +import json +from typing import Any, AsyncIterator +from unittest.mock import patch + +import httpx +import pytest + +from litellm import Router +from litellm.router_utils.fallback_event_handlers import ( + PRE_ROUTING_SELECTED_MODEL_KEY, + record_pre_routing_selection, +) + +REFUSAL_RESPONSE: dict[str, Any] = { + "id": "msg_refusal", + "type": "message", + "role": "assistant", + "model": "claude-fable-5", + "content": [], + "stop_reason": "refusal", + "stop_sequence": None, + "stop_details": {"category": "cyber", "explanation": "flagged"}, + "usage": {"input_tokens": 25, "output_tokens": 1}, +} + +PLAIN_REFUSAL_RESPONSE: dict[str, Any] = {k: v for k, v in REFUSAL_RESPONSE.items() if k != "stop_details"} + +OK_RESPONSE: dict[str, Any] = { + "id": "msg_ok", + "type": "message", + "role": "assistant", + "model": "claude-opus-5", + "content": [{"type": "text", "text": "hello"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 25, "output_tokens": 2}, +} + + +def _sse(event: str, data: dict[str, Any]) -> bytes: + return f"event: {event}\ndata: {json.dumps(data)}\n\n".encode() + + +REFUSAL_STREAM_FRAMES: tuple[bytes, ...] = ( + _sse("message_start", {"type": "message_start", "message": {**REFUSAL_RESPONSE, "stop_reason": None}}), + _sse( + "message_delta", + { + "type": "message_delta", + "delta": {"stop_reason": "refusal", "stop_details": {"category": "cyber"}}, + "usage": {"output_tokens": 1}, + }, + ), + _sse("message_stop", {"type": "message_stop"}), +) + +OK_STREAM_FRAMES: tuple[bytes, ...] = ( + _sse("message_start", {"type": "message_start", "message": {**OK_RESPONSE, "stop_reason": None}}), + _sse( + "content_block_delta", + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "hello"}}, + ), + _sse("message_stop", {"type": "message_stop"}), +) + + +def _split_frames_mid_data_line(frames: tuple[bytes, ...]) -> tuple[bytes, ...]: + """Split each frame's data line in half, modeling a transport chunk boundary.""" + return tuple(part for frame in frames for part in (frame[: len(frame) // 2], frame[len(frame) // 2 :])) + + +class _FrameStream(httpx.AsyncByteStream): + def __init__(self, frames: tuple[bytes, ...]) -> None: + self._frames = frames + + async def __aiter__(self) -> AsyncIterator[bytes]: + for frame in self._frames: + yield frame + + async def aclose(self) -> None: + return None + + +class FakeAnthropicUpstream: + """Intercepts the third-party transport (httpx.AsyncClient.send): refuses on fable + models, answers on others. The router deliberately does not forward caller-injected + clients, so the transport is the seam that exercises the real litellm pipeline.""" + + def __init__( + self, + refusal_body: dict[str, Any] = REFUSAL_RESPONSE, + refusal_frames: tuple[bytes, ...] = REFUSAL_STREAM_FRAMES, + ) -> None: + self.refusal_body = refusal_body + self.refusal_frames = refusal_frames + self.calls: list[str] = [] + self.bodies: list[dict[str, Any]] = [] + + async def send(self, request: httpx.Request, **kwargs: Any) -> httpx.Response: + body = json.loads(request.content or b"{}") + model = body.get("model", "") + self.calls.append(model) + self.bodies.append(body) + refuses = "fable" in model + if body.get("stream"): + frames = self.refusal_frames if refuses else OK_STREAM_FRAMES + return httpx.Response( + 200, + stream=_FrameStream(frames), + headers={"content-type": "text/event-stream"}, + request=request, + ) + return httpx.Response(200, json=self.refusal_body if refuses else OK_RESPONSE, request=request) + + def install(self): + async def _send(_client: httpx.AsyncClient, request: httpx.Request, **kwargs: Any) -> httpx.Response: + return await self.send(request, **kwargs) + + return patch("httpx.AsyncClient.send", new=_send) + + +FABLE_TIER = { + "model_name": "fable-tier", + "litellm_params": {"model": "anthropic/claude-fable-5", "api_key": "sk-test"}, +} +OPUS_TARGET = { + "model_name": "opus-target", + "litellm_params": {"model": "anthropic/claude-opus-5", "api_key": "sk-test"}, +} + + +def _router(content_policy_fallbacks: list | None) -> Router: + return Router(model_list=[FABLE_TIER, OPUS_TARGET], content_policy_fallbacks=content_policy_fallbacks) + + +async def _collect(stream: AsyncIterator[bytes]) -> bytes: + return b"".join([chunk async for chunk in stream]) + + +@pytest.mark.asyncio +async def test_non_streaming_refusal_with_fallback_row_returns_fallback_response(): + fake = FakeAnthropicUpstream() + router = _router(content_policy_fallbacks=[{"fable-tier": ["opus-target"]}]) + + with fake.install(): + response = await router.aanthropic_messages( + model="fable-tier", max_tokens=16, messages=[{"role": "user", "content": "hi"}] + ) + + assert response["stop_reason"] == "end_turn" + assert response["id"] == "msg_ok" + assert len(fake.calls) == 2 + assert "claude-opus-5" in fake.calls[1] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "content_policy_fallbacks, upstream_body", + [ + (None, REFUSAL_RESPONSE), + ([{"unrelated-group": ["opus-target"]}], REFUSAL_RESPONSE), + ([{"fable-tier": ["opus-target"]}], PLAIN_REFUSAL_RESPONSE), + ], + ids=["nothing-configured", "row-for-other-group", "refusal-without-stop-details"], +) +async def test_non_streaming_refusal_passes_through_untouched(content_policy_fallbacks, upstream_body): + fake = FakeAnthropicUpstream(refusal_body=upstream_body) + router = _router(content_policy_fallbacks=content_policy_fallbacks) + + with fake.install(): + response = await router.aanthropic_messages( + model="fable-tier", max_tokens=16, messages=[{"role": "user", "content": "hi"}] + ) + + assert response["stop_reason"] == "refusal" + assert response.get("stop_details") == upstream_body.get("stop_details") + assert len(fake.calls) == 1 + + +@pytest.mark.asyncio +async def test_streaming_refusal_with_fallback_row_streams_fallback_frames(): + fake = FakeAnthropicUpstream() + router = _router(content_policy_fallbacks=[{"fable-tier": ["opus-target"]}]) + + with fake.install(): + stream = await router.aanthropic_messages( + model="fable-tier", max_tokens=16, stream=True, messages=[{"role": "user", "content": "hi"}] + ) + body = await _collect(stream) + + assert b'"refusal"' not in body + assert b"text_delta" in body + assert len(fake.calls) == 2 + + +@pytest.mark.asyncio +async def test_streaming_refusal_split_across_chunks_still_falls_back(): + fake = FakeAnthropicUpstream(refusal_frames=_split_frames_mid_data_line(REFUSAL_STREAM_FRAMES)) + router = _router(content_policy_fallbacks=[{"fable-tier": ["opus-target"]}]) + + with fake.install(): + stream = await router.aanthropic_messages( + model="fable-tier", max_tokens=16, stream=True, messages=[{"role": "user", "content": "hi"}] + ) + body = await _collect(stream) + + assert b'"refusal"' not in body + assert b"text_delta" in body + assert len(fake.calls) == 2 + + +@pytest.mark.asyncio +async def test_streaming_refusal_without_fallback_row_passes_frames_through(): + fake = FakeAnthropicUpstream() + router = _router(content_policy_fallbacks=None) + + with fake.install(): + stream = await router.aanthropic_messages( + model="fable-tier", max_tokens=16, stream=True, messages=[{"role": "user", "content": "hi"}] + ) + body = await _collect(stream) + + assert b'"stop_reason": "refusal"' in body + assert b"stop_details" in body + assert len(fake.calls) == 1 + + +@pytest.mark.asyncio +async def test_streaming_refusal_on_routed_tier_matches_tier_keyed_row_without_inbound_metadata(): + """The pre-routing hook's tier stamp must reach the mid-stream fallback lookup even when the + request carries no metadata bucket at all (the snapshot is taken before the request runs).""" + fake = FakeAnthropicUpstream() + smart_router = { + "model_name": "smart-router", + "litellm_params": { + "model": "auto_router/complexity_router", + "complexity_router_config": { + "tiers": {"SIMPLE": "fable-tier", "MEDIUM": "fable-tier", "COMPLEX": "fable-tier"} + }, + "complexity_router_default_model": "fable-tier", + }, + "model_info": {"id": "router-1", "db_model": True}, + } + router = Router( + model_list=[FABLE_TIER, OPUS_TARGET, smart_router], + content_policy_fallbacks=[{"fable-tier": ["opus-target"]}], + ignore_invalid_deployments=True, + ) + + with fake.install(): + stream = await router.aanthropic_messages( + model="smart-router", max_tokens=16, stream=True, messages=[{"role": "user", "content": "hi"}] + ) + body = await _collect(stream) + + assert b'"refusal"' not in body + assert b"text_delta" in body + assert len(fake.calls) == 2 + + +@pytest.mark.asyncio +async def test_caller_forged_tier_stamp_cannot_pick_the_streaming_fallback_chain(): + fake = FakeAnthropicUpstream() + router = _router(content_policy_fallbacks=[{"forged-tier": ["opus-target"]}]) + + with fake.install(): + stream = await router.aanthropic_messages( + model="fable-tier", + max_tokens=16, + stream=True, + messages=[{"role": "user", "content": "hi"}], + litellm_metadata={PRE_ROUTING_SELECTED_MODEL_KEY: "forged-tier"}, + ) + body = await _collect(stream) + + assert b'"stop_reason": "refusal"' in body + assert len(fake.calls) == 1 + + +@pytest.mark.asyncio +async def test_tier_stamp_never_reaches_provider_bound_metadata(): + """On /v1/messages the top-level metadata dict is Anthropic's own request field, so the + routed-tier stamp must never appear in any upstream body even when the client sends one.""" + fake = FakeAnthropicUpstream() + smart_router = { + "model_name": "smart-router", + "litellm_params": { + "model": "auto_router/complexity_router", + "complexity_router_config": { + "tiers": {"SIMPLE": "fable-tier", "MEDIUM": "fable-tier", "COMPLEX": "fable-tier"} + }, + "complexity_router_default_model": "fable-tier", + }, + "model_info": {"id": "router-1", "db_model": True}, + } + router = Router( + model_list=[FABLE_TIER, OPUS_TARGET, smart_router], + content_policy_fallbacks=[{"fable-tier": ["opus-target"]}], + ignore_invalid_deployments=True, + ) + + with fake.install(): + response = await router.aanthropic_messages( + model="smart-router", + max_tokens=16, + messages=[{"role": "user", "content": "hi"}], + metadata={"user_id": "u1"}, + ) + + assert response["stop_reason"] == "end_turn" + assert len(fake.bodies) == 2 + for body in fake.bodies: + assert body.get("metadata") == {"user_id": "u1"} + + +def test_record_pre_routing_selection_writes_only_the_internal_bucket(): + """The Anthropic request's own metadata field must never carry the tier stamp.""" + kwargs = {"metadata": {"user_id": "u1"}, "litellm_metadata": {}} + + record_pre_routing_selection(kwargs, "tier-x") + + assert kwargs["litellm_metadata"] == {PRE_ROUTING_SELECTED_MODEL_KEY: "tier-x"} + assert kwargs["metadata"] == {"user_id": "u1"} + + +def test_refusal_gate_keys_on_pre_routing_tier_stamp(): + router = _router(content_policy_fallbacks=[{"tier-group": ["opus-target"]}]) + + def anthropic_messages(**kwargs: Any) -> None: + return None + + refusal_kwargs = {"litellm_metadata": {PRE_ROUTING_SELECTED_MODEL_KEY: "tier-group"}} + assert ( + router._should_raise_anthropic_refusal_error( + model="router-group", + original_generic_function=anthropic_messages, + response=dict(REFUSAL_RESPONSE), + kwargs=refusal_kwargs, + ) + is True + ) + assert ( + router._should_raise_anthropic_refusal_error( + model="router-group", + original_generic_function=anthropic_messages, + response=dict(REFUSAL_RESPONSE), + kwargs={}, + ) + is False + ) + + +def test_has_content_policy_fallback_default_fallbacks_arm(): + router = Router(model_list=[OPUS_TARGET], fallbacks=[{"*": ["opus-target"]}]) + + assert router._has_content_policy_fallback("any-group", {}) is True + assert router._has_content_policy_fallback("any-group", {"content_policy_fallbacks": [{"other": ["x"]}]}) is False + + +def test_get_fallback_model_group_for_lookup_groups_orders_tier_before_requested(): + router = _router(content_policy_fallbacks=None) + fallbacks = [{"tier1": ["backup-a"]}, {"smart-router": ["backup-b"]}] + + assert router._get_fallback_model_group_for_lookup_groups( + fallbacks=fallbacks, lookup_groups=("tier1", "smart-router") + ) == ["backup-a"] + assert router._get_fallback_model_group_for_lookup_groups( + fallbacks=fallbacks, lookup_groups=("tier9", "smart-router") + ) == ["backup-b"] + assert router._get_fallback_model_group_for_lookup_groups(fallbacks=fallbacks, lookup_groups=()) is None + + +def test_refusal_gate_ignores_other_generic_call_types(): + router = _router(content_policy_fallbacks=[{"fable-tier": ["opus-target"]}]) + + def aresponses(**kwargs: Any) -> None: + return None + + assert ( + router._should_raise_anthropic_refusal_error( + model="fable-tier", + original_generic_function=aresponses, + response=dict(REFUSAL_RESPONSE), + kwargs={}, + ) + is False + ) diff --git a/tests/test_litellm/router_utils/test_fallback_event_handlers.py b/tests/test_litellm/router_utils/test_fallback_event_handlers.py index 8336926c050..894b2d9e74f 100644 --- a/tests/test_litellm/router_utils/test_fallback_event_handlers.py +++ b/tests/test_litellm/router_utils/test_fallback_event_handlers.py @@ -11,7 +11,10 @@ from litellm.router_utils.fallback_event_handlers import ( AttemptedFallbackTargets, _trigger_cooldown_for_failed_deployment, fallback_attempt_key, + clear_pre_routing_selection, get_fallback_model_group, + get_pre_routing_selection, + record_pre_routing_selection, run_async_fallback, ) @@ -1090,3 +1093,119 @@ async def test_run_async_fallback_preserves_original_model_group_on_nested_fallb metadata = router.received_kwargs["metadata"] assert metadata["attempted_fallbacks"] == 2 assert metadata["original_model_group"] == "primary-model" + + +class TestPreRoutingSelectionCarriesToFallbacks: + """#38832: a complexity/auto router picks a tier behind the router name, but fallback + lookup kept using the router name, so the tier's configured chain never ran.""" + + def test_selection_is_recorded_in_the_metadata_bucket(self): + kwargs = {"model": "smart-router", "metadata": {}} + record_pre_routing_selection(kwargs, "tier1") + assert kwargs["metadata"]["pre_routing_selected_model"] == "tier1" + assert get_pre_routing_selection(kwargs) == "tier1" + + def test_selection_is_recorded_in_the_litellm_metadata_bucket(self): + kwargs = {"model": "smart-router", "litellm_metadata": {}} + record_pre_routing_selection(kwargs, "tier2") + assert get_pre_routing_selection(kwargs) == "tier2" + + def test_a_bucket_survives_the_kwargs_copy_that_fallbacks_run_on(self): + """The bucket is shared by reference, which is the whole reason this works.""" + outer = {"model": "smart-router", "metadata": {}} + inner = {**outer} + record_pre_routing_selection(inner, "tier1") + assert get_pre_routing_selection(outer) == "tier1" + + def test_no_selection_reads_as_none(self): + assert get_pre_routing_selection({"model": "smart-router", "metadata": {}}) is None + assert get_pre_routing_selection({"model": "smart-router"}) is None + + def test_missing_kwargs_is_a_no_op(self): + """A caller with no kwargs must not raise, and must not leak the selection anywhere.""" + record_pre_routing_selection(None, "tier1") + + assert get_pre_routing_selection({}) is None + + def test_a_non_dict_bucket_is_ignored(self): + kwargs = {"model": "smart-router", "metadata": "not-a-dict"} + record_pre_routing_selection(kwargs, "tier1") + assert get_pre_routing_selection(kwargs) is None + + def test_fallbacks_resolve_against_the_selected_tier(self): + """The lookup the router performs, keyed on the tier rather than the router name.""" + fallbacks = [{"tier1": ["backup-a", "backup-b"]}, {"tier2": ["backup-c"]}] + assert get_fallback_model_group(fallbacks=fallbacks, model_group="tier1")[0] == ["backup-a", "backup-b"] + assert get_fallback_model_group(fallbacks=fallbacks, model_group="smart-router")[0] is None + + +class TestPreRoutingSelectionIsPerHop: + """#38832 review: the buckets also carry whatever the caller sent, and a fallback hop + inherits the previous hop's tier, so a hop must start without a selection.""" + + def test_a_caller_supplied_selection_is_dropped(self): + kwargs = {"model": "plain", "metadata": {"pre_routing_selected_model": "tier1"}} + + clear_pre_routing_selection(kwargs) + + assert get_pre_routing_selection(kwargs) is None + assert "pre_routing_selected_model" not in kwargs["metadata"] + + def test_both_buckets_are_cleared(self): + kwargs = { + "metadata": {"pre_routing_selected_model": "tier1"}, + "litellm_metadata": {"pre_routing_selected_model": "tier2"}, + } + + clear_pre_routing_selection(kwargs) + + assert get_pre_routing_selection(kwargs) is None + + def test_the_rest_of_the_bucket_is_left_alone(self): + kwargs = {"metadata": {"pre_routing_selected_model": "tier1", "tags": ["a"]}} + + clear_pre_routing_selection(kwargs) + + assert kwargs["metadata"] == {"tags": ["a"]} + + def test_clearing_is_a_no_op_without_a_usable_bucket(self): + kwargs = {"model": "plain", "metadata": "not-a-dict"} + + clear_pre_routing_selection(None) + clear_pre_routing_selection(kwargs) + + assert kwargs == {"model": "plain", "metadata": "not-a-dict"} + + def test_a_selection_recorded_after_clearing_is_kept(self): + """Clearing runs before routing, so the hook's own write must survive it.""" + kwargs = {"model": "smart-router", "metadata": {"pre_routing_selected_model": "stale"}} + + clear_pre_routing_selection(kwargs) + record_pre_routing_selection(kwargs, "tier1") + + assert get_pre_routing_selection(kwargs) == "tier1" + + +class TestOrderedFallbackLookupGroups: + def test_tier_first_then_requested_group_deduped(self): + from litellm.router_utils.fallback_event_handlers import ( + PRE_ROUTING_SELECTED_MODEL_KEY, + fallback_lookup_groups, + ) + + kwargs = {"litellm_metadata": {PRE_ROUTING_SELECTED_MODEL_KEY: "tier1"}} + assert fallback_lookup_groups(kwargs, "smart-router") == ("tier1", "smart-router") + assert fallback_lookup_groups(kwargs, "tier1") == ("tier1",) + assert fallback_lookup_groups({}, "smart-router") == ("smart-router",) + assert fallback_lookup_groups({}, None) == () + + def test_first_resolving_group_wins_and_generic_idx_survives_a_miss(self): + from litellm.router_utils.fallback_event_handlers import ( + get_fallback_model_group_for_lookup_groups, + ) + + fallbacks = [{"tier1": ["backup-a"]}, {"smart-router": ["backup-b"]}, {"*": ["backup-c"]}] + assert get_fallback_model_group_for_lookup_groups(fallbacks, ("tier1", "smart-router")) == (["backup-a"], None) + assert get_fallback_model_group_for_lookup_groups(fallbacks, ("tier9", "smart-router")) == (["backup-b"], None) + assert get_fallback_model_group_for_lookup_groups(fallbacks, ("tier9", "no-such")) == (["backup-c"], 2) + assert get_fallback_model_group_for_lookup_groups([{"tier1": ["backup-a"]}], ("no", "nope")) == (None, None) diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 44c1cdbff06..84f6344be35 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -11595,3 +11595,138 @@ class TestTierParamsTheTargetAccepts: accepted = router._tier_params_the_target_accepts("no-such-group", {"reasoning_effort": "max"}, {}) assert accepted == {"reasoning_effort": "max"} + + +class TestPreRoutingTierDrivesFallbacks: + """#38832: a complexity/auto router picks a tier behind the router name, but fallback + lookup stayed on the router name, so the tier's configured chain never ran and a + provider failure on the tier's first hop was returned to the client.""" + + class _TierRouter(litellm.Router): + async def async_pre_routing_hook( + self, model, request_kwargs, messages=None, input=None, specific_deployment=False + ): + from litellm.types.router import PreRoutingHookResponse + + if model == "smart-router": + return PreRoutingHookResponse(model="tier1", messages=messages) + return None + + @classmethod + def _router(cls, fallbacks) -> "litellm.Router": + return cls._TierRouter( + model_list=[ + { + "model_name": "smart-router", + "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "sk-x"}, + }, + { + "model_name": "tier1", + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_key": "sk-x", + "mock_response": "litellm.RateLimitError", + }, + }, + { + "model_name": "backup-a", + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_key": "sk-x", + "mock_response": "from backup-a", + }, + }, + { + "model_name": "backup-b", + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_key": "sk-x", + "mock_response": "from backup-b", + }, + }, + { + "model_name": "failing-backup", + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_key": "sk-x", + "mock_response": "litellm.RateLimitError", + }, + }, + { + "model_name": "plain", + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_key": "sk-x", + "mock_response": "litellm.RateLimitError", + }, + }, + ], + fallbacks=fallbacks, + num_retries=0, + ) + + @pytest.mark.asyncio + async def test_the_selected_tier_fallback_chain_runs(self): + router = self._router([{"tier1": ["backup-a"]}]) + + response = await router.acompletion( + model="smart-router", messages=[{"role": "user", "content": "hi"}] + ) + + assert response.choices[0].message.content == "from backup-a" + + @pytest.mark.asyncio + async def test_a_chain_keyed_on_the_router_name_is_not_used(self): + """The router name has no chain of its own, so nothing should rescue this call.""" + router = self._router([{"tier2": ["backup-a"]}]) + + with pytest.raises(litellm.RateLimitError): + await router.acompletion(model="smart-router", messages=[{"role": "user", "content": "hi"}]) + + @pytest.mark.asyncio + async def test_a_chain_keyed_on_the_router_name_rescues_when_no_tier_chain_exists(self): + """The documented contract: configs keyed on the requested name keep working behind auto-routers.""" + router = self._router([{"smart-router": ["backup-a"]}]) + + response = await router.acompletion(model="smart-router", messages=[{"role": "user", "content": "hi"}]) + + assert response.choices[0].message.content == "from backup-a" + + @pytest.mark.asyncio + async def test_the_tier_chain_wins_over_the_router_name_chain(self): + router = self._router([{"tier1": ["backup-a"]}, {"smart-router": ["backup-b"]}]) + + response = await router.acompletion(model="smart-router", messages=[{"role": "user", "content": "hi"}]) + + assert response.choices[0].message.content == "from backup-a" + + @pytest.mark.asyncio + async def test_a_request_without_a_pre_routing_hook_still_uses_its_own_group(self): + router = self._router([{"tier1": ["backup-a"]}]) + + response = await router.acompletion( + model="tier1", messages=[{"role": "user", "content": "hi"}] + ) + + assert response.choices[0].message.content == "from backup-a" + + @pytest.mark.asyncio + async def test_a_caller_cannot_pick_the_chain_by_sending_the_selection(self): + """The metadata bucket carries caller-supplied keys, so only the hook may set the tier.""" + router = self._router([{"tier1": ["backup-a"]}]) + + with pytest.raises(litellm.RateLimitError): + await router.acompletion( + model="plain", + messages=[{"role": "user", "content": "hi"}], + metadata={"pre_routing_selected_model": "tier1"}, + ) + + @pytest.mark.asyncio + async def test_each_fallback_hop_resolves_its_own_chain(self): + """The second hop must key off the group it is running, not the tier that failed.""" + router = self._router([{"tier1": ["failing-backup"]}, {"failing-backup": ["backup-b"]}]) + + response = await router.acompletion(model="smart-router", messages=[{"role": "user", "content": "hi"}]) + + assert response.choices[0].message.content == "from backup-b" 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 36/97] 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 8c7fe00d80a9118baf5d2b8f87d24dc53644a8fd Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 1 Sep 2026 17:37:07 -0700 Subject: [PATCH 37/97] fix: compare stream event types by equality so typed completed events keep their usage --- .../base_llm/guardrail_translation/utils.py | 2 +- ...test_openai_responses_guardrail_handler.py | 32 +++++++++++++++++++ 2 files changed, 33 insertions(+), 1 deletion(-) diff --git a/litellm/llms/base_llm/guardrail_translation/utils.py b/litellm/llms/base_llm/guardrail_translation/utils.py index d30cdcecff5..9b6f9c47105 100644 --- a/litellm/llms/base_llm/guardrail_translation/utils.py +++ b/litellm/llms/base_llm/guardrail_translation/utils.py @@ -171,7 +171,7 @@ def blocked_responses_stream_usage(original_response: object) -> ResponseAPIUsag ( response for item in reversed(original_response) - if str(stream_item_field(item, "type") or "") == "response.completed" + if stream_item_field(item, "type") == "response.completed" and (response := stream_item_field(item, "response")) is not None ), None, diff --git a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py index d68ca0fdfb6..315b6948bd8 100644 --- a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py +++ b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py @@ -1389,6 +1389,38 @@ class TestBuildBlockSseChunks: assert completed["output"][0]["content"][0]["text"] == "Blocked by policy." assert completed["usage"] == {"input_tokens": 7, "output_tokens": 21, "total_tokens": 28} + def test_continuation_reads_usage_from_typed_completed_event(self): + from litellm.types.llms.openai import ( + ResponseCompletedEvent, + ResponsesAPIResponse, + ResponsesAPIStreamEvents, + ) + + handler = OpenAIResponsesHandler() + original = [ + ResponseCompletedEvent( + type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED, + response=ResponsesAPIResponse.model_validate( + { + "id": "resp_live", + "created_at": 1, + "model": "gpt-5.4-mini", + "output": [], + "usage": {"input_tokens": 7, "output_tokens": 21, "total_tokens": 28}, + } + ), + ) + ] + payloads = self._payloads( + handler.build_block_sse_chunks( + self._exc(original_response=original), stream_started=True, responses_so_far=[] + ) + ) + completed = payloads[-1]["response"] + assert completed["usage"]["input_tokens"] == 7 + assert completed["usage"]["output_tokens"] == 21 + assert completed["usage"]["total_tokens"] == 28 + def test_continuation_closes_open_item_given_pydantic_events_with_enum_types(self): from litellm.types.llms.openai import ( BaseLiteLLMOpenAIResponseObject, From b0041f32a2d7ba9a99e0d9eb0f5402df7237200a Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 1 Sep 2026 17:45:59 -0700 Subject: [PATCH 38/97] fix(helm): reuse the generated master key Secret on helm upgrade (#39219) The generated masterkey Secret rendered a fresh randAlphaNum value on every release, so any helm upgrade with masterkeySecretName and masterkey unset rotated the master key and invalidated every client holding the old one. Look up the existing Secret in the release namespace and reuse its value, falling back to a random key only on first install. Co-authored-by: yassin Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- helm/litellm-helm/README.md | 4 +- .../templates/secret-masterkey.yaml | 6 ++- .../tests/masterkey-secret_tests.yaml | 47 +++++++++++++++++++ 3 files changed, 54 insertions(+), 3 deletions(-) diff --git a/helm/litellm-helm/README.md b/helm/litellm-helm/README.md index b242373de5d..bf4089404db 100644 --- a/helm/litellm-helm/README.md +++ b/helm/litellm-helm/README.md @@ -26,7 +26,7 @@ If `db.useStackgresOperator` is used (not yet implemented): | `replicaCount` | The number of LiteLLM Proxy pods to be deployed | `1` | | `masterkeySecretName` | The name of the Kubernetes Secret that contains the Master API Key for LiteLLM. If not specified, use the generated secret name. | N/A | | `masterkeySecretKey` | The key within the Kubernetes Secret that contains the Master API Key for LiteLLM. If not specified, use `masterkey` as the key. | N/A | -| `masterkey` | The Master API Key for LiteLLM. If not specified, a random key in the `sk-...` format is generated. | N/A | +| `masterkey` | The Master API Key for LiteLLM. If not specified, a random key in the `sk-...` format is generated on first install and reused on upgrades. | N/A | | `environmentSecrets` | An optional array of Secret object names. The keys and values in these secrets will be presented to the LiteLLM proxy pod as environment variables. See below for an example Secret object. | `[]` | | `environmentConfigMaps` | An optional array of ConfigMap object names. The keys and values in these configmaps will be presented to the LiteLLM proxy pod as environment variables. See below for an example Secret object. | `[]` | | `image.repository` | LiteLLM Proxy image repository | `ghcr.io/berriai/litellm` | @@ -212,6 +212,8 @@ service, the **Proxy Endpoint** should be set to `http://-litellm:4000` The **Proxy Key** is the value specified for `masterkey` or, if a `masterkey` was not provided to the helm command line, the `masterkey` is a randomly generated string in the `sk-...` format stored in the `-litellm-masterkey` Kubernetes Secret. +The key is generated once on the first install; later `helm upgrade` runs reuse the +value already in that Secret, so upgrading never rotates the master key. ```bash kubectl -n litellm get secret -litellm-masterkey -o jsonpath="{.data.masterkey}" diff --git a/helm/litellm-helm/templates/secret-masterkey.yaml b/helm/litellm-helm/templates/secret-masterkey.yaml index 7c8560cc2cc..60ab4e74c6b 100644 --- a/helm/litellm-helm/templates/secret-masterkey.yaml +++ b/helm/litellm-helm/templates/secret-masterkey.yaml @@ -1,9 +1,11 @@ {{- if not .Values.masterkeySecretName }} -{{ $masterkey := (.Values.masterkey | default (printf "sk-%s" (randAlphaNum 18))) }} +{{- $secretName := printf "%s-masterkey" (include "litellm.fullname" .) }} +{{- $existing := lookup "v1" "Secret" .Release.Namespace $secretName }} +{{- $masterkey := .Values.masterkey | default (dig "data" "masterkey" "" $existing | b64dec) | default (printf "sk-%s" (randAlphaNum 18)) }} apiVersion: v1 kind: Secret metadata: - name: {{ include "litellm.fullname" . }}-masterkey + name: {{ $secretName }} data: masterkey: {{ $masterkey | b64enc }} type: Opaque diff --git a/helm/litellm-helm/tests/masterkey-secret_tests.yaml b/helm/litellm-helm/tests/masterkey-secret_tests.yaml index bbbade9d802..296f26755b8 100644 --- a/helm/litellm-helm/tests/masterkey-secret_tests.yaml +++ b/helm/litellm-helm/tests/masterkey-secret_tests.yaml @@ -15,6 +15,53 @@ tests: # Note: The masterkey is generated as "sk-<18-random-chars>" in plain text, # but stored as base64 encoded in Kubernetes secret (requirement). # "sk-" base64 encodes to "c2st", so we check for "^c2st" pattern. + - it: should reuse the master key already stored in the cluster instead of generating a new one on upgrade + template: secret-masterkey.yaml + set: + masterkeySecretName: "" + kubernetesProvider: + scheme: + "v1/Secret": + gvr: + version: "v1" + resource: "secrets" + namespaced: true + objects: + - kind: Secret + apiVersion: v1 + metadata: + name: RELEASE-NAME-litellm-masterkey + namespace: NAMESPACE + data: + masterkey: c2stZXhpc3Rpbmcta2V5 + asserts: + - equal: + path: data.masterkey + value: c2stZXhpc3Rpbmcta2V5 + - it: should let an explicit masterkey value override the one already stored in the cluster + template: secret-masterkey.yaml + set: + masterkeySecretName: "" + masterkey: sk-explicit + kubernetesProvider: + scheme: + "v1/Secret": + gvr: + version: "v1" + resource: "secrets" + namespaced: true + objects: + - kind: Secret + apiVersion: v1 + metadata: + name: RELEASE-NAME-litellm-masterkey + namespace: NAMESPACE + data: + masterkey: c2stZXhpc3Rpbmcta2V5 + asserts: + - equal: + path: data.masterkey + value: c2stZXhwbGljaXQ= - it: should not create a secret if masterkeySecretName is set template: secret-masterkey.yaml set: From 69029c139e3efb4f750594bceb536d6ef944cbbb Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 1 Sep 2026 17:48:01 -0700 Subject: [PATCH 39/97] fix(mcp): report per-server outcomes in aggregate REST tools/list (#39232) GET /mcp-rest/tools/list without server_id returned only the tools of the servers that answered and silently dropped any server whose listing failed (for example an OAuth-protected server without credentials), so clients could not tell a partial listing from a complete one. The aggregate response now carries a server_outcomes map keyed by server alias with the same classified outcome (ok/auth_required/forbidden/...) that the MCP protocol path already puts in _meta. Healthy tools and the HTTP 200 status are unchanged. Co-authored-by: yassin Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../mcp_server/rest_endpoints.py | 71 +++++++++++-------- .../mcp_server/test_rest_endpoints.py | 11 ++- 2 files changed, 51 insertions(+), 31 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index 2b89dba0e4f..d1ef73a15cd 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -1,7 +1,8 @@ import asyncio import importlib -from collections.abc import Awaitable, Callable, Mapping +from collections.abc import Awaitable, Callable, Mapping, Sequence from datetime import datetime +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal import anyio @@ -20,8 +21,11 @@ from litellm.proxy._experimental.mcp_server.exceptions import ( MCPUpstreamAuthError, ) from litellm.proxy._experimental.mcp_server.faults.list_outcomes import ( + ServerListOk, + ServerOutcome, classify_list_exception, list_fault_http_status, + outcome_wire_value, ) from litellm.proxy._experimental.mcp_server.ui_session_utils import ( acting_user_auth, @@ -99,6 +103,7 @@ if MCP_AVAILABLE: ListMCPToolsRestAPIResponseObject, MCPInfo, MCPServer, + _aggregate_server_key, # pyright: ignore[reportPrivateUsage] # same per-server key as the tools/list _meta outcomes _apply_toolset_scope, _fire_mcp_tool_call_logging, execute_mcp_tool, @@ -803,9 +808,6 @@ if MCP_AVAILABLE: list(allowed_server_ids_set), _rest_client_ip ) - list_tools_result: Final = [] - error_message = None - # If server_id is specified, only query that specific server if server_id: return await _list_tools_for_single_server( @@ -849,22 +851,19 @@ if MCP_AVAILABLE: else {} ) - # Query all servers the user has access to - errors: Final = [] - for allowed_server_id in allowed_server_ids: - server = global_mcp_server_manager.get_mcp_server_by_id(allowed_server_id) - if server is None: - continue - - server_auth_header = _get_server_auth_header(server, mcp_server_auth_headers, mcp_auth_header) - user_oauth_extra_headers = await _get_user_oauth_extra_headers( + async def list_server( + server: MCPServer, + ) -> tuple[Sequence[ListMCPToolsRestAPIResponseObject], ServerOutcome]: + server_auth_header: Final = _get_server_auth_header( + server, mcp_server_auth_headers, mcp_auth_header + ) + user_oauth_extra_headers: Final = await _get_user_oauth_extra_headers( server, user_api_key_dict, prefetched_creds=prefetched_oauth_creds, ) - try: - tools_result = await _get_tools_for_single_server( + tools_result: Final = await _get_tools_for_single_server( server, server_auth_header, raw_headers_from_request, @@ -872,24 +871,36 @@ if MCP_AVAILABLE: extra_headers=user_oauth_extra_headers, apply_tool_filters=apply_tool_filters, ) - list_tools_result.extend(tools_result) except Exception as e: verbose_logger.exception("Error getting tools from %s: %s", server.name, e) - errors.append( - f"{get_server_prefix(server)}: {classify_list_exception(e).tag}" - if isinstance(e, (MCPServerListError, MCPUpstreamAuthError)) - else f"{get_server_prefix(server)}: {e}" - ) - continue + return (), classify_list_exception(e) + return tools_result, ServerListOk(tool_count=len(tools_result)) - if errors and not list_tools_result: - error_message = "Failed to get tools from servers: " + "; ".join(errors) - - return { - "tools": list_tools_result, - "error": "partial_failure" if error_message else None, - "message": (error_message if error_message else "Successfully retrieved tools"), - } + # Query all servers the user has access to + queried_servers: Final = tuple( + server + for server in map(global_mcp_server_manager.get_mcp_server_by_id, allowed_server_ids) + if server is not None + ) + listings: Final = tuple([await list_server(server) for server in queried_servers]) + list_tools_result: Final = [tool for tools, _ in listings for tool in tools] + server_outcomes: Final = MappingProxyType( + {_aggregate_server_key(server): outcome for server, (_, outcome) in zip(queried_servers, listings)} + ) + errors: Final = tuple( + f"{key}: {outcome.tag}" for key, outcome in server_outcomes.items() if outcome.tag != "ok" + ) + error_message: Final = ( + "Failed to get tools from servers: " + "; ".join(errors) + if errors and not list_tools_result + else None + ) + return { + "tools": list_tools_result, + "error": "partial_failure" if error_message else None, + "message": (error_message if error_message else "Successfully retrieved tools"), + "server_outcomes": {key: outcome_wire_value(outcome) for key, outcome in server_outcomes.items()}, + } except MCPUpstreamAuthError as e: # Surface upstream pass-through 401/403 challenges to the client so diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py index e36bef229f6..0480bbc40a7 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py @@ -1431,7 +1431,11 @@ class TestListToolsRestAPI: async def test_aggregate_list_absorbs_one_server_auth_failure(self, monkeypatch): """The multi-server aggregate listing degrades a server whose upstream rejects auth to an empty contribution and still returns the healthy - server's tools with a 200, rather than surfacing a 401.""" + server's tools with a 200, rather than surfacing a 401. The absorbed + server must still show up as a classified per-server outcome so a REST + caller can tell "needs upstream auth" apart from "has no tools".""" + from pydantic import TypeAdapter + from litellm.proxy._experimental.mcp_server.exceptions import ( MCPUpstreamAuthError, ) @@ -1497,6 +1501,11 @@ class TestListToolsRestAPI: assert result["tools"] == ["good-tool"] assert result["error"] is None + wire_body = json.loads(TypeAdapter(dict).dump_json(result)) + assert wire_body["server_outcomes"] == { + "good": {"status": "ok", "tool_count": 1}, + "bad": {"status": "auth_required", "http_status": 401}, + } async def test_name_resolution_finds_server_by_uuid(self, monkeypatch): """When server_id is a name string, it should be resolved to its UUID From 04a25083a64fb4a2287d39db535ce680c3a17d93 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 1 Sep 2026 17:49:06 -0700 Subject: [PATCH 40/97] fix(cost-map): retry transient boot fetch failures and recover config deployments dropped by a stale cost map (#39230) Co-authored-by: yassin Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../litellm_core_utils/get_model_cost_map.py | 126 +++++++++++++----- litellm/router.py | 23 +++- .../test_get_model_cost_map.py | 106 ++++++++++++++- .../test_router_model_cost_isolation.py | 80 +++++++++++ 4 files changed, 291 insertions(+), 44 deletions(-) diff --git a/litellm/litellm_core_utils/get_model_cost_map.py b/litellm/litellm_core_utils/get_model_cost_map.py index 2043a9e2f89..9cba5db8ab7 100644 --- a/litellm/litellm_core_utils/get_model_cost_map.py +++ b/litellm/litellm_core_utils/get_model_cost_map.py @@ -12,6 +12,7 @@ import asyncio import json import os import random +import time from collections.abc import Awaitable, Callable from dataclasses import dataclass from datetime import datetime, timezone @@ -154,18 +155,6 @@ class GetModelCostMap: return True - @staticmethod - def fetch_remote_model_cost_map(url: str, timeout: int = 5) -> dict: - """ - Fetch the model cost map from a remote URL. - - Returns the parsed JSON dict. Raises on network/parse errors - (caller is expected to handle). - """ - response: Final = httpx.get(url, timeout=timeout) - response.raise_for_status() - return response.json() - RETRYABLE_FETCH_STATUS_CODES: Final = frozenset({429, 500, 502, 503, 504}) MODEL_COST_MAP_FETCH_MAX_ATTEMPTS: Final = 3 @@ -212,6 +201,13 @@ class _AsyncGetClient(Protocol): def get(self, url: str, *, timeout: float | None = None) -> Awaitable[httpx.Response]: ... +class _SyncGetClient(Protocol): + def get(self, url: str, *, timeout: float | None = None) -> httpx.Response: ... + + +_FetchAttemptOutcome = ModelCostMapReloaded | ModelCostMapReloadUnavailable | _FetchAttemptRetryable + + def _default_reload_client() -> _AsyncGetClient: from litellm.llms.custom_httpx.http_handler import get_async_httpx_client from litellm.types.llms.custom_http import httpxSpecialProvider @@ -219,13 +215,30 @@ def _default_reload_client() -> _AsyncGetClient: return get_async_httpx_client(llm_provider=httpxSpecialProvider.ModelCostMap) -async def _attempt_fetch( - client: _AsyncGetClient, url: str, timeout: int -) -> ModelCostMapReloaded | ModelCostMapReloadUnavailable | _FetchAttemptRetryable: +def _classify_fetch_error(error: httpx.HTTPError | httpx.InvalidURL, url: str) -> _FetchAttemptOutcome: + reason: Final = f"{type(error).__name__} fetching {url}: {error}" + if isinstance(error, (httpx.InvalidURL, httpx.UnsupportedProtocol)): + return ModelCostMapReloadUnavailable(reason=reason) + return _FetchAttemptRetryable(reason=reason, retry_after_seconds=None) + + +async def _attempt_fetch(client: _AsyncGetClient, url: str, timeout: int) -> _FetchAttemptOutcome: try: response: Final = await client.get(url, timeout=timeout) - except httpx.HTTPError as e: - return _FetchAttemptRetryable(reason=f"{type(e).__name__} fetching {url}: {e}", retry_after_seconds=None) + except (httpx.HTTPError, httpx.InvalidURL) as e: + return _classify_fetch_error(e, url) + return _classify_fetch_response(response, url) + + +def _attempt_fetch_sync(client: _SyncGetClient, url: str, timeout: int) -> _FetchAttemptOutcome: + try: + response: Final = client.get(url, timeout=timeout) + except (httpx.HTTPError, httpx.InvalidURL) as e: + return _classify_fetch_error(e, url) + return _classify_fetch_response(response, url) + + +def _classify_fetch_response(response: httpx.Response, url: str) -> _FetchAttemptOutcome: if response.status_code in RETRYABLE_FETCH_STATUS_CODES: return _FetchAttemptRetryable( reason=f"HTTP {response.status_code} from {url}", @@ -242,6 +255,22 @@ async def _attempt_fetch( return ModelCostMapReloaded(model_cost_map=parsed) +def _next_retry_wait( + outcome: _FetchAttemptRetryable, attempt: int, max_attempts: int, rng: random.Random +) -> float | ModelCostMapReloadUnavailable: + if attempt == max_attempts: + return ModelCostMapReloadUnavailable(reason=f"{outcome.reason} (after {max_attempts} attempts)") + wait_seconds: Final = _retry_wait_seconds(outcome=outcome, attempt=attempt, rng=rng) + verbose_logger.warning( + "LiteLLM: model cost map fetch attempt %d/%d failed (%s); retrying in %.1fs", + attempt, + max_attempts, + outcome.reason, + wait_seconds, + ) + return wait_seconds + + async def _fetch_remote_model_cost_map_with_retry( url: str, timeout: int, @@ -254,20 +283,32 @@ async def _fetch_remote_model_cost_map_with_retry( outcome = await _attempt_fetch(client=client, url=url, timeout=timeout) if not isinstance(outcome, _FetchAttemptRetryable): return outcome - if attempt == max_attempts: - return ModelCostMapReloadUnavailable(reason=f"{outcome.reason} (after {max_attempts} attempts)") - wait_seconds = _retry_wait_seconds(outcome=outcome, attempt=attempt, rng=rng) - verbose_logger.warning( - "LiteLLM: model cost map fetch attempt %d/%d failed (%s); retrying in %.1fs", - attempt, - max_attempts, - outcome.reason, - wait_seconds, - ) + wait_seconds = _next_retry_wait(outcome=outcome, attempt=attempt, max_attempts=max_attempts, rng=rng) + if isinstance(wait_seconds, ModelCostMapReloadUnavailable): + return wait_seconds await sleep(wait_seconds) return ModelCostMapReloadUnavailable(reason="model cost map fetch failed") +def _fetch_remote_model_cost_map_with_retry_sync( + url: str, + timeout: int, + max_attempts: int, + sleep: Callable[[float], None], + rng: random.Random, + client: _SyncGetClient, +) -> ModelCostMapReloadResult: + for attempt in range(1, max_attempts + 1): + outcome = _attempt_fetch_sync(client=client, url=url, timeout=timeout) + if not isinstance(outcome, _FetchAttemptRetryable): + return outcome + wait_seconds = _next_retry_wait(outcome=outcome, attempt=attempt, max_attempts=max_attempts, rng=rng) + if isinstance(wait_seconds, ModelCostMapReloadUnavailable): + return wait_seconds + sleep(wait_seconds) + return ModelCostMapReloadUnavailable(reason="model cost map fetch failed") + + async def refetch_model_cost_map( url: str, timeout: int = 5, @@ -423,13 +464,21 @@ def _finalize_model_cost_map(model_cost: dict) -> dict: return _expand_model_aliases(model_cost) -def get_model_cost_map(url: str) -> dict: +def get_model_cost_map( + url: str, + timeout: int = 5, + max_attempts: int = MODEL_COST_MAP_FETCH_MAX_ATTEMPTS, + sleep: Callable[[float], None] = time.sleep, + rng: random.Random | None = None, + client: "_SyncGetClient | None" = None, +) -> dict: """ Public entry point — returns the model cost map dict. 1. If ``LITELLM_LOCAL_MODEL_COST_MAP`` is set, uses the local backup only. - 2. Otherwise fetches from ``url``, validates integrity, and falls back - to the local backup on any failure. + 2. Otherwise fetches from ``url``, retrying transient HTTP errors + (429/5xx/transport) with Retry-After-aware backoff, validates + integrity, and falls back to the local backup on any failure. Only the backup model count is cached (a single int) for validation. The full backup dict is only parsed when it must be *returned* as a @@ -448,17 +497,24 @@ def get_model_cost_map(url: str) -> dict: _cost_map_source_info.url = url _cost_map_source_info.is_env_forced = False - try: - content: Final = GetModelCostMap.fetch_remote_model_cost_map(url) - except Exception as e: + result: Final = _fetch_remote_model_cost_map_with_retry_sync( + url=url, + timeout=timeout, + max_attempts=max_attempts, + sleep=sleep, + rng=rng if rng is not None else random.Random(), + client=client if client is not None else httpx, + ) + if isinstance(result, ModelCostMapReloadUnavailable): verbose_logger.warning( "LiteLLM: Failed to fetch remote model cost map from %s: %s. Falling back to local backup.", url, - str(e), + result.reason, ) _cost_map_source_info.source = "local" - _cost_map_source_info.fallback_reason = f"Remote fetch failed: {e}" + _cost_map_source_info.fallback_reason = f"Remote fetch failed: {result.reason}" return _finalize_model_cost_map(GetModelCostMap.load_local_model_cost_map()) + content: Final = result.model_cost_map # Validate using cached count (cheap int comparison, no file I/O) if not GetModelCostMap.validate_model_cost_map( diff --git a/litellm/router.py b/litellm/router.py index 6e4405ebfef..23d8907fb49 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -22,7 +22,7 @@ import traceback import weakref from collections import defaultdict from collections.abc import AsyncGenerator, AsyncIterator, Callable, Generator, Mapping, Sequence -from functools import lru_cache +from functools import lru_cache, partial from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, Optional, TypeAlias, TypeVar, Union, cast @@ -825,6 +825,7 @@ class Router: self._zero_cost_cache: dict[str, bool] = {} self._routing_group_rows: tuple[DeploymentTypedDict, ...] | None = None self._init_routing_groups(None) + self._provider_unresolved_deployments: tuple[Callable[[], Deployment | None], ...] = () self.deployment_affinity_ttl_seconds = deployment_affinity_ttl_seconds self.model_group_affinity_config = model_group_affinity_config @@ -8472,6 +8473,19 @@ class Router: return deployment except Exception as e: if self.ignore_invalid_deployments: + if isinstance(e, litellm.BadRequestError): + self._provider_unresolved_deployments = ( + *self._provider_unresolved_deployments, + partial( + self._create_deployment, + deployment_info=deployment_info, + _model_name=_model_name, + _litellm_params=_litellm_params, + _model_info=_model_info, + declared_id=declared_id, + duplicate_ids=duplicate_ids, + ), + ) verbose_router_logger.exception( "Error creating deployment: %s, ignoring and continuing with other deployments.", e ) @@ -8901,6 +8915,7 @@ class Router: self.quality_routers = {} self.complexity_routers = {} self.auto_routers = {} + self._provider_unresolved_deployments = () self._invalidate_model_group_info_cache() self._invalidate_access_groups_cache() # we add api_base/api_key each model so load balancing between azure/gpt on api_base1 and api_base2 works @@ -9523,8 +9538,12 @@ class Router: """Re-assert this router's deployments onto a freshly fetched catalog. Reads ``model_list`` at call time, so only deployments the router still - serves are restored. + serves are restored, plus any config deployment the fresh catalog now resolves. """ + provider_unresolved: Final = self._provider_unresolved_deployments + self._provider_unresolved_deployments = () + for create_deployment in provider_unresolved: + create_deployment() for entry in tuple(self.model_list): try: deployment = entry if isinstance(entry, Deployment) else Deployment(**entry) diff --git a/tests/test_litellm/litellm_core_utils/test_get_model_cost_map.py b/tests/test_litellm/litellm_core_utils/test_get_model_cost_map.py index 8c0e8ee5d02..a374e03d1c7 100644 --- a/tests/test_litellm/litellm_core_utils/test_get_model_cost_map.py +++ b/tests/test_litellm/litellm_core_utils/test_get_model_cost_map.py @@ -256,14 +256,12 @@ def test_get_model_cost_map_stamps_loaded_at(monkeypatch): from litellm.litellm_core_utils import get_model_cost_map as module monkeypatch.setattr(module._cost_map_source_info, "loaded_at", None) - monkeypatch.setattr( - module.GetModelCostMap, - "fetch_remote_model_cost_map", - staticmethod(lambda url, timeout=5: _load_root_cost_map()), + client, _calls = _mock_client( + [httpx.Response(200, content=_real_map_bytes())], client_cls=httpx.Client ) before = datetime.now(timezone.utc) - module.get_model_cost_map(url="https://example.invalid/cost_map.json") + module.get_model_cost_map(url="https://example.invalid/cost_map.json", client=client) loaded_at = module.get_model_cost_map_loaded_at() assert loaded_at is not None @@ -308,7 +306,7 @@ def _unset_local_cost_map_env(monkeypatch): monkeypatch.delenv("LITELLM_LOCAL_MODEL_COST_MAP", raising=False) -def _mock_client(outcomes): +def _mock_client(outcomes, client_cls=httpx.AsyncClient): """httpx client over a MockTransport serving one outcome per request; an exception instance is raised.""" calls = {"count": 0} @@ -320,7 +318,7 @@ def _mock_client(outcomes): raise outcome return outcome - return httpx.AsyncClient(transport=httpx.MockTransport(handler)), calls + return client_cls(transport=httpx.MockTransport(handler)), calls @pytest.mark.asyncio @@ -450,3 +448,97 @@ async def test_refetch_respects_local_env_override(monkeypatch): ) assert isinstance(result, ModelCostMapReloaded) assert len(result.model_cost_map) > 100 + + +# --------------------------------------------------------------------------- +# get_model_cost_map: the boot-time load retries transient failures like a reload does +# --------------------------------------------------------------------------- + +from litellm.litellm_core_utils.get_model_cost_map import ( + get_model_cost_map, + get_model_cost_map_source_info, +) + + +class _SyncSleepRecorder: + """Injected in place of time.sleep so the boot path's waits are asserted without delay.""" + + def __init__(self): + self.waits = [] + + def __call__(self, seconds: float) -> None: + self.waits.append(seconds) + + +def test_boot_load_retries_transient_failures_instead_of_falling_back(): + """A refused connection then a 503 at pod boot used to pin the process to the bundled + backup for its lifetime; both are transient and must be retried before giving up.""" + client, calls = _mock_client( + [ + httpx.ConnectError("connection refused"), + httpx.Response(503), + httpx.Response(200, content=_real_map_bytes()), + ], + client_cls=httpx.Client, + ) + sleeper = _SyncSleepRecorder() + + cost_map = get_model_cost_map(url=_URL, sleep=sleeper, rng=random.Random(0), client=client) + + assert calls["count"] == 3 + assert len(sleeper.waits) == 2 + assert 2.0 <= sleeper.waits[0] < 3.0 + assert 4.0 <= sleeper.waits[1] < 5.0 + source = get_model_cost_map_source_info() + assert source["source"] == "remote" + assert source["fallback_reason"] is None + assert cost_map.keys() >= _load_root_cost_map().keys() - {"sample_spec", FALLBACK_GENERALIZATIONS_KEY} + + +def test_boot_load_honors_retry_after_then_falls_back_after_max_attempts(): + """An outage longer than the retry budget still ends on the bundled backup, and the + recorded fallback reason says how many attempts were spent so operators can tell.""" + client, calls = _mock_client( + [httpx.Response(429, headers={"Retry-After": "7"})], client_cls=httpx.Client + ) + sleeper = _SyncSleepRecorder() + + cost_map = get_model_cost_map(url=_URL, sleep=sleeper, rng=random.Random(0), client=client) + + assert calls["count"] == 3 + assert sleeper.waits == [7.0, 7.0] + source = get_model_cost_map_source_info() + assert source["source"] == "local" + assert "after 3 attempts" in source["fallback_reason"] + assert len(cost_map) > 100 + + +def test_boot_load_does_not_retry_permanent_failures(): + """A 404 or a malformed URL cannot heal by waiting: one attempt, no sleeps, backup.""" + client, calls = _mock_client([httpx.Response(404)], client_cls=httpx.Client) + sleeper = _SyncSleepRecorder() + + get_model_cost_map(url=_URL, sleep=sleeper, rng=random.Random(0), client=client) + assert calls["count"] == 1 + assert sleeper.waits == [] + assert get_model_cost_map_source_info()["source"] == "local" + + get_model_cost_map(url="not a url", sleep=sleeper, rng=random.Random(0)) + assert sleeper.waits == [] + assert get_model_cost_map_source_info()["source"] == "local" + + +def test_boot_load_respects_local_env_override(monkeypatch): + """LITELLM_LOCAL_MODEL_COST_MAP=True still short-circuits to the backup with zero HTTP.""" + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + + def _fail(request): + raise AssertionError("no HTTP request should be made when local map is forced") + + cost_map = get_model_cost_map( + url=_URL, + sleep=_SyncSleepRecorder(), + client=httpx.Client(transport=httpx.MockTransport(_fail)), + ) + assert len(cost_map) > 100 + assert get_model_cost_map_source_info()["is_env_forced"] is True diff --git a/tests/test_litellm/test_router_model_cost_isolation.py b/tests/test_litellm/test_router_model_cost_isolation.py index 7b7a962bf00..30b265905f3 100644 --- a/tests/test_litellm/test_router_model_cost_isolation.py +++ b/tests/test_litellm/test_router_model_cost_isolation.py @@ -2281,3 +2281,83 @@ def test_every_declaring_deployment_is_named(caplog): assert "azure-ptu-east" in warnings[0] assert "azure-ptu-west" in warnings[0] assert "plain-gpt-4o" not in warnings[0] + + +def _simulate_price_data_reload_with_provider_sets(monkeypatch, fetched_catalog): + """Like `_simulate_price_data_reload`, plus the provider model-set refresh the proxy's + `_swap_in_model_cost_map` does before replaying, so bare names in the new catalog resolve.""" + monkeypatch.setattr(litellm, "model_cost", fetched_catalog) + _invalidate_model_cost_lowercase_map() + litellm.add_known_models(model_cost_map=fetched_catalog) + reapply_runtime_model_cost_registrations() + + +def test_a_config_deployment_dropped_by_a_stale_cost_map_comes_back_on_reload(monkeypatch): + """ + Booting on the bundled backup, a bare model that only the remote catalog knows + cannot be provider-resolved, so the proxy router (ignore_invalid_deployments) drops + it. Once a reload brings in a catalog that knows the model, the deployment must be + served again with its access groups, and exactly once however many reloads follow. + """ + backend = "lit-5766-only-in-remote-catalog" + try: + router = Router( + model_list=[ + { + "model_name": "new-model", + "litellm_params": {"model": backend, "api_key": "k"}, + "model_info": {"id": "new-id", "access_groups": ["team-models"]}, + }, + { + "model_name": "control-model", + "litellm_params": {"model": "hosted_vllm/control-backend", "api_key": "k"}, + "model_info": {"id": "control-id", "access_groups": ["team-models"]}, + }, + ], + ignore_invalid_deployments=True, + ) + assert router.get_model_names() == ["control-model"] + assert router.get_model_access_groups(model_name="new-model") == {} + + fresh_catalog = {**litellm.model_cost, backend: {"litellm_provider": "openai", "mode": "chat"}} + _simulate_price_data_reload_with_provider_sets(monkeypatch, fresh_catalog) + _simulate_price_data_reload_with_provider_sets(monkeypatch, fresh_catalog) + + assert sorted(router.get_model_names()) == ["control-model", "new-model"] + assert router.get_model_access_groups(model_name="new-model") == {"team-models": ["new-model"]} + assert [d["model_info"]["id"] for d in router.model_list] == ["control-id", "new-id"] + assert "new-id" in litellm.model_cost + finally: + litellm.open_ai_chat_completion_models.discard(backend) + litellm.models_by_provider["openai"].discard(backend) + + +def test_a_config_deployment_dropped_for_a_permanent_reason_is_not_retried_on_reload(monkeypatch): + """ + Only provider-resolution drops can be healed by a fresh catalog. A deployment that + fails after its provider resolved (here a pass-through vertex entry with no project) + has already touched router state, so replaying it on every reload would leak into + `deployment_names` each time. + """ + router = Router( + model_list=[ + { + "model_name": "vertex-passthrough", + "litellm_params": {"model": "vertex_ai/gemini-2.5-flash", "use_in_pass_through": True}, + "model_info": {"id": "vertex-id"}, + }, + { + "model_name": "control-model", + "litellm_params": {"model": "hosted_vllm/control-backend", "api_key": "k"}, + "model_info": {"id": "control-id"}, + }, + ], + ignore_invalid_deployments=True, + ) + assert router.get_model_names() == ["control-model"] + names_after_boot = list(router.deployment_names) + + _simulate_price_data_reload_with_provider_sets(monkeypatch, dict(litellm.model_cost)) + + assert router.get_model_names() == ["control-model"] + assert router.deployment_names == names_after_boot From 47b9d838aa05bbc5404c357454ef612dea3ae370 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 1 Sep 2026 17:49:40 -0700 Subject: [PATCH 41/97] perf(scim): resolve group members with one user table read per member (#39228) Co-authored-by: yassin Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../management_endpoints/scim/scim_v2.py | 40 ++++- .../scim/test_scim_v2_endpoints.py | 150 ++++++++++++++---- 2 files changed, 155 insertions(+), 35 deletions(-) diff --git a/litellm/proxy/management_endpoints/scim/scim_v2.py b/litellm/proxy/management_endpoints/scim/scim_v2.py index ded57815e91..8a0436f42dd 100644 --- a/litellm/proxy/management_endpoints/scim/scim_v2.py +++ b/litellm/proxy/management_endpoints/scim/scim_v2.py @@ -585,6 +585,37 @@ async def _users_named_by_member_value( return tuple(dict.fromkeys(row.user_id for row in rows)) +async def _accounts_named_by_member_value(value: str, prisma_client: PrismaClient) -> tuple[str, ...]: + """Every user id this member value names, by user id, SSO identity or email. + + Classification needs to know whether the value is one account's ``user_id`` and + whether it names any other account, so all three fields are read in one pass. The + id is compared exactly and unstripped, as a primary key lookup would; the + identities compare as ``_users_named_by_member_value`` describes. Two rows are + enough to tell one account from several, so the read stops there. Only a full + read that lacks the row keyed by the value leaves that row's existence open, and + only then is the id read on its own. + """ + subject: Final = value.strip() + email: Final[_CaseInsensitiveMatch] = {"equals": subject, "mode": "insensitive"} + users: Final = _table(UserRepository(prisma_client)) + rows: Final = await users.find_many( + where={ # mutable-ok: Prisma filter + "OR": [ # mutable-ok: Prisma filter + {"user_id": value}, # mutable-ok: Prisma filter + {"sso_user_id": subject}, # mutable-ok: Prisma filter + {"user_email": email}, # mutable-ok: Prisma filter + ], + }, + take=2, + ) + named: Final = tuple(dict.fromkeys(row.user_id for row in rows)) + if len(named) < 2 or value in named: + return named + keyed: Final = await users.find_unique(where={"user_id": value}) + return named if keyed is None else (value, *named) + + async def _classify_group_member(member: SCIMMember, prisma_client: PrismaClient) -> _ClassifiedGroupMember: """ Decide what a single SCIM group member refers to. @@ -627,11 +658,9 @@ async def _classify_group_member(member: SCIMMember, prisma_client: PrismaClient if member_type == "group": return _SkippedGroupMember(value=value, reason="nested_group") - user: Final = await _table(UserRepository(prisma_client)).find_unique(where={"user_id": value}) - if user is not None: - shared_with: Final = tuple( - other for other in await _users_named_by_member_value(value, prisma_client) if other != value - ) + named: Final = await _accounts_named_by_member_value(value, prisma_client) + if value in named: + shared_with: Final = tuple(other for other in named if other != value) if shared_with: verbose_proxy_logger.warning( "SCIM: group member '%s' is one account's user id and is also account '%s' by SSO identity or email, " @@ -651,7 +680,6 @@ async def _classify_group_member(member: SCIMMember, prisma_client: PrismaClient if team is not None and _team_metadata_has_scim_provenance(team.metadata): return _SkippedGroupMember(value=value, reason="existing_team") - named: Final = await _users_named_by_member_value(value, prisma_client) if len(named) == 1: verbose_proxy_logger.info( "SCIM: group member '%s' matched user_id '%s' by SSO identity or email", diff --git a/tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py b/tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py index 957f9fde645..d8fe22c4979 100644 --- a/tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py @@ -1,6 +1,6 @@ import logging import time -from collections.abc import Mapping +from collections.abc import Callable, Mapping from itertools import chain from typing import Final from unittest.mock import AsyncMock, MagicMock, call @@ -1645,6 +1645,25 @@ async def test_update_group_e2e(mocker): ScimTransformations.transform_litellm_team_to_scim_group.assert_called_once_with(updated_team) +def _rows_by_exact_id( + user_row: Callable[[Mapping[str, str]], LiteLLM_UserTable | MagicMock | None], +) -> Callable[..., tuple[LiteLLM_UserTable | MagicMock, ...]]: + """``find_many`` stand-in for the classifier's cross-field read on a table where a + member value only ever matches as an exact ``user_id``.""" + + def rows(where: Mapping[str, object], take: int | None = None) -> tuple[LiteLLM_UserTable | MagicMock, ...]: + clauses: Final = where["OR"] + assert isinstance(clauses, list) + found: Final = tuple(user_row(clause) for clause in clauses if "user_id" in clause) + return tuple(row for row in found if row is not None) + + return rows + + +def _user_row_for(where: Mapping[str, str]) -> LiteLLM_UserTable: + return LiteLLM_UserTable(user_id=where["user_id"]) + + @pytest.mark.asyncio async def test_create_group_with_nonexistent_users_rejects(mocker, monkeypatch): """ @@ -1696,9 +1715,8 @@ async def test_create_group_with_nonexistent_users_rejects(mocker, monkeypatch): return mock_user return None # new-user-1 and new-user-2 don't exist - mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(side_effect=mock_user_lookup) mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=None) - mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=[]) + mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(side_effect=_rows_by_exact_id(mock_user_lookup)) # Mock dependencies mocker.patch( @@ -1782,9 +1800,8 @@ async def test_update_group_with_nonexistent_users_rejects(mocker, monkeypatch): return mock_user return None # new-user-3 and new-user-4 don't exist - mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(side_effect=mock_user_lookup) mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=None) - mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=[]) + mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(side_effect=_rows_by_exact_id(mock_user_lookup)) # Mock dependencies mocker.patch( @@ -1853,9 +1870,8 @@ async def test_create_group_with_nonexistent_users_creates_when_flag_true(mocker return mock_user return None # new-user-1 and new-user-2 don't exist - mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(side_effect=mock_user_lookup) mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=None) - mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=[]) + mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(side_effect=_rows_by_exact_id(mock_user_lookup)) # Mock user creation created_user_1 = NewUserResponse(user_id="new-user-1", key="test-key-1") @@ -1943,9 +1959,8 @@ async def test_extract_group_member_ids_with_flag_true_creates_users(mocker, mon return mock_user return None # new-user-1 doesn't exist - mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(side_effect=mock_user_lookup) mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=None) - mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=[]) + mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(side_effect=_rows_by_exact_id(mock_user_lookup)) # Mock user creation created_user = NewUserResponse(user_id="new-user-1", key="test-key-1") @@ -2013,9 +2028,8 @@ async def test_extract_group_member_ids_with_flag_false_rejects(mocker, monkeypa return mock_user return None # new-user-1 doesn't exist - mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(side_effect=mock_user_lookup) mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=None) - mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=[]) + mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(side_effect=_rows_by_exact_id(mock_user_lookup)) # Mock dependencies mocker.patch( @@ -3121,8 +3135,7 @@ async def test_process_group_patch_operations_add_retains_existing_members(mocke mock_prisma_client.db = mocker.MagicMock() mock_prisma_client.db.litellm_usertable = mocker.MagicMock() # new-user already exists in the DB - mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=mocker.MagicMock(user_id="new-user")) - mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=()) + mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=(mocker.MagicMock(user_id="new-user"),)) _, final_members, _ = await _process_group_patch_operations( patch_ops=patch_ops, @@ -3415,8 +3428,7 @@ async def test_patch_group_add_applies_delta_and_keeps_concurrent_add(mocker): ) mock_prisma_client.db.litellm_teamtable.update = AsyncMock(return_value=final_team) mock_prisma_client.db.litellm_usertable = mocker.MagicMock() - mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=mocker.MagicMock()) - mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=()) + mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(side_effect=_rows_by_exact_id(_user_row_for)) mocker.patch( "litellm.proxy.management_endpoints.scim.scim_v2._get_prisma_client_or_raise_exception", @@ -3509,8 +3521,7 @@ async def test_patch_group_replace_stays_absolute_against_concurrent_roster(mock ) mock_prisma_client.db.litellm_teamtable.update = AsyncMock(return_value=final_team) mock_prisma_client.db.litellm_usertable = mocker.MagicMock() - mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=mocker.MagicMock()) - mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=()) + mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(side_effect=_rows_by_exact_id(_user_row_for)) mocker.patch( "litellm.proxy.management_endpoints.scim.scim_v2._get_prisma_client_or_raise_exception", @@ -3640,8 +3651,7 @@ async def test_process_group_patch_add_filtered_path_without_value(mocker): prisma_client = mocker.MagicMock() prisma_client.db = mocker.MagicMock() prisma_client.db.litellm_usertable = mocker.MagicMock() - prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=LiteLLM_UserTable(user_id="user-3")) - prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=()) + prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=(LiteLLM_UserTable(user_id="user-3"),)) _, final_members, _ = await _process_group_patch_operations( patch_ops=patch_ops, @@ -3733,12 +3743,14 @@ def _member_resolution_prisma( starts folding it, fails here instead of passing. A caller that must know which accounts match rather than merely how many - passes take=None, so an unbounded read returns every match. + passes take=None, so an unbounded read returns every match. The row keyed by + the value comes last, the order a bounded read is least prepared for, since + the database promises no order at all. """ clauses: Final = where["OR"] assert isinstance(clauses, list) fields: Final = tuple(next(iter(clause)) for clause in clauses) - assert fields == ("sso_user_id", "user_email"), fields + assert fields in (("user_id", "sso_user_id", "user_email"), ("sso_user_id", "user_email")), fields def comparison(clause: Mapping[str, object]) -> tuple[str, bool]: """The needle and whether production asked for a case-insensitive compare, @@ -3749,8 +3761,9 @@ def _member_resolution_prisma( assert isinstance(criterion, dict), criterion return criterion["equals"], criterion.get("mode") == "insensitive" - sso_needle, sso_insensitive = comparison(clauses[0]) - email_needle, email_insensitive = comparison(clauses[1]) + by_field: Final = dict(zip(fields, (comparison(clause) for clause in clauses))) + sso_needle, sso_insensitive = by_field["sso_user_id"] + email_needle, email_insensitive = by_field["user_email"] def same(stored: str, needle: str, insensitive: bool) -> bool: return stored.casefold() == needle.casefold() if insensitive else stored == needle @@ -3768,6 +3781,11 @@ def _member_resolution_prisma( if same(email, email_needle, email_insensitive) for user_id in user_ids ), + ( + user_id + for user_id in users + if "user_id" in by_field and same(user_id, by_field["user_id"][0], by_field["user_id"][1]) + ), ) ) found: Final = tuple(dict.fromkeys(matched)) @@ -4611,9 +4629,15 @@ async def test_resolve_group_member_ids_dedupes_repeated_member(mocker, scim_ups def _identity_lookup(value: str) -> object: - """The single cross-field lookup the classifier is expected to issue.""" + """The single cross-field lookup the classifier is expected to issue per member.""" return call( - where={"OR": [{"sso_user_id": value}, {"user_email": {"equals": value, "mode": "insensitive"}}]}, + where={ + "OR": [ + {"user_id": value}, + {"sso_user_id": value}, + {"user_email": {"equals": value, "mode": "insensitive"}}, + ] + }, take=2, ) @@ -5152,6 +5176,77 @@ async def test_resolve_group_member_ids_refuses_a_user_id_that_names_another_acc ) +@pytest.mark.asyncio +async def test_resolve_group_member_ids_reads_the_exact_id_when_two_other_accounts_fill_the_lookup( + mocker, scim_upsert_user_enabled +): + """A value that is one account's id and two other accounts' identities fills the + bounded lookup with the other two. The account keyed by the value must still be + found, or the id would lose its precedence and a non-canonical type would skip + a member that names a real user.""" + prisma_client = _member_resolution_prisma( + mocker, + users={"shared"}, + teams=set(), + sso_user_id_to_user_id={"shared": "by-sso"}, + email_to_user_id={"shared": "by-email"}, + ) + create_user_mock = mocker.patch( # test-quality-ok: user creation is module-level, not injectable into the resolver + "litellm.proxy.management_endpoints.scim.scim_v2._create_user_if_not_exists", + AsyncMock(return_value=None), + ) + + with pytest.raises(HTTPException) as exc_info: + await _resolve_group_member_ids( + members=[SCIMMember(value="shared", type="direct")], + created_via="scim_group_membership", + prisma_client=prisma_client, + ) + + assert exc_info.value.status_code == 400 + assert "shared" in str(exc_info.value.detail) + create_user_mock.assert_not_called() + assert prisma_client.db.litellm_usertable.find_many.await_args_list == [_identity_lookup("shared")] + prisma_client.db.litellm_usertable.find_unique.assert_awaited_once_with(where={"user_id": "shared"}) + + +@pytest.mark.asyncio +async def test_resolve_group_member_ids_reads_the_user_table_once_per_member(mocker, scim_upsert_user_enabled): + """Every member costs one read of the user table, however it resolves: by its exact + id (which still outranks a non-canonical type), by identity, as a SCIM team, or not + at all. Looking the exact id up on its own before the identity read doubled the + reads of a push, and the identity read is a scan.""" + prisma_client = _member_resolution_prisma( + mocker, + users={"by-id"}, + teams={"by-team"}, + email_to_user_id={"by-email@example.com": "email-user"}, + ) + mocker.patch( # test-quality-ok: user creation is module-level, not injectable into the resolver + "litellm.proxy.management_endpoints.scim.scim_v2._create_user_if_not_exists", + AsyncMock(return_value=NewUserResponse(user_id="nobody", key="key")), + ) + + result = await _resolve_group_member_ids( + members=[ + SCIMMember(value="by-id", type="direct"), + SCIMMember(value="by-email@example.com"), + SCIMMember(value="by-team"), + SCIMMember(value="nobody"), + ], + created_via="scim_group_membership", + prisma_client=prisma_client, + ) + + assert result.all_member_ids == ["by-id", "email-user", "nobody"] + prisma_client.db.litellm_usertable.find_unique.assert_not_awaited() + assert prisma_client.db.litellm_usertable.find_many.await_args_list == [ + _identity_lookup("by-id"), + _identity_lookup("by-email@example.com"), + _identity_lookup("by-team"), + _identity_lookup("nobody"), + ] + @pytest.mark.asyncio async def test_resolve_group_member_ids_warns_before_creating_unmatched_placeholder( @@ -5536,10 +5631,7 @@ async def test_resolve_group_member_ids_admits_member_created_concurrently(mocke the member is still admitted: the id resolves to a real user row, so failing or dropping it would be wrong either way.""" prisma_client = _member_resolution_prisma(mocker, users=set(), teams=set()) - prisma_client.db.litellm_usertable.find_unique = AsyncMock( - side_effect=[None, LiteLLM_UserTable(user_id="raced-user")] - ) - prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=()) + prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=LiteLLM_UserTable(user_id="raced-user")) mocker.patch( "litellm.proxy.management_endpoints.scim.scim_v2._create_user_if_not_exists", AsyncMock(return_value=None), From 93219a9257ddd74137d7c2476ace01171ce484b9 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 1 Sep 2026 17:51:55 -0700 Subject: [PATCH 42/97] fix(docker): install bedrock-realtime extra in monolith proxy images (#39223) The Dockerfile, docker/Dockerfile.non_root and docker/Dockerfile.database uv sync stages never passed --extra bedrock-realtime, so aws-sdk-bedrock-runtime was absent from the image venv and Bedrock Nova Sonic /v1/realtime sessions failed with 'Missing aws_sdk_bedrock_runtime'. gateway/Dockerfile already had the extra (PR #34426). Adds a static check over every uv sync in the proxy Dockerfiles and an image-level import probe that the image-scan workflow runs against the built root, non-root and gateway images. Co-authored-by: yassin Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .github/workflows/image-scan.yml | 6 +- Dockerfile | 2 + docker/Dockerfile.database | 2 + docker/Dockerfile.non_root | 3 + .../test_image_bedrock_realtime_extra.py | 58 +++++++++++++++++++ .../test_dockerfile_bedrock_realtime_extra.py | 56 ++++++++++++++++++ 6 files changed, 124 insertions(+), 3 deletions(-) create mode 100644 tests/proxy_migration_tests/test_image_bedrock_realtime_extra.py create mode 100644 tests/test_litellm/test_dockerfile_bedrock_realtime_extra.py diff --git a/.github/workflows/image-scan.yml b/.github/workflows/image-scan.yml index bb04563c1a8..206bb809e0c 100644 --- a/.github/workflows/image-scan.yml +++ b/.github/workflows/image-scan.yml @@ -80,7 +80,7 @@ jobs: LITELLM_IMAGE: litellm-image-scan:${{ github.sha }} run: | python -m pip install "pytest==9.0.3" - python -m pytest tests/proxy_migration_tests/test_offline_image_migration.py -v + python -m pytest tests/proxy_migration_tests/test_offline_image_migration.py tests/proxy_migration_tests/test_image_bedrock_realtime_extra.py -v # Scans the whole shipped artifact: OS/apk plus every language package # baked into the image, including ones no lockfile declares (e.g. prisma's @@ -124,7 +124,7 @@ jobs: LITELLM_IMAGE: litellm-runtime-scan:${{ github.sha }} run: | python -m pip install "pytest==9.0.3" - python -m pytest tests/proxy_migration_tests/test_offline_image_migration.py -v + python -m pytest tests/proxy_migration_tests/test_offline_image_migration.py tests/proxy_migration_tests/test_image_bedrock_realtime_extra.py -v migrations-image: name: migrations-image @@ -185,7 +185,7 @@ jobs: LITELLM_COMPONENT_PORT: "4000" run: | python -m pip install "pytest==9.0.3" - python -m pytest tests/proxy_migration_tests/test_component_image_serves_offline.py -v + python -m pytest tests/proxy_migration_tests/test_component_image_serves_offline.py tests/proxy_migration_tests/test_image_bedrock_realtime_extra.py -v ui-image: name: ui-image diff --git a/Dockerfile b/Dockerfile index 29a085a4ef9..0a92aa9a68c 100644 --- a/Dockerfile +++ b/Dockerfile @@ -66,6 +66,7 @@ RUN uv sync --frozen --no-install-project --no-install-workspace --no-default-gr --extra extra_proxy \ --extra semantic-router \ --extra saml \ + --extra bedrock-realtime \ --python python3.13 # Copy full source tree @@ -87,6 +88,7 @@ RUN uv sync --frozen --no-default-groups --no-editable \ --extra extra_proxy \ --extra semantic-router \ --extra saml \ + --extra bedrock-realtime \ --python python3.13 RUN HOME=/opt/prisma XDG_CACHE_HOME=/opt/prisma/.cache PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries \ diff --git a/docker/Dockerfile.database b/docker/Dockerfile.database index c1348f68231..e9ad2849bb2 100644 --- a/docker/Dockerfile.database +++ b/docker/Dockerfile.database @@ -64,6 +64,7 @@ RUN uv sync --frozen --no-install-project --no-install-workspace --no-default-gr --extra extra_proxy \ --extra semantic-router \ --extra saml \ + --extra bedrock-realtime \ --python python3.13 # Copy full source tree @@ -85,6 +86,7 @@ RUN uv sync --frozen --no-default-groups --no-editable \ --extra extra_proxy \ --extra semantic-router \ --extra saml \ + --extra bedrock-realtime \ --python python3.13 RUN HOME=/opt/prisma XDG_CACHE_HOME=/opt/prisma/.cache PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries \ diff --git a/docker/Dockerfile.non_root b/docker/Dockerfile.non_root index 2221435a83a..edf20e8bbff 100644 --- a/docker/Dockerfile.non_root +++ b/docker/Dockerfile.non_root @@ -70,6 +70,7 @@ RUN --mount=type=cache,target=/app/.cache/uv,id=litellm-uv-cache \ --extra extra_proxy \ --extra semantic-router \ --extra saml \ + --extra bedrock-realtime \ --python python3.13 # Copy full source tree @@ -97,6 +98,7 @@ RUN --mount=type=cache,target=/app/.cache/uv,id=litellm-uv-cache \ --extra extra_proxy \ --extra semantic-router \ --extra saml \ + --extra bedrock-realtime \ --python python3.13 \ --no-sources-package litellm-proxy-extras; \ else \ @@ -106,6 +108,7 @@ RUN --mount=type=cache,target=/app/.cache/uv,id=litellm-uv-cache \ --extra extra_proxy \ --extra semantic-router \ --extra saml \ + --extra bedrock-realtime \ --python python3.13; \ fi diff --git a/tests/proxy_migration_tests/test_image_bedrock_realtime_extra.py b/tests/proxy_migration_tests/test_image_bedrock_realtime_extra.py new file mode 100644 index 00000000000..ed21734c5fc --- /dev/null +++ b/tests/proxy_migration_tests/test_image_bedrock_realtime_extra.py @@ -0,0 +1,58 @@ +"""Image-level check that the built proxy image can import the Bedrock realtime SDK. + +Bedrock Nova Sonic (`/v1/realtime`) imports `aws_sdk_bedrock_runtime` lazily on the +first session, so an image whose `uv sync` stages skip the `bedrock-realtime` extra +boots, passes health checks, and then fails every Nova Sonic session with +"Missing aws_sdk_bedrock_runtime". Importing inside the built image is what catches +that class of regression (missing extra, lockfile drift, a stage that syncs a +different set of extras), which a static Dockerfile check cannot. + +Gated on LITELLM_IMAGE like the other image checks in this directory; exercised +where an image has been built (the image-scan workflow). Requires a working docker CLI. +""" + +import os +import shutil +import subprocess +from typing import Final + +import pytest + +IMAGE: Final = os.getenv("LITELLM_IMAGE") +NON_ROOT_UID: Final = "12345:0" +IMPORT_PROBE: Final = "import aws_sdk_bedrock_runtime, smithy_aws_core; print('bedrock-realtime ok')" + +pytestmark = [ + pytest.mark.skipif(IMAGE is None, reason="requires a built image (set LITELLM_IMAGE)"), + pytest.mark.skipif(shutil.which("docker") is None, reason="requires the docker CLI"), +] + + +def test_image_imports_bedrock_realtime_sdk(): + assert IMAGE is not None + + probe: Final = subprocess.run( + [ + "docker", + "run", + "--rm", + "--network", + "none", + "--user", + NON_ROOT_UID, + "--entrypoint", + "python", + IMAGE, + "-c", + IMPORT_PROBE, + ], + capture_output=True, + text=True, + check=False, + ) + + assert probe.returncode == 0 and "bedrock-realtime ok" in probe.stdout, ( + f"{IMAGE} cannot import aws_sdk_bedrock_runtime as uid {NON_ROOT_UID}, so Bedrock Nova Sonic " + "/v1/realtime sessions fail with 'Missing aws_sdk_bedrock_runtime'. Is `--extra bedrock-realtime` " + f"passed to every `uv sync` in its Dockerfile?\nstdout:\n{probe.stdout}\nstderr:\n{probe.stderr}" + ) diff --git a/tests/test_litellm/test_dockerfile_bedrock_realtime_extra.py b/tests/test_litellm/test_dockerfile_bedrock_realtime_extra.py new file mode 100644 index 00000000000..44572aed08e --- /dev/null +++ b/tests/test_litellm/test_dockerfile_bedrock_realtime_extra.py @@ -0,0 +1,56 @@ +""" +Static checks that every proxy Docker image installs the `bedrock-realtime` extra. + +Bedrock Nova Sonic speech-to-speech (`/v1/realtime`) needs `aws-sdk-bedrock-runtime`, +which only ships in the `bedrock-realtime` extra. An image whose `uv sync` stages +omit the extra fails every Nova Sonic realtime session with +"Missing aws_sdk_bedrock_runtime. Install with: pip install aws-sdk-bedrock-runtime". +""" + +import os +import re +from typing import Final + +import pytest + +REPO_ROOT: Final = os.path.join(os.path.dirname(__file__), "..", "..") + +PROXY_DOCKERFILES: Final = ( + "Dockerfile", + os.path.join("docker", "Dockerfile.non_root"), + os.path.join("docker", "Dockerfile.database"), + os.path.join("gateway", "Dockerfile"), +) + +CONTINUED_LINE_RE: Final = re.compile(r"(?:\\\n|[^\n])+") +UV_SYNC_BOUNDARY_RE: Final = re.compile(r"(?=uv sync)") + + +def _uv_sync_invocations(dockerfile_text: str) -> tuple[str, ...]: + """Return each `uv sync ...` command, split apart when one RUN holds several (if/else branches).""" + return tuple( + part + for line in CONTINUED_LINE_RE.finditer(dockerfile_text) + for part in UV_SYNC_BOUNDARY_RE.split(line.group(0)) + if part.startswith("uv sync") + ) + + +@pytest.mark.parametrize("relative_path", PROXY_DOCKERFILES) +def test_every_uv_sync_installs_bedrock_realtime_extra(relative_path: str): + dockerfile_path: Final = os.path.join(REPO_ROOT, relative_path) + if not os.path.exists(dockerfile_path): + pytest.skip(f"{relative_path} not present in this checkout") + + with open(dockerfile_path, "r", encoding="utf-8") as f: + contents: Final = f.read() + + invocations: Final = _uv_sync_invocations(contents) + assert invocations, f"{relative_path} has no `uv sync` invocation" + + missing: Final = tuple(invocation for invocation in invocations if "--extra bedrock-realtime" not in invocation) + assert not missing, ( + f"{relative_path}: {len(missing)} of {len(invocations)} `uv sync` invocations omit " + "`--extra bedrock-realtime`, so aws-sdk-bedrock-runtime is absent and Bedrock Nova Sonic " + "/v1/realtime sessions fail with 'Missing aws_sdk_bedrock_runtime'" + ) From c0019751520837bd3a0a6203297292bab31171de Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 1 Sep 2026 17:55:47 -0700 Subject: [PATCH 43/97] fix(aiohttp_transport): map transport-internal CancelledError to a retryable ConnectError (#39240) aiohttp shields its DNS resolution task; when the connector closes it cancels that child, so the request task sees CancelledError without ever being cancelled itself. map_aiohttp_exceptions() only caught Exception, so the BaseException skipped transport mapping, router retries and proxy error handling, and /v1/responses answered 500 "No response returned". Catch CancelledError in the mapper, re-raise when the current task is really being cancelled (Task.cancelling() > 0), and otherwise map it to httpx.ConnectError so the usual retry, fallback and error mapping apply. Co-authored-by: yassin Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../llms/custom_httpx/aiohttp_transport.py | 13 +++++ .../custom_httpx/test_aiohttp_transport.py | 56 +++++++++++++++++++ 2 files changed, 69 insertions(+) diff --git a/litellm/llms/custom_httpx/aiohttp_transport.py b/litellm/llms/custom_httpx/aiohttp_transport.py index b6586481fd3..73adf9c7455 100644 --- a/litellm/llms/custom_httpx/aiohttp_transport.py +++ b/litellm/llms/custom_httpx/aiohttp_transport.py @@ -3,6 +3,7 @@ import concurrent.futures import contextlib import os import ssl +import sys import typing import urllib.request from collections.abc import Callable, Generator @@ -75,10 +76,22 @@ except ImportError: pass +def _current_task_is_cancelling() -> bool: + task: Final = asyncio.current_task() + if task is None or sys.version_info < (3, 11): + return True + return task.cancelling() > 0 + + @contextlib.contextmanager def map_aiohttp_exceptions() -> Generator[None, None, None]: try: yield + except asyncio.CancelledError as exc: + # a closing connector cancels its shielded DNS task; that surfaces here without the request task being cancelled + if _current_task_is_cancelling(): + raise + raise httpx.ConnectError("aiohttp transport cancelled the request internally") from exc except Exception as exc: mapped_exc: type[Exception] | None = None diff --git a/tests/test_litellm/llms/custom_httpx/test_aiohttp_transport.py b/tests/test_litellm/llms/custom_httpx/test_aiohttp_transport.py index 4c92c52d556..7509e35e3f7 100644 --- a/tests/test_litellm/llms/custom_httpx/test_aiohttp_transport.py +++ b/tests/test_litellm/llms/custom_httpx/test_aiohttp_transport.py @@ -1,7 +1,11 @@ import asyncio import concurrent.futures +import socket +import sys +from typing import Final import aiohttp +import aiohttp.abc import aiohttp.client_exceptions import aiohttp.http_exceptions import httpx @@ -1140,3 +1144,55 @@ async def test_stopped_loop_session_disposed_synchronously_on_recycle(): finally: await new_session.close() result["loop"].close() + + +class _CancellingResolver(aiohttp.abc.AbstractResolver): + """Cancels the given task (or, by default, aiohttp's shielded DNS child task) mid-lookup.""" + + def __init__(self, task_to_cancel: "asyncio.Task[object] | None" = None): + self._task_to_cancel: Final = task_to_cancel + + async def resolve( + self, host: str, port: int = 0, family: socket.AddressFamily = socket.AF_INET + ) -> list[aiohttp.abc.ResolveResult]: + target: Final = self._task_to_cancel or asyncio.current_task() + assert target is not None + target.cancel() + await asyncio.sleep(0) + raise OSError("resolver finished after the task was cancelled") + + async def close(self) -> None: + return None + + +@pytest.mark.asyncio +@pytest.mark.skipif( + sys.version_info < (3, 11), reason="Task.cancelling() is needed to tell the two cancellations apart" +) +async def test_internal_dns_cancellation_maps_to_connect_error(): + """A CancelledError the request task never asked for must surface as a mapped httpx transport error.""" + session = aiohttp.ClientSession(connector=aiohttp.TCPConnector(resolver=_CancellingResolver())) + transport = LiteLLMAiohttpTransport(client=session) + try: + with pytest.raises(httpx.ConnectError): + await transport.handle_async_request(httpx.Request("GET", "http://example.invalid/")) + current = asyncio.current_task() + assert current is not None and current.cancelling() == 0 + finally: + await transport.aclose() + + +@pytest.mark.asyncio +async def test_genuine_request_cancellation_still_propagates(): + """Cancelling the request task itself (client disconnect, shutdown) must still propagate unmapped.""" + current = asyncio.current_task() + assert current is not None + session = aiohttp.ClientSession(connector=aiohttp.TCPConnector(resolver=_CancellingResolver(current))) + transport = LiteLLMAiohttpTransport(client=session) + try: + with pytest.raises(asyncio.CancelledError): + await transport.handle_async_request(httpx.Request("GET", "http://example.invalid/")) + finally: + if sys.version_info >= (3, 11): + current.uncancel() + await transport.aclose() From 48dd06e841074e6f3362c85bd4fd679c34530f97 Mon Sep 17 00:00:00 2001 From: tin-berri Date: Tue, 1 Sep 2026 18:00:31 -0700 Subject: [PATCH 44/97] fix(bedrock): gate Converse cachePoint emission on model prompt caching support (#39210) Bedrock rejects requests carrying cachePoint blocks for models whose entry in the cost map does not declare supports_prompt_caching (403 "You invoked an unsupported model or your request did not allow prompt caching"). Clients like Claude Code attach cache_control to every request, so any such model behind the gateway failed on every call. The new bedrock_model_accepts_cache_points predicate drops cachePoint emission for map-known non-caching models at all three emission funnels, keeps emitting for unmapped ids (application inference profile ARNs), and skips the gateway injection credit when the tool_config point is not placed. --- .../prompt_templates/factory.py | 7 +- .../bedrock/chat/converse_transformation.py | 5 +- litellm/llms/bedrock/common_utils.py | 24 ++++++ ...llm_core_utils_prompt_templates_factory.py | 22 +++++ .../chat/test_converse_transformation.py | 80 ++++++++++++++++++- 5 files changed, 133 insertions(+), 5 deletions(-) diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index e6402e8c1bd..ba59e3fa997 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -4957,10 +4957,13 @@ def make_valid_bedrock_tool_name(input_tool_name: str) -> str: def add_cache_point_tool_block(tool: dict, model: str | None = None) -> BedrockToolBlock | None: - from litellm.llms.bedrock.common_utils import is_claude_4_5_on_bedrock + from litellm.llms.bedrock.common_utils import ( + bedrock_model_accepts_cache_points, + is_claude_4_5_on_bedrock, + ) cache_control: Final = tool.get("cache_control", None) - if cache_control is not None: + if cache_control is not None and bedrock_model_accepts_cache_points(model): cache_point: Final = cache_control.get("type", "ephemeral") if cache_point == "ephemeral": cache_point_block: Final[CachePointBlock] = {"type": "default"} diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index 7fefeaeaf04..5363c3c0366 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -87,6 +87,7 @@ from ..common_utils import ( BedrockError, BedrockModelInfo, bedrock_converse_supports_parallel_tool_use_config, + bedrock_model_accepts_cache_points, get_anthropic_beta_from_headers, get_bedrock_tool_name, is_bedrock_application_inference_profile_arn, @@ -1149,7 +1150,7 @@ class AmazonConverseConfig(BaseConfig): model: str | None = None, ) -> SystemContentBlock | ContentBlock | None: cache_control: Final = message_block.get("cache_control", None) - if cache_control is None: + if cache_control is None or not bedrock_model_accepts_cache_points(model): return None cache_point: Final = self._build_cache_point_block(cache_control, model) @@ -1613,7 +1614,7 @@ class AmazonConverseConfig(BaseConfig): # Append cachePoint to tools if cache_control_injection_points has tool_config cache_injection_points: Final = additional_request_params.pop("cache_control_injection_points", None) - if cache_injection_points and len(bedrock_tools) > 0: + if cache_injection_points and len(bedrock_tools) > 0 and bedrock_model_accepts_cache_points(model): for point in cache_injection_points: if point.get("location") == "tool_config": cache_point = self._build_cache_point_block(point.get("control"), model) diff --git a/litellm/llms/bedrock/common_utils.py b/litellm/llms/bedrock/common_utils.py index 30a77d57f24..66ee5f10679 100644 --- a/litellm/llms/bedrock/common_utils.py +++ b/litellm/llms/bedrock/common_utils.py @@ -816,6 +816,30 @@ def bedrock_converse_supports_parallel_tool_use_config(model: str) -> bool: ) +def bedrock_model_accepts_cache_points(model: str | None) -> bool: + """ + Whether Converse ``cachePoint`` blocks may be sent to this model. + + Bedrock rejects requests carrying cachePoint blocks for models without prompt + caching support ("You invoked an unsupported model or your request did not allow + prompt caching"), so a model whose cost-map entry does not declare + ``supports_prompt_caching`` must not receive them. A model absent from the map + (an application inference profile ARN, a model newer than the map) keeps emitting + so existing caching setups never silently degrade. ``litellm.utils.supports_prompt_caching`` + is not reusable here: it returns False for unmapped models, the opposite polarity. + """ + if model is None: + return True + entries: Final = tuple( + entry + for candidate in (model, get_bedrock_base_model(model)) + if (entry := litellm.model_cost.get(candidate)) is not None + ) + if not entries: + return True + return any(entry.get("supports_prompt_caching") is True for entry in entries) + + def is_claude_4_5_on_bedrock(model: str) -> bool: """ Check if the model supports Bedrock prompt caching with an extended '1h' TTL diff --git a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py index 64c96b575c5..dd2d45f00c6 100644 --- a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py +++ b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py @@ -2932,6 +2932,28 @@ def test_add_cache_point_tool_block_passes_ttl_for_claude_4_5(monkeypatch): monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", old_env) +def test_add_cache_point_tool_block_stands_down_for_model_without_prompt_caching(monkeypatch): + """A tool carrying cache_control must not become a cachePoint for a Bedrock model + whose cost-map entry lacks prompt caching support, since Bedrock rejects the whole + request. An unmapped id keeps emitting so ARN deployments do not lose caching.""" + from litellm.litellm_core_utils.prompt_templates.factory import ( + add_cache_point_tool_block, + ) + + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) + tool = {"cache_control": {"type": "ephemeral"}} + + assert add_cache_point_tool_block(tool, model="nvidia.nemotron-super-3-120b") is None + assert add_cache_point_tool_block(tool, model="us.nvidia.nemotron-super-3-120b") is None + assert add_cache_point_tool_block( + tool, model="arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/abc123" + ) == {"cachePoint": {"type": "default"}} + assert add_cache_point_tool_block(tool, model="us.anthropic.claude-sonnet-4-5-20250929-v1:0") == { + "cachePoint": {"type": "default"} + } + + def test_bedrock_tools_pt_passes_ttl_for_claude_4_5(monkeypatch): """ End-to-end: _bedrock_tools_pt should produce cachePoint blocks with ttl diff --git a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py index 4f53d3481de..70f3153ed7e 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py @@ -5248,6 +5248,84 @@ def test_cache_control_injection_tool_config_drops_ttl_for_unsupported_model(): assert tools[-1] == {"cachePoint": {"type": "default"}} +@pytest.mark.parametrize( + ("model", "expects_cache_points"), + [ + pytest.param("nvidia.nemotron-super-3-120b", False, id="mapped-model-without-prompt-caching"), + pytest.param("us.nvidia.nemotron-super-3-120b", False, id="regional-prefix-resolves-through-base-model"), + pytest.param( + "us.anthropic.claude-3-5-sonnet-20240620-v1:0", False, id="claude-named-but-not-caching-on-bedrock" + ), + pytest.param("us.anthropic.claude-sonnet-4-5-20250929-v1:0", True, id="mapped-model-with-prompt-caching"), + pytest.param( + "arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/abc123", + True, + id="unmapped-arn-keeps-emitting", + ), + ], +) +def test_cache_points_emitted_only_for_models_that_support_prompt_caching(model, expects_cache_points, monkeypatch): + """Bedrock rejects cachePoint blocks for models without prompt caching support + ("You invoked an unsupported model or your request did not allow prompt caching"), + and clients like Claude Code attach cache_control to every request, so a map-known + model without the capability must not receive them. Unmapped ids (application + inference profile ARNs, models newer than the map) keep emitting so existing + caching setups never silently degrade.""" + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) + + body = AmazonConverseConfig().transform_request( + model=model, + messages=[ + {"role": "system", "content": [{"type": "text", "text": "sys", "cache_control": {"type": "ephemeral"}}]}, + {"role": "user", "content": [{"type": "text", "text": "hi", "cache_control": {"type": "ephemeral"}}]}, + ], + optional_params={}, + litellm_params={}, + headers={}, + ) + + assert ("cachePoint" in json.dumps(body)) is expects_cache_points + assert body["system"][0]["text"] == "sys" + assert body["messages"][0]["content"][0]["text"] == "hi" + + +def test_tool_config_cachepoint_not_placed_or_credited_for_model_without_prompt_caching(monkeypatch): + """The tool_config injection point must stand down with the rest of the cachePoint + emission when the model cannot cache, and spend attribution must not credit the + gateway for a breakpoint that was never placed.""" + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) + + bucket: dict = {"user_api_key": "sk-test"} + data = AmazonConverseConfig()._transform_request_helper( + model="nvidia.nemotron-super-3-120b", + system_content_blocks=[], + optional_params={ + "tools": [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get weather", + "parameters": { + "type": "object", + "properties": {"location": {"type": "string"}}, + "required": ["location"], + }, + }, + } + ], + "cache_control_injection_points": [{"location": "tool_config"}], + }, + messages=[{"role": "user", "content": "hi"}], + litellm_params={"metadata": bucket, "litellm_metadata": None, "model_info": {"id": "dep-bedrock"}}, + ) + + assert "cachePoint" not in json.dumps(data.get("toolConfig", {})) + assert "litellm_gateway_injected_cache" not in bucket + + def test_translate_response_format_json_schema_still_injects_tool(): """ response_format with an explicit json_schema should still use the @@ -6211,7 +6289,7 @@ def test_message_level_cache_control_drops_ttl_for_unsupported_model(ttl_target) result = _bedrock_converse_messages_pt( messages=_agentic_messages_with_ttl(ttl_target), - model="anthropic.claude-3-5-sonnet-20240620-v1:0", + model="anthropic.claude-3-5-sonnet-20241022-v2:0", llm_provider="bedrock_converse", ) From 81277252e1e3a9f8ac7c8c73a6c72151c5b596a4 Mon Sep 17 00:00:00 2001 From: tin-berri Date: Tue, 1 Sep 2026 18:01:13 -0700 Subject: [PATCH 45/97] fix(datadog_llm_obs): send tool calls, tool results and cache tokens in DD's own fields (#39222) The LLM Obs callback copied litellm's OpenAI-shaped objects into the span verbatim, so every field Datadog names differently landed somewhere it does not read: tool calls kept their nested `function` wrapper instead of DD's name/arguments/tool_id, tool messages carried no result linking them to their call, the request's tools were never sent, and prompt-cache counts sat inside meta.metadata rather than the span metrics its cache dashboards chart. One rule governs the message mapper: add the fields Datadog declares, and never destroy content it did not understand. Content collapses to its text only when it has text, so a content list carrying tool or image blocks rides along unchanged, and absent messages map to an empty input rather than a fabricated turn. Tool calls and results are read from both dialects, the OpenAI `tool_calls` / `role: tool` shape and the Anthropic `tool_use` / `tool_result` content blocks, so /v1/messages sessions gain tool linking they never had. Cache counts come from the same owners the savings dashboard uses, so every provider spelling resolves through one place rather than a second local guess. The three cache metrics partition the input count: litellm's normalized prompt total includes both cache categories, as the cost calculator's pricing helper documents, so the non-cached residual subtracts reads AND writes. Counting a primed prefix as ordinary input had inflated non-cached usage by exactly the cache-write count on every priming request. Correlating a result to its call reads ids and names structurally and parses no arguments, so a tool call's arguments are decoded once per span rather than once per pass, and arguments past a size bound ship as the raw string instead of paying a decode that multiplies memory on hostile compact JSON. The flat `output_tool_calls.*` metadata copies go away with this: they were a second representation of a fact that now has its own field on the same span. --- .../integrations/datadog/datadog_llm_obs.py | 393 +++++++++------ litellm/types/integrations/datadog_llm_obs.py | 49 +- .../datadog/test_datadog_llm_obs.py | 469 ++++++++++++++++++ 3 files changed, 762 insertions(+), 149 deletions(-) create mode 100644 tests/test_litellm/integrations/datadog/test_datadog_llm_obs.py diff --git a/litellm/integrations/datadog/datadog_llm_obs.py b/litellm/integrations/datadog/datadog_llm_obs.py index e5789965c6e..5e116b7301a 100644 --- a/litellm/integrations/datadog/datadog_llm_obs.py +++ b/litellm/integrations/datadog/datadog_llm_obs.py @@ -11,6 +11,7 @@ import json import os from collections.abc import Mapping, Sequence from datetime import datetime +from types import MappingProxyType from typing import Any, Final, Literal import httpx @@ -30,12 +31,16 @@ from litellm.integrations.datadog.datadog_mock_client import ( ) from litellm.litellm_core_utils.dd_tracing import tracer from litellm.litellm_core_utils.prompt_templates.common_utils import ( + convert_content_list_to_str, handle_any_messages_to_chat_completion_str_messages_conversion, ) +from litellm.litellm_core_utils.safe_json_dumps import safe_dumps +from litellm.litellm_core_utils.safe_json_loads import safe_json_loads from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, httpxSpecialProvider, ) +from litellm.proxy.spend_tracking.savings import extract_cache_creation_tokens, extract_cache_read_tokens from litellm.types.integrations.datadog_llm_obs import * from litellm.types.utils import ( CallTypes, @@ -44,6 +49,189 @@ from litellm.types.utils import ( StandardLoggingPayloadErrorInformation, ) +_EMPTY_MAPPING: Final[Mapping[str, Any]] = MappingProxyType({}) +_EMPTY_MESSAGE: Final[Message] = {"role": "", "content": ""} +_MAX_PARSED_TOOL_ARGUMENT_CHARS: Final = 256 * 1024 + + +def _mapping_field(source: Mapping[str, Any], key: str) -> Mapping[str, Any]: + """The value at `key` when it is a mapping, else an empty one.""" + value: Final = source.get(key) + return value if isinstance(value, dict) else _EMPTY_MAPPING + + +def _content_blocks(message: Mapping[str, Any]) -> tuple[Mapping[str, Any], ...]: + content: Final = message.get("content") + if not isinstance(content, list): + return () + return tuple(block for block in content if isinstance(block, dict)) + + +def _to_dd_arguments(raw_arguments: object) -> dict[str, Any] | str: + """ + Arguments as the object LLM Obs types them as, or the raw string when they are not one. + + Strings past the size bound ship unparsed: decoding multiplies memory on hostile compact + JSON, and the raw string is what the intake receives either way. + """ + if not isinstance(raw_arguments, str): + return raw_arguments if isinstance(raw_arguments, dict) else str(raw_arguments) + if len(raw_arguments) > _MAX_PARSED_TOOL_ARGUMENT_CHARS: + return raw_arguments + parsed: Final = safe_json_loads(raw_arguments) + return parsed if isinstance(parsed, dict) else raw_arguments + + +def _to_dd_tool_calls(message: Mapping[str, Any]) -> tuple[ToolCall, ...]: + """ + The tool calls a message carries, in LLM Obs' ToolCall schema, from either dialect. + + OpenAI puts them in `tool_calls` with the callee nested under `function` and `arguments` + serialized; Anthropic puts them in `content` as `tool_use` blocks with `input` already an + object. LLM Obs reads `name` / `arguments` / `tool_id` either way. + """ + raw_tool_calls: Final = message.get("tool_calls") + openai_calls: Final = tuple( + ToolCall( + name=function.get("name", ""), + arguments=_to_dd_arguments(function.get("arguments", "")), + tool_id=tool_call.get("id", ""), + type=tool_call.get("type", "function"), + ) + for tool_call in (raw_tool_calls if isinstance(raw_tool_calls, list) else ()) + if isinstance(tool_call, dict) + for function in [_mapping_field(tool_call, "function")] + ) + anthropic_calls: Final = tuple( + ToolCall( + name=block.get("name", ""), + arguments=_to_dd_arguments(block.get("input") or {}), + tool_id=block.get("id", ""), + type="tool_use", + ) + for block in _content_blocks(message) + if block.get("type") == "tool_use" + ) + return openai_calls + anthropic_calls + + +def _to_dd_tool_results(message: Mapping[str, Any], tool_call_names: Mapping[str, str]) -> tuple[ToolResult, ...]: + """ + The tool results a message carries, linked back to the call each answers. + + OpenAI models a result as a whole `role: "tool"` message keyed by `tool_call_id`; + Anthropic nests `tool_result` blocks inside a user message, keyed by `tool_use_id`. + """ + + def to_result(tool_id: str, result: object) -> ToolResult: + return ToolResult( + name=tool_call_names.get(tool_id, ""), + result=result if isinstance(result, str) else safe_dumps(result), + tool_id=tool_id, + type="function", + ) + + if message.get("role") == "tool": + return (to_result(str(message.get("tool_call_id", "")), message.get("content") or ""),) + return tuple( + to_result(str(block.get("tool_use_id", "")), block.get("content") or "") + for block in _content_blocks(message) + if block.get("type") == "tool_result" + ) + + +def _tool_call_names_by_id(messages: Sequence[object]) -> Mapping[str, str]: + """Ids to tool names for result linking; reads names structurally and parses nothing.""" + openai_pairs: Final = tuple( + (tool_call.get("id"), function.get("name", "")) + for message in messages + if isinstance(message, dict) and isinstance(message.get("tool_calls"), list) + for tool_call in message["tool_calls"] + if isinstance(tool_call, dict) + for function in [_mapping_field(tool_call, "function")] + ) + anthropic_pairs: Final = tuple( + (block.get("id"), block.get("name", "")) + for message in messages + if isinstance(message, dict) + for block in _content_blocks(message) + if block.get("type") == "tool_use" + ) + return MappingProxyType({str(tool_id): str(name) for tool_id, name in openai_pairs + anthropic_pairs if tool_id}) + + +def _to_dd_message(message: object, tool_call_names: Mapping[str, str]) -> Message: + """ + Map one chat message onto LLM Obs' Message schema, adding fields and never destroying content. + + Content collapses to its text only when it has text; a content list with none (tool blocks, + images) rides along unchanged so nothing the caller logged is lost. Tool calls and results + move into the fields the LLM Obs Tools panel reads, from both the OpenAI and Anthropic shapes. + """ + if not isinstance(message, dict): + converted: Final = handle_any_messages_to_chat_completion_str_messages_conversion(message) + return converted[0] if converted else _EMPTY_MESSAGE + + text: Final = convert_content_list_to_str(message) # pyright: ignore[reportArgumentType] # caller-supplied dict + original_content: Final = message.get("content") + content: Final = ( + text if text or not isinstance(original_content, list) or not original_content else original_content + ) + reasoning: Final = message.get("reasoning_content") + tool_calls: Final = _to_dd_tool_calls(message) + tool_results: Final = _to_dd_tool_results(message, tool_call_names) + dd_message: Final[Message] = { + "role": message.get("role", ""), + "content": content, + **({"reasoning_content": reasoning} if reasoning is not None else {}), + **({"tool_calls": tool_calls} if tool_calls else {}), + **({"tool_results": tool_results} if tool_results else {}), + } + return dd_message + + +def _to_dd_messages(messages: object) -> tuple[Message, ...]: + """Map a whole conversation, resolving each tool result against the calls that precede it.""" + if messages is None: + return () + if not isinstance(messages, list): + return tuple(handle_any_messages_to_chat_completion_str_messages_conversion(messages)) + tool_call_names: Final = _tool_call_names_by_id(messages) + return tuple(_to_dd_message(message, tool_call_names) for message in messages) + + +def _to_dd_tool_definition(entry: Mapping[str, Any]) -> ToolDefinition | None: + function: Final = entry.get("function") + declared: Final[Mapping[str, Any]] = function if isinstance(function, dict) else entry + name: Final = declared.get("name") + if not name: + return None + schema: Final = declared.get("parameters") or declared.get("input_schema") + description: Final = declared.get("description", "") + if not isinstance(schema, dict): + return ToolDefinition(name=name, description=description) + return ToolDefinition(name=name, description=description, schema=schema) + + +def _to_dd_tool_definitions(model_parameters: object) -> tuple[ToolDefinition, ...]: + """ + Map the request's declared tools onto LLM Obs' ToolDefinition schema. + + Handles the wrapped chat-completions shape and the bare shape the Anthropic and + Responses surfaces use, since both reach this logger through `model_parameters`. + """ + if not isinstance(model_parameters, dict): + return () + raw_tools: Final = model_parameters.get("tools") or model_parameters.get("functions") + if not isinstance(raw_tools, list): + return () + return tuple( + definition + for entry in raw_tools + if isinstance(entry, dict) + if (definition := _to_dd_tool_definition(entry)) is not None + ) + class DataDogLLMObsLogger(CustomBatchLogger): def __init__(self, **kwargs): @@ -222,12 +410,9 @@ class DataDogLLMObsLogger(CustomBatchLogger): if standard_logging_payload is None: raise Exception("DataDogLLMObs: standard_logging_object is not set") - messages = standard_logging_payload["messages"] - messages = self._ensure_string_content(messages=messages) - metadata: Final = kwargs.get("litellm_params", {}).get("metadata", {}) - input_meta: Final = InputMeta(messages=handle_any_messages_to_chat_completion_str_messages_conversion(messages)) + input_meta: Final = InputMeta(messages=_to_dd_messages(standard_logging_payload["messages"])) output_meta: Final = OutputMeta( messages=self._get_response_messages( standard_logging_payload=standard_logging_payload, @@ -241,22 +426,20 @@ class DataDogLLMObsLogger(CustomBatchLogger): if isinstance(metadata, dict): metadata_parent_id = metadata.get("parent_id") - meta: Final = Meta( - kind=self._get_datadog_span_kind(standard_logging_payload.get("call_type"), metadata_parent_id), - input=input_meta, - output=output_meta, - metadata=self._get_dd_llm_obs_payload_metadata(standard_logging_payload), - error=error_info, - ) + tool_definitions: Final = _to_dd_tool_definitions(standard_logging_payload.get("model_parameters")) + span_kind: Final = self._get_datadog_span_kind(standard_logging_payload.get("call_type"), metadata_parent_id) + payload_metadata: Final = self._get_dd_llm_obs_payload_metadata(standard_logging_payload) - # Calculate metrics (you may need to adjust these based on available data) - metrics: Final = LLMMetrics( - input_tokens=float(standard_logging_payload.get("prompt_tokens", 0)), - output_tokens=float(standard_logging_payload.get("completion_tokens", 0)), - total_tokens=float(standard_logging_payload.get("total_tokens", 0)), - total_cost=float(standard_logging_payload.get("response_cost", 0)), - time_to_first_token=self._get_time_to_first_token_seconds(standard_logging_payload), - ) + meta: Final[Meta] = { + "kind": span_kind, + "input": input_meta, + "output": output_meta, + "metadata": payload_metadata, + "error": error_info, + **({"tool_definitions": tool_definitions} if tool_definitions else {}), + } + + metrics: Final = self._assemble_metrics(standard_logging_payload) payload: Final[LLMObsPayload] = LLMObsPayload( parent_id=metadata_parent_id if metadata_parent_id else "undefined", @@ -314,6 +497,45 @@ class DataDogLLMObsLogger(CustomBatchLogger): ) return error_info + def _assemble_metrics(self, standard_logging_payload: StandardLoggingPayload) -> LLMMetrics: + """ + Build the span metrics, including the prompt-cache counts LLM Obs charts cache savings from. + + Cache counts resolve through the same owners the savings dashboard uses, so every provider + spelling is covered, and `non_cached_input_tokens` subtracts BOTH cache categories because + litellm's normalized prompt count includes both (the invariant the cost calculator's custom + pricing helper documents). A zero residual on a fully cached request is real data and is + emitted; a zero read or write count is absence and is not. + """ + prompt_tokens: Final = float(standard_logging_payload.get("prompt_tokens", 0)) + completion_tokens: Final = float(standard_logging_payload.get("completion_tokens", 0)) + total_tokens: Final = float(standard_logging_payload.get("total_tokens", 0)) + total_cost: Final = float(standard_logging_payload.get("response_cost", 0)) + time_to_first_token: Final = self._get_time_to_first_token_seconds(standard_logging_payload) + + raw_usage: Final = (standard_logging_payload.get("metadata") or {}).get("usage_object") + usage_object: Final = raw_usage if isinstance(raw_usage, dict) else None + cache_read: Final = float(extract_cache_read_tokens(usage_object)) + cache_write: Final = float(extract_cache_creation_tokens(usage_object)) + + metrics: Final[LLMMetrics] = { + "input_tokens": prompt_tokens, + "output_tokens": completion_tokens, + "total_tokens": total_tokens, + "total_cost": total_cost, + "time_to_first_token": time_to_first_token, + **( + { + **({"cache_read_input_tokens": cache_read} if cache_read else {}), + **({"cache_write_input_tokens": cache_write} if cache_write else {}), + "non_cached_input_tokens": max(prompt_tokens - cache_read - cache_write, 0.0), + } + if cache_read or cache_write + else {} + ), + } + return metrics + def _get_time_to_first_token_seconds(self, standard_logging_payload: StandardLoggingPayload) -> float: """ Get the time to first token in seconds @@ -335,7 +557,7 @@ class DataDogLLMObsLogger(CustomBatchLogger): def _get_response_messages( self, standard_logging_payload: StandardLoggingPayload, call_type: str | None - ) -> list[object]: + ) -> tuple[Message, ...]: """ Get the messages from the response object @@ -344,7 +566,7 @@ class DataDogLLMObsLogger(CustomBatchLogger): response_obj = standard_logging_payload.get("response") if response_obj is None: - return [] + return () # edge case: handle response_obj is a string representation of a dict if isinstance(response_obj, str): @@ -357,7 +579,7 @@ class DataDogLLMObsLogger(CustomBatchLogger): # fallback to json parsing response_obj = json.loads(str(response_obj)) except json.JSONDecodeError: - return [] + return () if call_type in [ CallTypes.completion.value, @@ -375,12 +597,12 @@ class DataDogLLMObsLogger(CustomBatchLogger): if isinstance(response_obj, dict) and "choices" in response_obj: choices: Final = response_obj["choices"] if choices and len(choices) > 0 and "message" in choices[0]: - return [choices[0]["message"]] - return [] + return _to_dd_messages([choices[0]["message"]]) + return () except (KeyError, IndexError, TypeError): # In case of any error accessing the response structure, return empty list - return [] - return [] + return () + return () def _get_datadog_span_kind( self, call_type: str | None, parent_id: str | None = None @@ -485,17 +707,6 @@ class DataDogLLMObsLogger(CustomBatchLogger): # Default fallback for unknown or passthrough operations return "llm" - def _ensure_string_content(self, messages: str | Sequence[object] | Mapping[object, object] | None) -> list[object]: - if messages is None: - return [] - if isinstance(messages, str): - return [messages] - elif isinstance(messages, list): - return [message for message in messages] - elif isinstance(messages, dict): - return [str(messages.get("content", ""))] - return [] - def _get_dd_llm_obs_payload_metadata(self, standard_logging_payload: StandardLoggingPayload) -> dict[str, object]: """ Fields to track in DD LLM Observability metadata from litellm standard logging payload @@ -524,10 +735,6 @@ class DataDogLLMObsLogger(CustomBatchLogger): spend_metrics: Final = self._get_spend_metrics(standard_logging_payload) _metadata.update({"spend_metrics": dict(spend_metrics)}) - ## extract tool calls and add to metadata - tool_call_metadata: Final = self._extract_tool_call_metadata(standard_logging_payload) - _metadata.update(tool_call_metadata) - _standard_logging_metadata: Final[dict] = dict(standard_logging_payload.get("metadata", {})) or {} _metadata.update(_standard_logging_metadata) return _metadata @@ -647,107 +854,3 @@ class DataDogLLMObsLogger(CustomBatchLogger): verbose_logger.debug("Original value: %s", user_api_key_budget_reset_at) return spend_metrics - - def _process_input_messages_preserving_tool_calls(self, messages: Sequence[object]) -> list[dict[str, object]]: - """ - Process input messages while preserving tool_calls and tool message types. - - This bypasses the lossy string conversion when tool calls are present, - allowing complex nested tool_calls objects to be preserved for Datadog. - """ - processed: Final = [] - for msg in messages: - if isinstance(msg, dict): - # Preserve messages with tool_calls or tool role as-is - if "tool_calls" in msg or msg.get("role") == "tool": - processed.append(msg) - else: - # For regular messages, still apply string conversion - converted = handle_any_messages_to_chat_completion_str_messages_conversion([msg]) - processed.extend(converted) - else: - # For non-dict messages, apply string conversion - converted = handle_any_messages_to_chat_completion_str_messages_conversion([msg]) - processed.extend(converted) - return processed - - @staticmethod - def _tool_calls_kv_pair(tool_calls: list[dict[str, Any]]) -> dict[str, object]: - """ - Extract tool call information into key-value pairs for Datadog metadata. - - Similar to OpenTelemetry's implementation but adapted for Datadog's format. - """ - kv_pairs: Final[dict[str, object]] = {} - for idx, tool_call in enumerate(tool_calls): - try: - # Extract tool call ID - tool_id = tool_call.get("id") - if tool_id: - kv_pairs[f"tool_calls.{idx}.id"] = tool_id - - # Extract tool call type - tool_type = tool_call.get("type") - if tool_type: - kv_pairs[f"tool_calls.{idx}.type"] = tool_type - - # Extract function information - function = tool_call.get("function") - if function: - function_name = function.get("name") - if function_name: - kv_pairs[f"tool_calls.{idx}.function.name"] = function_name - - function_arguments = function.get("arguments") - if function_arguments: - # Store arguments as JSON string for Datadog - if isinstance(function_arguments, str): - kv_pairs[f"tool_calls.{idx}.function.arguments"] = function_arguments - else: - import json - - kv_pairs[f"tool_calls.{idx}.function.arguments"] = json.dumps(function_arguments) - except (KeyError, TypeError, ValueError) as e: - verbose_logger.debug("DataDogLLMObs: Error processing tool call %s: %s", idx, e) - continue - - return kv_pairs - - def _extract_tool_call_metadata(self, standard_logging_payload: StandardLoggingPayload) -> dict[str, object]: - """ - Extract tool call information from both input messages and response for Datadog metadata. - """ - tool_call_metadata: Final[dict[str, object]] = {} - - try: - # Extract tool calls from input messages - messages: Final = standard_logging_payload.get("messages", []) - if messages and isinstance(messages, list): - for message in messages: - if isinstance(message, dict) and "tool_calls" in message: - tool_calls = message.get("tool_calls") - if tool_calls: - input_tool_calls_kv = self._tool_calls_kv_pair(tool_calls) - # Prefix with "input_" to distinguish from response tool calls - for key, value in input_tool_calls_kv.items(): - tool_call_metadata[f"input_{key}"] = value - - # Extract tool calls from response - response_obj: Final = standard_logging_payload.get("response") - if response_obj and isinstance(response_obj, dict): - choices: Final = response_obj.get("choices", []) - for choice in choices: - if isinstance(choice, dict): - message = choice.get("message") - if message and isinstance(message, dict): - tool_calls = message.get("tool_calls") - if tool_calls: - response_tool_calls_kv = self._tool_calls_kv_pair(tool_calls) - # Prefix with "output_" to distinguish from input tool calls - for key, value in response_tool_calls_kv.items(): - tool_call_metadata[f"output_{key}"] = value - - except Exception as e: - verbose_logger.debug("DataDogLLMObs: Error extracting tool call metadata: %s", e) - - return tool_call_metadata diff --git a/litellm/types/integrations/datadog_llm_obs.py b/litellm/types/integrations/datadog_llm_obs.py index 7853dda1213..bae876dfdd9 100644 --- a/litellm/types/integrations/datadog_llm_obs.py +++ b/litellm/types/integrations/datadog_llm_obs.py @@ -4,21 +4,58 @@ Payloads for Datadog LLM Observability Service (LLMObs) API Reference: https://docs.datadoghq.com/llm_observability/setup/api/?tab=example#api-standards """ +from collections.abc import Sequence from typing import Any, Literal -from typing_extensions import TypedDict +from typing_extensions import ReadOnly, TypedDict from litellm.types.integrations.custom_logger import StandardCustomLoggerInitParams +class ToolCall(TypedDict, total=False): + """A tool call on a message, as LLM Obs names its fields.""" + + name: ReadOnly[str] + arguments: ReadOnly[dict[str, Any] | str] # parsed object, or the raw string when it will not parse to one + tool_id: ReadOnly[str] + type: ReadOnly[str] + + +class ToolResult(TypedDict, total=False): + """The result of a tool call, as LLM Obs names its fields.""" + + name: ReadOnly[str] + result: ReadOnly[str] + tool_id: ReadOnly[str] + type: ReadOnly[str] + + +class ToolDefinition(TypedDict, total=False): + """A tool the model was offered on the request.""" + + name: ReadOnly[str] + description: ReadOnly[str] + schema: ReadOnly[dict[str, Any]] + + +class Message(TypedDict, total=False): + """A message on a span, as LLM Obs names its fields.""" + + content: ReadOnly[str] + role: ReadOnly[str] + reasoning_content: ReadOnly[str] + tool_calls: ReadOnly[Sequence[ToolCall]] + tool_results: ReadOnly[Sequence[ToolResult]] + + class InputMeta(TypedDict): - messages: list[ - dict[str, Any] # changed to fit with tool calls + messages: Sequence[ + Message | dict[str, Any] # changed to fit with tool calls ] # Relevant Issue: https://github.com/BerriAI/litellm/issues/9494 class OutputMeta(TypedDict): - messages: list[Any] + messages: Sequence[Any] class DDLLMObsError(TypedDict, total=False): @@ -36,6 +73,7 @@ class Meta(TypedDict, total=False): output: OutputMeta # The span's output information. metadata: dict[str, Any] error: DDLLMObsError | None # Error information on the span + tool_definitions: ReadOnly[Sequence[ToolDefinition]] # The tools offered to the model on this request class LLMMetrics(TypedDict, total=False): @@ -45,6 +83,9 @@ class LLMMetrics(TypedDict, total=False): time_to_first_token: float time_per_output_token: float total_cost: float + cache_read_input_tokens: ReadOnly[float] + cache_write_input_tokens: ReadOnly[float] + non_cached_input_tokens: ReadOnly[float] class LLMObsPayload(TypedDict, total=False): diff --git a/tests/test_litellm/integrations/datadog/test_datadog_llm_obs.py b/tests/test_litellm/integrations/datadog/test_datadog_llm_obs.py new file mode 100644 index 00000000000..2d0605e3b7f --- /dev/null +++ b/tests/test_litellm/integrations/datadog/test_datadog_llm_obs.py @@ -0,0 +1,469 @@ +""" +Regression tests for the Datadog LLM Observability payload schema (issue #35786). + +Datadog renders tool calls, tool results and prompt-cache savings only from the fields its +own schema names. These assert on the payload `create_llm_obs_payload` actually hands the +intake, so a regression that moves data back into `meta.metadata` fails here. + +Fixtures mirror what a live proxy run recorded on the callback, including the provider +spelling of prompt-cache counts (`prompt_tokens_details.cached_tokens`). +""" + +import json +import os +from datetime import datetime, timedelta +from typing import Any +from unittest.mock import patch + +import pytest + +from litellm.integrations.datadog.datadog_llm_obs import DataDogLLMObsLogger +from litellm.litellm_core_utils.safe_json_dumps import safe_dumps + +TOOL_DEFINITION: dict[str, Any] = { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get current weather for a city", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + }, + }, +} + +ASSISTANT_TOOL_CALL: dict[str, Any] = { + "id": "call_abc123", + "type": "function", + "function": {"name": "get_weather", "arguments": '{"city":"Paris","unit":"c"}'}, +} + + +@pytest.fixture +def logger() -> DataDogLLMObsLogger: + with patch.dict(os.environ, {"DD_API_KEY": "k", "DD_SITE": "us5.datadoghq.com"}, clear=True): + with patch("asyncio.create_task"): + return DataDogLLMObsLogger() + + +NOT_GIVEN: Any = object() + + +def build_payload( + messages: Any = NOT_GIVEN, + response_message: dict[str, Any] | None = None, + usage_object: dict[str, Any] | None = None, + model_parameters: dict[str, Any] | None = None, + prompt_tokens: int = 4447, +) -> dict[str, Any]: + return { + "standard_logging_object": { + "call_type": "acompletion", + "messages": [{"role": "user", "content": "hi"}] if messages is NOT_GIVEN else messages, + "response": {"choices": [{"message": response_message or {"role": "assistant", "content": "hello"}}]}, + "model_parameters": model_parameters or {}, + "metadata": {"usage_object": usage_object} if usage_object is not None else {}, + "prompt_tokens": prompt_tokens, + "completion_tokens": 507, + "total_tokens": prompt_tokens + 507, + "response_cost": 0.02, + "status": "success", + }, + "litellm_params": {"metadata": {}}, + } + + +def build(logger: DataDogLLMObsLogger, **kwargs: Any) -> dict[str, Any]: + """Build a span and read it back as the JSON the intake receives, not as Python objects.""" + start = datetime(2026, 9, 1, 12, 0, 0) + payload = logger.create_llm_obs_payload(build_payload(**kwargs), start, start + timedelta(seconds=2)) + return json.loads(safe_dumps(payload)) + + +def test_output_tool_calls_use_the_datadog_tool_call_schema(logger: DataDogLLMObsLogger) -> None: + """Datadog reads name/arguments/tool_id off the tool call; OpenAI nests them under `function`.""" + payload = build( + logger, + response_message={"role": "assistant", "content": None, "tool_calls": [ASSISTANT_TOOL_CALL]}, + ) + + message = payload["meta"]["output"]["messages"][0] + assert message["tool_calls"] == [ + { + "name": "get_weather", + "arguments": {"city": "Paris", "unit": "c"}, + "tool_id": "call_abc123", + "type": "function", + } + ] + assert "function" not in message["tool_calls"][0] + + +def test_tool_calls_are_not_duplicated_into_metadata(logger: DataDogLLMObsLogger) -> None: + """The flat `output_tool_calls.*` keys were a second copy of a fact that now has its own field.""" + payload = build( + logger, + response_message={"role": "assistant", "content": None, "tool_calls": [ASSISTANT_TOOL_CALL]}, + ) + + assert [key for key in payload["meta"]["metadata"] if "tool_calls." in key] == [] + + +def test_tool_result_message_links_back_to_its_tool_call(logger: DataDogLLMObsLogger) -> None: + """Datadog pairs a result with its call through tool_id, and names the tool from the call.""" + payload = build( + logger, + messages=[ + {"role": "user", "content": "Weather in Paris?"}, + {"role": "assistant", "content": None, "tool_calls": [ASSISTANT_TOOL_CALL]}, + {"role": "tool", "tool_call_id": "call_abc123", "content": '{"temp_c": 18}'}, + ], + ) + + tool_message = payload["meta"]["input"]["messages"][2] + assert tool_message["tool_results"] == [ + {"name": "get_weather", "result": '{"temp_c": 18}', "tool_id": "call_abc123", "type": "function"} + ] + + +def test_tool_result_without_a_matching_call_still_reports_its_id(logger: DataDogLLMObsLogger) -> None: + """A truncated conversation loses the call, so the name is unknown but the link must survive.""" + payload = build( + logger, + messages=[{"role": "tool", "tool_call_id": "call_orphan", "content": "42"}], + ) + + assert payload["meta"]["input"]["messages"][0]["tool_results"] == [ + {"name": "", "result": "42", "tool_id": "call_orphan", "type": "function"} + ] + + +def test_cache_tokens_are_reported_as_span_metrics(logger: DataDogLLMObsLogger) -> None: + """ + Datadog charts cache savings from span metrics; nested usage_object is not read for it. + + litellm's normalized prompt count includes both cache categories, so the three cache + metrics must partition input_tokens: read + write + non_cached == input. + """ + payload = build( + logger, + usage_object={"prompt_tokens_details": {"cached_tokens": 4300, "cache_write_tokens": 95}}, + ) + + metrics = payload["metrics"] + assert metrics["cache_read_input_tokens"] == 4300.0 + assert metrics["cache_write_input_tokens"] == 95.0 + assert metrics["non_cached_input_tokens"] == 4447.0 - 4300.0 - 95.0 + assert ( + metrics["cache_read_input_tokens"] + metrics["cache_write_input_tokens"] + metrics["non_cached_input_tokens"] + == metrics["input_tokens"] + ) + + +def test_cache_write_tokens_are_not_counted_as_non_cached(logger: DataDogLLMObsLogger) -> None: + """A cache-priming request must not report its primed prefix as full-price uncached input.""" + payload = build(logger, usage_object={"prompt_tokens_details": {"cache_write_tokens": 4000}}) + + assert payload["metrics"]["cache_write_input_tokens"] == 4000.0 + assert payload["metrics"]["non_cached_input_tokens"] == 4447.0 - 4000.0 + assert "cache_read_input_tokens" not in payload["metrics"] + + +def test_a_fully_cached_request_reports_a_zero_non_cached_count(logger: DataDogLLMObsLogger) -> None: + """Zero residual is real data: everything was served from cache. Inconsistent counts clamp to it.""" + payload = build( + logger, + usage_object={"prompt_tokens_details": {"cached_tokens": 4352, "cache_write_tokens": 95}}, + ) + + assert payload["metrics"]["non_cached_input_tokens"] == 0.0 + + +def test_anthropic_top_level_cache_keys_are_read(logger: DataDogLLMObsLogger) -> None: + """A raw Anthropic usage dict records the counts top level, not under prompt_tokens_details.""" + payload = build( + logger, + usage_object={"cache_read_input_tokens": 4300, "cache_creation_input_tokens": 95}, + ) + + metrics = payload["metrics"] + assert metrics["cache_read_input_tokens"] == 4300.0 + assert metrics["cache_write_input_tokens"] == 95.0 + assert metrics["non_cached_input_tokens"] == 4447.0 - 4300.0 - 95.0 + + +def test_cache_metrics_come_from_the_normalized_field_not_the_anthropic_one(logger: DataDogLLMObsLogger) -> None: + """ + litellm normalizes every provider's cache counters into prompt_tokens_details. + + A real cached request from a non-Anthropic provider carries only `cached_tokens`, so + reading the Anthropic-specific `cache_read_input_tokens` key reports nothing for it. + """ + payload = build( + logger, + usage_object={"prompt_tokens_details": {"audio_tokens": None, "cached_tokens": 4096}}, + prompt_tokens=4335, + ) + + assert payload["metrics"]["cache_read_input_tokens"] == 4096.0 + assert payload["metrics"]["non_cached_input_tokens"] == 4335.0 - 4096.0 + + +@pytest.mark.parametrize( + "usage_object", + [ + {"prompt_tokens_details": {"cache_write_tokens": 95}}, + {"prompt_tokens_details": {"cache_creation_tokens": 95}}, + {"cache_creation_input_tokens": 95}, + ], +) +def test_every_spelling_of_cache_write_tokens_is_read( + logger: DataDogLLMObsLogger, usage_object: dict[str, Any] +) -> None: + """A raw usage dict that bypassed litellm's normalizer can carry any provider's spelling.""" + payload = build(logger, usage_object=usage_object) + + assert payload["metrics"]["cache_write_input_tokens"] == 95.0 + + +def test_a_cache_read_does_not_emit_a_zero_cache_write(logger: DataDogLLMObsLogger) -> None: + """A zero write on every cache-read span would drag Datadog's cache-write average to nothing.""" + payload = build(logger, usage_object={"prompt_tokens_details": {"cached_tokens": 4096}}) + + assert payload["metrics"]["cache_read_input_tokens"] == 4096.0 + assert "cache_write_input_tokens" not in payload["metrics"] + + +def test_no_cache_keys_when_the_provider_reports_no_caching(logger: DataDogLLMObsLogger) -> None: + """An uncached request must not gain zero-valued cache metrics that dilute cache dashboards.""" + payload = build(logger, usage_object={"prompt_tokens_details": None}) + + assert "cache_read_input_tokens" not in payload["metrics"] + assert "cache_write_input_tokens" not in payload["metrics"] + assert "non_cached_input_tokens" not in payload["metrics"] + + +def test_tool_definitions_are_sent_on_meta(logger: DataDogLLMObsLogger) -> None: + payload = build(logger, model_parameters={"tools": [TOOL_DEFINITION]}) + + assert payload["meta"]["tool_definitions"] == [ + { + "name": "get_weather", + "description": "Get current weather for a city", + "schema": {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]}, + } + ] + + +def test_tool_definitions_accept_the_bare_anthropic_shape(logger: DataDogLLMObsLogger) -> None: + """The Anthropic surface declares tools unwrapped, with input_schema instead of parameters.""" + payload = build( + logger, + model_parameters={"tools": [{"name": "get_weather", "description": "d", "input_schema": {"type": "object"}}]}, + ) + + assert payload["meta"]["tool_definitions"] == [ + {"name": "get_weather", "description": "d", "schema": {"type": "object"}} + ] + + +def test_meta_omits_tool_definitions_when_no_tools_were_offered(logger: DataDogLLMObsLogger) -> None: + assert "tool_definitions" not in build(logger)["meta"] + + +def test_unparseable_tool_arguments_are_preserved_rather_than_dropped(logger: DataDogLLMObsLogger) -> None: + """A truncated argument string is still the only record of what the model tried to call.""" + payload = build( + logger, + response_message={ + "role": "assistant", + "content": None, + "tool_calls": [{"id": "c1", "type": "function", "function": {"name": "f", "arguments": '{"city":'}}], + }, + ) + + assert payload["meta"]["output"]["messages"][0]["tool_calls"][0]["arguments"] == '{"city":' + + +def test_oversized_tool_arguments_ship_unparsed(logger: DataDogLLMObsLogger) -> None: + """ + Decoding attacker-sized compact JSON multiplies memory for a span that is only logging. + + This payload is perfectly valid JSON, so the only reason it arrives as a string is the + size bound; a smaller copy of the same shape comes back as an object below. + """ + oversized = '{"a":"' + "x" * 300_000 + '"}' + + payload = build( + logger, + response_message={ + "role": "assistant", + "content": None, + "tool_calls": [{"id": "c1", "type": "function", "function": {"name": "f", "arguments": oversized}}], + }, + ) + + assert payload["meta"]["output"]["messages"][0]["tool_calls"][0]["arguments"] == oversized + + +def test_valid_arguments_below_the_bound_still_parse(logger: DataDogLLMObsLogger) -> None: + """The size bound must not swallow ordinary arguments; this is the oversized test's control.""" + payload = build( + logger, + response_message={ + "role": "assistant", + "content": None, + "tool_calls": [ + {"id": "c1", "type": "function", "function": {"name": "f", "arguments": '{"a":"' + "x" * 64 + '"}'}} + ], + }, + ) + + assert payload["meta"]["output"]["messages"][0]["tool_calls"][0]["arguments"] == {"a": "x" * 64} + + +def test_a_result_is_named_even_when_its_call_had_unparseable_arguments(logger: DataDogLLMObsLogger) -> None: + """Correlating a result to its call reads ids and names, so bad arguments cannot break linking.""" + payload = build( + logger, + messages=[ + { + "role": "assistant", + "content": None, + "tool_calls": [ + {"id": "call_abc123", "type": "function", "function": {"name": "get_weather", "arguments": "{"}} + ], + }, + {"role": "tool", "tool_call_id": "call_abc123", "content": "18C"}, + ], + ) + + assert payload["meta"]["input"]["messages"][1]["tool_results"] == [ + {"name": "get_weather", "result": "18C", "tool_id": "call_abc123", "type": "function"} + ] + + +def test_deeply_nested_tool_arguments_do_not_drop_the_span(logger: DataDogLLMObsLogger) -> None: + """json.loads raises RecursionError, not JSONDecodeError, on hostile nesting.""" + payload = build( + logger, + response_message={ + "role": "assistant", + "content": None, + "tool_calls": [{"id": "c1", "type": "function", "function": {"name": "f", "arguments": "[" * 50_000}}], + }, + ) + + assert payload["meta"]["output"]["messages"][0]["tool_calls"][0]["arguments"] == "[" * 50_000 + + +def test_tool_arguments_that_parse_to_a_non_object_stay_a_string(logger: DataDogLLMObsLogger) -> None: + """Datadog types arguments as an object, so a bare JSON scalar must not land there as one.""" + payload = build( + logger, + response_message={ + "role": "assistant", + "content": None, + "tool_calls": [{"id": "c1", "type": "function", "function": {"name": "f", "arguments": "42"}}], + }, + ) + + assert payload["meta"]["output"]["messages"][0]["tool_calls"][0]["arguments"] == "42" + + +def test_a_tool_without_a_name_is_not_offered_as_a_definition(logger: DataDogLLMObsLogger) -> None: + """A nameless tool cannot be matched to a call, so it is dropped rather than sent blank.""" + payload = build(logger, model_parameters={"tools": [{"function": {"description": "no name"}}, TOOL_DEFINITION]}) + + assert [tool["name"] for tool in payload["meta"]["tool_definitions"]] == ["get_weather"] + + +def test_a_tool_definition_without_a_schema_omits_the_field(logger: DataDogLLMObsLogger) -> None: + """An empty schema object would read as a tool that takes no arguments, which is a different claim.""" + payload = build(logger, model_parameters={"tools": [{"name": "ping", "description": "d"}]}) + + assert payload["meta"]["tool_definitions"] == [{"name": "ping", "description": "d"}] + + +def test_a_non_dict_message_still_reaches_datadog(logger: DataDogLLMObsLogger) -> None: + """Callers can log arbitrary message payloads, and dropping the span over one loses the request.""" + payload = build(logger, messages=["just a bare string"]) + + assert payload["meta"]["input"]["messages"] == [{"input": "just a bare string"}] + + +def test_messages_logged_as_a_bare_string_still_reach_datadog(logger: DataDogLLMObsLogger) -> None: + payload = build(logger, messages="the whole prompt as one string") + + assert payload["meta"]["input"]["messages"] == [{"input": "the whole prompt as one string"}] + + +def test_non_chat_call_types_log_an_empty_input(logger: DataDogLLMObsLogger) -> None: + """Embedding and image calls carry no messages; fabricating an "None" turn misreads in Datadog.""" + payload = build(logger, messages=None) + + assert payload["meta"]["input"]["messages"] == [] + + +def test_anthropic_tool_blocks_map_to_tool_calls_and_results(logger: DataDogLLMObsLogger) -> None: + """/v1/messages carries tool traffic as content blocks, not OpenAI fields.""" + payload = build( + logger, + messages=[ + {"role": "user", "content": [{"type": "text", "text": "Weather in Tokyo?"}]}, + { + "role": "assistant", + "content": [{"type": "tool_use", "id": "toolu_1", "name": "get_weather", "input": {"city": "Tokyo"}}], + }, + {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "toolu_1", "content": "18C"}]}, + ], + ) + + assistant, result_turn = payload["meta"]["input"]["messages"][1:3] + assert assistant["tool_calls"] == [ + {"name": "get_weather", "arguments": {"city": "Tokyo"}, "tool_id": "toolu_1", "type": "tool_use"} + ] + assert result_turn["tool_results"] == [ + {"name": "get_weather", "result": "18C", "tool_id": "toolu_1", "type": "function"} + ] + + +def test_content_with_no_text_parts_is_preserved_not_blanked(logger: DataDogLLMObsLogger) -> None: + """A content list the mapper does not understand must ride along, not be erased.""" + blocks = [{"type": "image_url", "image_url": {"url": "https://example.com/x.png"}}] + payload = build(logger, messages=[{"role": "user", "content": blocks}]) + + assert payload["meta"]["input"]["messages"][0]["content"] == blocks + + +def test_multimodal_content_parts_are_flattened_to_text(logger: DataDogLLMObsLogger) -> None: + """Datadog types Message.content as a string, so content lists collapse to their text.""" + payload = build( + logger, + messages=[ + {"role": "user", "content": [{"type": "text", "text": "describe "}, {"type": "text", "text": "this"}]} + ], + ) + + assert payload["meta"]["input"]["messages"][0]["content"] == "describe this" + + +def test_mapping_input_messages_does_not_mutate_the_shared_payload(logger: DataDogLLMObsLogger) -> None: + """Sibling callbacks read the same messages list, so flattening must not write through it.""" + messages: list[dict[str, Any]] = [{"role": "user", "content": [{"type": "text", "text": "hi"}]}] + kwargs = build_payload(messages=messages) + start = datetime(2026, 9, 1, 12, 0, 0) + + logger.create_llm_obs_payload(kwargs, start, start + timedelta(seconds=1)) + + assert messages[0]["content"] == [{"type": "text", "text": "hi"}] + + +def test_reasoning_content_survives_the_mapping(logger: DataDogLLMObsLogger) -> None: + payload = build( + logger, + response_message={"role": "assistant", "content": "answer", "reasoning_content": "thinking"}, + ) + + assert payload["meta"]["output"]["messages"][0]["reasoning_content"] == "thinking" From 6d0367ce350402e93651f1391afdde1ee2553b0f Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 1 Sep 2026 18:03:18 -0700 Subject: [PATCH 46/97] feat(prometheus): expose per-key and per-team rate limit allowed and used gauges (#39236) Co-authored-by: yassin Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/integrations/prometheus.py | 120 ++++++++- .../bounded_prometheus_series_tracker.py | 4 + litellm/types/integrations/prometheus.py | 20 ++ .../test_prometheus_client_ip_user_agent.py | 1 + .../test_prometheus_rate_limit_labels.py | 252 ++++++++++++++++++ 5 files changed, 395 insertions(+), 2 deletions(-) diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index add91033ff3..975a9bd8639 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -8,6 +8,7 @@ import math import os import sys from collections.abc import Awaitable, Callable, Mapping, Sequence +from dataclasses import replace from datetime import datetime, timedelta from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, TypeVar, cast @@ -58,6 +59,7 @@ from litellm.types.utils import ( if TYPE_CHECKING: from apscheduler.schedulers.asyncio import AsyncIOScheduler + from prometheus_client import Gauge from prometheus_client.metrics import MetricWrapperBase from litellm.router import Router @@ -476,6 +478,30 @@ class PrometheusLogger(CustomLogger): labelnames=self.get_labels_for_metric("litellm_remaining_api_key_tokens_for_model"), ) + self.litellm_api_key_rate_limit_allowed_metric = self._gauge_factory( + "litellm_api_key_rate_limit_allowed_metric", + "Configured rate limit for the API Key in the current window (rpm_limit / tpm_limit), by rate_limit_type", + labelnames=self.get_labels_for_metric("litellm_api_key_rate_limit_allowed_metric"), + ) + + self.litellm_api_key_rate_limit_used_metric = self._gauge_factory( + "litellm_api_key_rate_limit_used_metric", + "Requests or tokens the API Key has consumed in the current rate limit window, by rate_limit_type", + labelnames=self.get_labels_for_metric("litellm_api_key_rate_limit_used_metric"), + ) + + self.litellm_team_rate_limit_allowed_metric = self._gauge_factory( + "litellm_team_rate_limit_allowed_metric", + "Configured rate limit for the Team in the current window (team rpm_limit / tpm_limit), by rate_limit_type", + labelnames=self.get_labels_for_metric("litellm_team_rate_limit_allowed_metric"), + ) + + self.litellm_team_rate_limit_used_metric = self._gauge_factory( + "litellm_team_rate_limit_used_metric", + "Requests or tokens the Team has consumed in the current rate limit window, by rate_limit_type", + labelnames=self.get_labels_for_metric("litellm_team_rate_limit_used_metric"), + ) + ######################################## # LLM API Deployment Metrics / analytics ######################################## @@ -1475,6 +1501,11 @@ class PrometheusLogger(CustomLogger): model_id=enum_values.model_id, ) + self._set_key_and_team_rate_limit_metrics( + standard_logging_payload=standard_logging_payload, # pyright: ignore[reportArgumentType] # isinstance(dict) above narrows the TypedDict to dict[Unknown, Unknown] + enum_values=enum_values, + ) + # set latency metrics self._set_latency_metrics( kwargs=kwargs, @@ -2002,17 +2033,102 @@ class PrometheusLogger(CustomLogger): """ if standard_logging_payload is None: return None + return PrometheusLogger._get_int_from_v3_rate_limit_headers( + standard_logging_payload=standard_logging_payload, + header_name=f"x-ratelimit-model_per_key-remaining-{rate_limit_type}", + ) + + @staticmethod + def _get_int_from_v3_rate_limit_headers( + standard_logging_payload: StandardLoggingPayload, + header_name: str, + ) -> int | None: hidden_params: Final = standard_logging_payload.get("hidden_params") if hidden_params is None: return None - additional_headers: Final = hidden_params.get("additional_headers") + additional_headers: Final[Mapping[str, object] | None] = hidden_params.get("additional_headers") if additional_headers is None: return None - value: Final = dict(additional_headers).get(f"x-ratelimit-model_per_key-remaining-{rate_limit_type}") + value: Final = additional_headers.get(header_name) if isinstance(value, bool) or not isinstance(value, int): return None return value + def _set_key_and_team_rate_limit_metrics( + self, + standard_logging_payload: StandardLoggingPayload, + enum_values: UserAPIKeyLabelValues, + ) -> None: + """ + Export the key-level and team-level RPM / TPM limit and current window + usage from the ``x-ratelimit-{api_key,team}-{limit,remaining}-*`` + headers the v3 rate limiter mirrors into the logging payload. The + limiter already read these counters (from Redis when configured) on + the request path, so no extra store lookup happens here. Descriptors + without a configured limit emit no header, so their series is removed + rather than left at the value from before the limit was dropped. + """ + descriptor_gauges: Final[ + tuple[tuple[Literal["api_key", "team"], DEFINED_PROMETHEUS_METRICS, Gauge, Gauge], ...] + ] = ( + ( + "api_key", + "litellm_api_key_rate_limit_allowed_metric", + self.litellm_api_key_rate_limit_allowed_metric, + self.litellm_api_key_rate_limit_used_metric, + ), + ( + "team", + "litellm_team_rate_limit_allowed_metric", + self.litellm_team_rate_limit_allowed_metric, + self.litellm_team_rate_limit_used_metric, + ), + ) + for descriptor_key, metric_name, allowed_gauge, used_gauge in descriptor_gauges: + for rate_limit_type in ("requests", "tokens"): + self._set_rate_limit_allowed_and_used_gauges( + standard_logging_payload=standard_logging_payload, + enum_values=enum_values, + descriptor_key=descriptor_key, + metric_name=metric_name, + allowed_gauge=allowed_gauge, + used_gauge=used_gauge, + rate_limit_type=rate_limit_type, + ) + + def _set_rate_limit_allowed_and_used_gauges( + self, + standard_logging_payload: StandardLoggingPayload, + enum_values: UserAPIKeyLabelValues, + descriptor_key: Literal["api_key", "team"], + metric_name: DEFINED_PROMETHEUS_METRICS, + allowed_gauge: Gauge, + used_gauge: Gauge, + rate_limit_type: Literal["requests", "tokens"], + ) -> None: + limit: Final = self._get_int_from_v3_rate_limit_headers( + standard_logging_payload=standard_logging_payload, + header_name=f"x-ratelimit-{descriptor_key}-limit-{rate_limit_type}", + ) + remaining: Final = self._get_int_from_v3_rate_limit_headers( + standard_logging_payload=standard_logging_payload, + header_name=f"x-ratelimit-{descriptor_key}-remaining-{rate_limit_type}", + ) + labelled_values: Final = replace(enum_values, rate_limit_type=rate_limit_type) + labelnames: Final = self.get_labels_for_metric(metric_name) + labels: Final = prometheus_label_factory( + supported_enum_labels=labelnames, + enum_values=labelled_values, + label_context=PrometheusLabelFactoryContext(labelled_values), + ) + if limit is None or remaining is None: + label_values: Final = tuple(labels.get(label) for label in labelnames) + self._bounded_prometheus_series_tracker.remove_series(allowed_gauge, label_values) + self._bounded_prometheus_series_tracker.remove_series(used_gauge, label_values) + return + allowed_gauge.labels(**labels).set(limit) + used_gauge.labels(**labels).set(limit - remaining) + def _set_virtual_key_rate_limit_metrics( self, user_api_key: str | None, diff --git a/litellm/integrations/prometheus_helpers/bounded_prometheus_series_tracker.py b/litellm/integrations/prometheus_helpers/bounded_prometheus_series_tracker.py index c54790b8ae7..c1ccf09d5d6 100644 --- a/litellm/integrations/prometheus_helpers/bounded_prometheus_series_tracker.py +++ b/litellm/integrations/prometheus_helpers/bounded_prometheus_series_tracker.py @@ -60,6 +60,10 @@ class BoundedPrometheusSeriesTracker: break del series[tracked_label_values] + def remove_series(self, metric: object, label_values: tuple[str | None, ...]) -> bool: + """Drop one child series, True when it is gone (removed or never existed).""" + return self._remove_metric_child(metric, label_values) + def _should_run_ttl_cleanup( self, metric_name: str, diff --git a/litellm/types/integrations/prometheus.py b/litellm/types/integrations/prometheus.py index 01ed8b08571..8498b6f6d00 100644 --- a/litellm/types/integrations/prometheus.py +++ b/litellm/types/integrations/prometheus.py @@ -270,6 +270,10 @@ DEFINED_PROMETHEUS_METRICS = Literal[ "litellm_deployment_rpm_limit", "litellm_remaining_api_key_requests_for_model", "litellm_remaining_api_key_tokens_for_model", + "litellm_api_key_rate_limit_allowed_metric", + "litellm_api_key_rate_limit_used_metric", + "litellm_team_rate_limit_allowed_metric", + "litellm_team_rate_limit_used_metric", "litellm_llm_api_failed_requests_metric", "litellm_callback_logging_failures_metric", "litellm_in_flight_requests", @@ -775,6 +779,22 @@ class PrometheusMetricLabels: UserAPIKeyLabelNames.MODEL_ID.value, ] + litellm_api_key_rate_limit_allowed_metric: ClassVar[tuple[str, ...]] = ( + UserAPIKeyLabelNames.API_KEY_HASH.value, + UserAPIKeyLabelNames.API_KEY_ALIAS.value, + UserAPIKeyLabelNames.RATE_LIMIT_TYPE.value, + ) + + litellm_api_key_rate_limit_used_metric = litellm_api_key_rate_limit_allowed_metric + + litellm_team_rate_limit_allowed_metric: ClassVar[tuple[str, ...]] = ( + UserAPIKeyLabelNames.TEAM.value, + UserAPIKeyLabelNames.TEAM_ALIAS.value, + UserAPIKeyLabelNames.RATE_LIMIT_TYPE.value, + ) + + litellm_team_rate_limit_used_metric = litellm_team_rate_limit_allowed_metric + litellm_llm_api_failed_requests_metric = [ UserAPIKeyLabelNames.END_USER.value, UserAPIKeyLabelNames.API_KEY_HASH.value, diff --git a/tests/test_litellm/integrations/test_prometheus_client_ip_user_agent.py b/tests/test_litellm/integrations/test_prometheus_client_ip_user_agent.py index 029b097cb75..ea661d2ea78 100644 --- a/tests/test_litellm/integrations/test_prometheus_client_ip_user_agent.py +++ b/tests/test_litellm/integrations/test_prometheus_client_ip_user_agent.py @@ -93,6 +93,7 @@ async def test_async_post_call_success_hook_includes_client_ip_user_agent(): logger._increment_token_metrics = MagicMock() logger._increment_remaining_budget_metrics = AsyncMock() logger._set_virtual_key_rate_limit_metrics = MagicMock() + logger._set_key_and_team_rate_limit_metrics = MagicMock() logger._set_latency_metrics = MagicMock() logger.set_llm_deployment_success_metrics = MagicMock() logger._increment_cache_metrics = MagicMock() diff --git a/tests/test_litellm/integrations/test_prometheus_rate_limit_labels.py b/tests/test_litellm/integrations/test_prometheus_rate_limit_labels.py index 9c6d2e018ff..bf1d68c7714 100644 --- a/tests/test_litellm/integrations/test_prometheus_rate_limit_labels.py +++ b/tests/test_litellm/integrations/test_prometheus_rate_limit_labels.py @@ -13,6 +13,7 @@ Covers two follow-up gaps to the unified rate-limit error work: 429s don't silently break when the new class lands. """ +from collections.abc import Mapping from unittest.mock import MagicMock, patch import pytest @@ -471,3 +472,254 @@ def test_should_ignore_non_int_v3_header_values(bad_value): logger.litellm_remaining_api_key_tokens_for_model.labels.return_value.set.assert_called_once_with( sys.maxsize ) + + +KEY_AND_TEAM_RATE_LIMIT_METRICS = ( + "litellm_api_key_rate_limit_allowed_metric", + "litellm_api_key_rate_limit_used_metric", + "litellm_team_rate_limit_allowed_metric", + "litellm_team_rate_limit_used_metric", +) + + +def _clear_prometheus_registry() -> None: + from prometheus_client import REGISTRY + + for collector in list(REGISTRY._collector_to_names.keys()): + try: + REGISTRY.unregister(collector) + except Exception: + pass + + +def _collected_samples(metric_name: str) -> dict[tuple[tuple[str, str], ...], float]: + from prometheus_client import REGISTRY + + return { + tuple(sorted(sample.labels.items())): sample.value + for metric in REGISTRY.collect() + for sample in metric.samples + if sample.name == metric_name + } + + +def _success_kwargs_with_rate_limit_headers(additional_headers: Mapping[str, object] | None) -> dict[str, object]: + return { + "model": "claude-haiku-4-5", + "litellm_params": {"metadata": {}}, + "standard_logging_object": { + "id": "t", + "call_type": "completion", + "response_cost": 0.001, + "status": "success", + "total_tokens": 20, + "prompt_tokens": 15, + "completion_tokens": 5, + "startTime": 1.0, + "endTime": 2.0, + "completionStartTime": 1.5, + "model": "claude-haiku-4-5", + "model_id": "model-123", + "model_group": "anthropic-haiku-4-5", + "api_base": "https://api.anthropic.com", + "custom_llm_provider": "anthropic", + "request_tags": [], + "end_user": None, + "cache_hit": False, + "stream": False, + "response": None, + "model_parameters": None, + "metadata": { + "user_api_key_hash": "key-hash", + "user_api_key_alias": "key-alias", + "user_api_key_team_id": "team-id", + "user_api_key_team_alias": "team-alias", + "user_api_key_user_id": "u", + "user_api_key_user_email": "e@x.com", + "user_api_key_org_id": None, + "user_api_key_org_alias": None, + "requester_metadata": None, + "user_api_key_end_user_id": None, + "usage_object": None, + }, + "hidden_params": { + "litellm_overhead_time_ms": None, + "additional_headers": additional_headers, + }, + }, + } + + +async def _run_success_event( + additional_headers: Mapping[str, object] | None, logger: PrometheusLogger | None = None +) -> None: + import datetime + + now = datetime.datetime.now() + await (logger or PrometheusLogger()).async_log_success_event( + _success_kwargs_with_rate_limit_headers(additional_headers), None, now, now + ) + + +@pytest.mark.asyncio +async def test_should_emit_key_and_team_rate_limit_allowed_and_used_from_v3_headers(): + """ + LIT-1672: the v3 limiter mirrors ``x-ratelimit-{api_key,team}-{limit,remaining}-*`` + into the logging payload. The gauges must expose the configured limit as-is + and the window consumption as ``limit - remaining`` for each key / team + dimension, split by ``rate_limit_type``. + """ + _clear_prometheus_registry() + try: + await _run_success_event( + { + "x-ratelimit-api_key-limit-requests": 10, + "x-ratelimit-api_key-remaining-requests": 7, + "x-ratelimit-api_key-limit-tokens": 20000, + "x-ratelimit-api_key-remaining-tokens": 19947, + "x-ratelimit-team-limit-requests": 50, + "x-ratelimit-team-remaining-requests": 47, + "x-ratelimit-team-limit-tokens": 40000, + "x-ratelimit-team-remaining-tokens": 39960, + "x-ratelimit-model_per_key-limit-requests": 5, + "x-ratelimit-model_per_key-remaining-requests": 1, + } + ) + + key_requests = ( + ("api_key_alias", "key-alias"), + ("hashed_api_key", "key-hash"), + ("rate_limit_type", "requests"), + ) + key_tokens = ( + ("api_key_alias", "key-alias"), + ("hashed_api_key", "key-hash"), + ("rate_limit_type", "tokens"), + ) + team_requests = ( + ("rate_limit_type", "requests"), + ("team", "team-id"), + ("team_alias", "team-alias"), + ) + team_tokens = ( + ("rate_limit_type", "tokens"), + ("team", "team-id"), + ("team_alias", "team-alias"), + ) + + assert _collected_samples("litellm_api_key_rate_limit_allowed_metric") == { + key_requests: 10, + key_tokens: 20000, + } + assert _collected_samples("litellm_api_key_rate_limit_used_metric") == { + key_requests: 3, + key_tokens: 53, + } + assert _collected_samples("litellm_team_rate_limit_allowed_metric") == { + team_requests: 50, + team_tokens: 40000, + } + assert _collected_samples("litellm_team_rate_limit_used_metric") == { + team_requests: 3, + team_tokens: 40, + } + finally: + _clear_prometheus_registry() + + +@pytest.mark.asyncio +async def test_should_emit_only_the_dimensions_the_limiter_enforced(): + """ + A key with only ``rpm_limit`` set and no team limits produces only the + key/requests headers, so no tokens series and no team series may appear + (a phantom 0 or sys.maxsize series would misreport an unlimited dimension). + """ + _clear_prometheus_registry() + try: + await _run_success_event( + { + "x-ratelimit-api_key-limit-requests": 10, + "x-ratelimit-api_key-remaining-requests": 10, + } + ) + + key_requests = ( + ("api_key_alias", "key-alias"), + ("hashed_api_key", "key-hash"), + ("rate_limit_type", "requests"), + ) + assert _collected_samples("litellm_api_key_rate_limit_allowed_metric") == {key_requests: 10} + assert _collected_samples("litellm_api_key_rate_limit_used_metric") == {key_requests: 0} + assert _collected_samples("litellm_team_rate_limit_allowed_metric") == {} + assert _collected_samples("litellm_team_rate_limit_used_metric") == {} + finally: + _clear_prometheus_registry() + + +@pytest.mark.asyncio +async def test_should_drop_key_and_team_series_once_the_limiter_stops_reporting_a_limit(): + """ + Removing a key's ``rpm_limit`` / ``tpm_limit`` (or a team's ``tpm_limit``) + makes the v3 limiter stop emitting that descriptor's headers on later + requests. The old allowed/used samples must disappear instead of keeping + a limit that no longer exists on the scrape. + """ + _clear_prometheus_registry() + try: + logger = PrometheusLogger() + await _run_success_event( + { + "x-ratelimit-api_key-limit-requests": 10, + "x-ratelimit-api_key-remaining-requests": 7, + "x-ratelimit-api_key-limit-tokens": 20000, + "x-ratelimit-api_key-remaining-tokens": 19947, + "x-ratelimit-team-limit-requests": 50, + "x-ratelimit-team-remaining-requests": 47, + "x-ratelimit-team-limit-tokens": 40000, + "x-ratelimit-team-remaining-tokens": 39960, + }, + logger=logger, + ) + await _run_success_event( + { + "x-ratelimit-team-limit-requests": 50, + "x-ratelimit-team-remaining-requests": 46, + }, + logger=logger, + ) + + team_requests = ( + ("rate_limit_type", "requests"), + ("team", "team-id"), + ("team_alias", "team-alias"), + ) + assert _collected_samples("litellm_api_key_rate_limit_allowed_metric") == {} + assert _collected_samples("litellm_api_key_rate_limit_used_metric") == {} + assert _collected_samples("litellm_team_rate_limit_allowed_metric") == {team_requests: 50} + assert _collected_samples("litellm_team_rate_limit_used_metric") == {team_requests: 4} + finally: + _clear_prometheus_registry() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "additional_headers", + [ + None, + {"x-ratelimit-model_per_key-remaining-requests": 42}, + {"x-ratelimit-api_key-limit-requests": 10}, + {"x-ratelimit-api_key-limit-requests": "10", "x-ratelimit-api_key-remaining-requests": "7"}, + {"x-ratelimit-team-limit-tokens": True, "x-ratelimit-team-remaining-tokens": 5}, + ], +) +async def test_should_emit_no_key_or_team_rate_limit_series_without_a_complete_int_pair( + additional_headers, +): + _clear_prometheus_registry() + try: + await _run_success_event(additional_headers) + + for metric_name in KEY_AND_TEAM_RATE_LIMIT_METRICS: + assert _collected_samples(metric_name) == {}, metric_name + finally: + _clear_prometheus_registry() From 2b616fc479c873692829605c72999c99700c952b 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 01:04:53 +0000 Subject: [PATCH 47/97] feat(scim): add placeholder listing and merge so a shadowed account can be healed (#39231) Co-authored-by: yassin Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/models/user.py | 10 +- litellm/proxy/_lazy_openapi_snapshot.json | 179 +++++++++++++ .../management_endpoints/scim/scim_v2.py | 84 ++++++ litellm/repositories/user_repository.py | 26 +- .../proxy/management_endpoints/scim_v2.py | 6 + .../scim/test_scim_v2_endpoints.py | 240 ++++++++++++++++-- ui/litellm-dashboard/src/lib/http/schema.d.ts | 137 ++++++++++ 7 files changed, 659 insertions(+), 23 deletions(-) diff --git a/litellm/models/user.py b/litellm/models/user.py index 259c3440d87..82f78c28078 100644 --- a/litellm/models/user.py +++ b/litellm/models/user.py @@ -7,7 +7,7 @@ Canonical definition for ``litellm_usertable``. Re-exported from from datetime import datetime -from pydantic import ConfigDict, Field, model_validator +from pydantic import BaseModel, ConfigDict, Field, model_validator from litellm.models.object_permission import LiteLLM_ObjectPermissionTable from litellm.models.organization_membership import ( @@ -67,3 +67,11 @@ class LiteLLM_UserTable(LiteLLMPydanticObjectBase): if not self.models: return True return model_name in self.models + + +class SCIMPlaceholder(BaseModel): + """A user row keyed by a value that names another account by SSO identity or email.""" + + placeholder_user_id: str + resolved_user_ids: tuple[str, ...] + team_ids: tuple[str, ...] diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 5af45b29226..13c7a4c7cfa 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -32002,6 +32002,62 @@ "title": "SCIMPatchOperation", "type": "object" }, + "SCIMPlaceholder": { + "description": "A user row keyed by a value that names another account by SSO identity or email.", + "properties": { + "placeholder_user_id": { + "title": "Placeholder User Id", + "type": "string" + }, + "resolved_user_ids": { + "items": { + "type": "string" + }, + "title": "Resolved User Ids", + "type": "array" + }, + "team_ids": { + "items": { + "type": "string" + }, + "title": "Team Ids", + "type": "array" + } + }, + "required": [ + "placeholder_user_id", + "resolved_user_ids", + "team_ids" + ], + "title": "SCIMPlaceholder", + "type": "object" + }, + "SCIMPlaceholderMergeResult": { + "properties": { + "merged_into_user_id": { + "title": "Merged Into User Id", + "type": "string" + }, + "placeholder_user_id": { + "title": "Placeholder User Id", + "type": "string" + }, + "team_ids": { + "items": { + "type": "string" + }, + "title": "Team Ids", + "type": "array" + } + }, + "required": [ + "placeholder_user_id", + "merged_into_user_id", + "team_ids" + ], + "title": "SCIMPlaceholderMergeResult", + "type": "object" + }, "SCIMServiceProviderConfig": { "properties": { "authenticationSchemes": { @@ -33641,6 +33697,129 @@ "scim" ] } + }, + "/scim/v2/placeholders": { + "get": { + "description": "List user rows whose id is another account's SSO identity or email.\n\nAn earlier release provisioned a group member it could not match as a user keyed\nby the raw member value, and that row now shadows the account the value really\nnames, so every push of that member is refused. This lists those rows so an\noperator can fold each one into the account it shadows with\n``POST /scim/v2/placeholders/{user_id}/merge``. A row that has an SSO identity of\nits own or owns virtual keys is left out: someone uses that account.", + "operationId": "list_placeholders_scim_v2_placeholders_get", + "parameters": [ + { + "in": "query", + "name": "feature", + "required": false, + "schema": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Feature" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "items": { + "$ref": "#/components/schemas/SCIMPlaceholder" + }, + "title": "Response List Placeholders Scim V2 Placeholders Get", + "type": "array" + } + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "List Placeholders", + "tags": [ + "scim" + ] + } + }, + "/scim/v2/placeholders/{user_id}/merge": { + "post": { + "description": "Fold a placeholder user into the one account its id names by SSO identity or email.\n\nThe account is added to every team the placeholder is on, then the placeholder is\ndeleted the way ``DELETE /scim/v2/Users/{id}`` deletes a user, so the next group\npush resolves the member value to the real account. Refused with 409 when the row\nhas an SSO identity of its own, owns virtual keys, or names no account or several.", + "operationId": "merge_placeholder_scim_v2_placeholders__user_id__merge_post", + "parameters": [ + { + "in": "path", + "name": "user_id", + "required": true, + "schema": { + "title": "User ID", + "type": "string" + } + }, + { + "in": "query", + "name": "feature", + "required": false, + "schema": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Feature" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/SCIMPlaceholderMergeResult" + } + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Merge Placeholder", + "tags": [ + "scim" + ] + } } } }, diff --git a/litellm/proxy/management_endpoints/scim/scim_v2.py b/litellm/proxy/management_endpoints/scim/scim_v2.py index 8a0436f42dd..069f86c852c 100644 --- a/litellm/proxy/management_endpoints/scim/scim_v2.py +++ b/litellm/proxy/management_endpoints/scim/scim_v2.py @@ -29,6 +29,7 @@ import litellm from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid from litellm.litellm_core_utils.safe_json_dumps import safe_dumps +from litellm.models.user import SCIMPlaceholder from litellm.proxy._types import ( LiteLLM_TeamTable, LiteLLM_UserTable, @@ -1862,6 +1863,89 @@ async def delete_user( raise handle_exception_on_proxy(e) +@scim_router.get( + "/placeholders", + response_model=tuple[SCIMPlaceholder, ...], + dependencies=(Depends(user_api_key_auth),), +) +async def list_placeholders() -> tuple[SCIMPlaceholder, ...]: + """ + List user rows whose id is another account's SSO identity or email. + + An earlier release provisioned a group member it could not match as a user keyed + by the raw member value, and that row now shadows the account the value really + names, so every push of that member is refused. This lists those rows so an + operator can fold each one into the account it shadows with + ``POST /scim/v2/placeholders/{user_id}/merge``. A row that has an SSO identity of + its own or owns virtual keys is left out: someone uses that account. + """ + try: + prisma_client: Final = await _get_prisma_client_or_raise_exception() + async with prisma_client.tx() as tx: + return await UserRepository(prisma_client).find_shadowing_placeholders(tx) + except Exception as e: + raise handle_exception_on_proxy(e) + + +def _placeholder_rejection(placeholder: LiteLLM_UserTable, resolved: tuple[str, ...], key_count: int) -> str | None: + if placeholder.sso_user_id is not None: + return f"User '{placeholder.user_id}' has an SSO identity of its own, so it is an account someone signs in to" + if key_count: + return f"User '{placeholder.user_id}' owns {key_count} virtual keys. Move or delete them before merging it" + if not resolved: + return f"User '{placeholder.user_id}' shadows no account: no other user has that id as SSO identity or email" + if len(resolved) > 1: + return ( + f"User '{placeholder.user_id}' names {len(resolved)} accounts ({', '.join(resolved)}). Resolve that first" + ) + return None + + +@scim_router.post( + "/placeholders/{user_id}/merge", + response_model=SCIMPlaceholderMergeResult, + dependencies=(Depends(user_api_key_auth),), +) +async def merge_placeholder( + user_id: str = Path(..., title="User ID"), +) -> SCIMPlaceholderMergeResult: + """ + Fold a placeholder user into the one account its id names by SSO identity or email. + + The account is added to every team the placeholder is on, then the placeholder is + deleted the way ``DELETE /scim/v2/Users/{id}`` deletes a user, so the next group + push resolves the member value to the real account. Refused with 409 when the row + has an SSO identity of its own, owns virtual keys, or names no account or several. + """ + try: + prisma_client: Final = await _get_prisma_client_or_raise_exception() + placeholder: Final = await _check_user_exists(user_id) + resolved: Final = tuple( + other for other in await _users_named_by_member_value(user_id, prisma_client, take=None) if other != user_id + ) + owned_keys: Final[_UserIdWhere] = {"user_id": user_id} + keys: Final = await _table(VerificationTokenRepository(prisma_client)).find_many(where=owned_keys) + rejection: Final = _placeholder_rejection(placeholder, resolved, len(keys)) + if rejection is not None: + detail: Final[_ScimErrorDetail] = {"error": rejection} + raise HTTPException(status_code=409, detail=detail) + + target_user_id: Final = resolved[0] + team_ids: Final = tuple(placeholder.teams) + for team_id in team_ids: + await _add_user_to_team(user_id=target_user_id, team_id=team_id) + await delete_user(user_id=user_id) + await _recompute_scim_member_roles(prisma_client, (target_user_id,)) + verbose_proxy_logger.info( + "SCIM: merged placeholder user '%s' into '%s', moving teams %s", user_id, target_user_id, team_ids + ) + return SCIMPlaceholderMergeResult( + placeholder_user_id=user_id, merged_into_user_id=target_user_id, team_ids=team_ids + ) + except Exception as e: + raise handle_exception_on_proxy(e) + + def _parse_member_entry(entry: object) -> SCIMMember | None: """Parse one entry of a SCIM patch value, or None when it carries no id.""" if isinstance(entry, str): diff --git a/litellm/repositories/user_repository.py b/litellm/repositories/user_repository.py index 9df1bceac9c..87eb45f262d 100644 --- a/litellm/repositories/user_repository.py +++ b/litellm/repositories/user_repository.py @@ -6,15 +6,34 @@ import json from collections.abc import Mapping from typing import TYPE_CHECKING, Final -from litellm.models.user import LiteLLM_UserTable +from pydantic import TypeAdapter + +from litellm.models.user import LiteLLM_UserTable, SCIMPlaceholder from litellm.repositories.base_repository import BaseRepository, DbRecord, record_to_dict from litellm.repositories.prisma_protocols import TableActions if TYPE_CHECKING: + from prisma import Prisma from prisma import models as prisma_models _JSON_ENCODED_COLUMNS: Final = frozenset({"metadata", "model_spend", "model_max_budget"}) +_SHADOWING_PLACEHOLDERS_SQL: Final = """ +SELECT p.user_id AS placeholder_user_id, + array_agg(r.user_id ORDER BY r.user_id) AS resolved_user_ids, + p.teams AS team_ids +FROM "LiteLLM_UserTable" p +JOIN "LiteLLM_UserTable" r + ON r.user_id <> p.user_id + AND (r.sso_user_id = p.user_id OR LOWER(r.user_email) = LOWER(p.user_id)) +WHERE p.sso_user_id IS NULL + AND NOT EXISTS (SELECT 1 FROM "LiteLLM_VerificationToken" k WHERE k.user_id = p.user_id) +GROUP BY p.user_id, p.teams +ORDER BY p.user_id +""" + +_PLACEHOLDER_ROWS_ADAPTER: Final = TypeAdapter(tuple[SCIMPlaceholder, ...]) + class UserRepository(BaseRepository[LiteLLM_UserTable]): """Repository for user database operations.""" @@ -59,6 +78,11 @@ class UserRepository(BaseRepository[LiteLLM_UserTable]): """Find all users in a team.""" return await self.find_many(where={"teams": {"has": team_id}}) + async def find_shadowing_placeholders(self, tx: "Prisma") -> tuple[SCIMPlaceholder, ...]: + """Users with no SSO id and no virtual keys whose id is another user's SSO id or email.""" + rows: Final = await tx.query_raw(_SHADOWING_PLACEHOLDERS_SQL) + return _PLACEHOLDER_ROWS_ADAPTER.validate_python(rows) + async def count_billable_users(self) -> int: """Number of users that count toward the license seat limit. diff --git a/litellm/types/proxy/management_endpoints/scim_v2.py b/litellm/types/proxy/management_endpoints/scim_v2.py index 1612ea03817..7825684cfe5 100644 --- a/litellm/types/proxy/management_endpoints/scim_v2.py +++ b/litellm/types/proxy/management_endpoints/scim_v2.py @@ -150,6 +150,12 @@ class SCIMGroup(SCIMResource): members: list[SCIMMember] | None = None +class SCIMPlaceholderMergeResult(BaseModel): + placeholder_user_id: str + merged_into_user_id: str + team_ids: tuple[str, ...] + + # SCIM List Response Models class SCIMListResponse(BaseModel): schemas: list[str] = ["urn:ietf:params:scim:api:messages:2.0:ListResponse"] diff --git a/tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py b/tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py index d8fe22c4979..1697b77b99a 100644 --- a/tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py @@ -1,7 +1,8 @@ import logging import time -from collections.abc import Callable, Mapping +from collections.abc import Callable, Mapping, Sequence from itertools import chain +from types import MappingProxyType from typing import Final from unittest.mock import AsyncMock, MagicMock, call @@ -38,6 +39,7 @@ from litellm.proxy.management_endpoints.scim.scim_v2 import ( get_groups, get_users, get_service_provider_config, + merge_placeholder, patch_group, patch_team_membership, patch_user, @@ -52,6 +54,7 @@ from litellm.types.proxy.management_endpoints.scim_v2 import ( SCIMMember, SCIMPatchOp, SCIMPatchOperation, + SCIMPlaceholderMergeResult, SCIMServiceProviderConfig, SCIMUser, SCIMUserEmail, @@ -778,13 +781,17 @@ async def test_handle_existing_user_by_email_without_teams_preserves_memberships "litellm.proxy.management_endpoints.scim.scim_v2.ScimTransformations.transform_litellm_user_to_scim_user", AsyncMock(return_value=None), ) - mock_team_member_add = mocker.patch( # test-quality-ok: roster helpers are module-level, not injectable into the helper - "litellm.proxy.management_endpoints.scim.scim_v2.team_member_add", - AsyncMock(), + mock_team_member_add = ( + mocker.patch( # test-quality-ok: roster helpers are module-level, not injectable into the helper + "litellm.proxy.management_endpoints.scim.scim_v2.team_member_add", + AsyncMock(), + ) ) - mock_team_member_delete = mocker.patch( # test-quality-ok: roster helpers are module-level, not injectable into the helper - "litellm.proxy.management_endpoints.scim.scim_v2.team_member_delete", - AsyncMock(), + mock_team_member_delete = ( + mocker.patch( # test-quality-ok: roster helpers are module-level, not injectable into the helper + "litellm.proxy.management_endpoints.scim.scim_v2.team_member_delete", + AsyncMock(), + ) ) new_user_request = NewUserRequest( @@ -4470,9 +4477,11 @@ async def test_create_group_applies_default_team_params( "litellm.proxy.management_endpoints.scim.scim_v2._get_prisma_client_or_raise_exception", AsyncMock(return_value=_member_resolution_prisma(mocker, users=set(), teams=set())), ) - new_team_mock = mocker.patch( # test-quality-ok: endpoint collaborators are module-level, not injectable into create_group - "litellm.proxy.management_endpoints.scim.scim_v2.new_team", - AsyncMock(return_value=mocker.MagicMock()), + new_team_mock = ( + mocker.patch( # test-quality-ok: endpoint collaborators are module-level, not injectable into create_group + "litellm.proxy.management_endpoints.scim.scim_v2.new_team", + AsyncMock(return_value=mocker.MagicMock()), + ) ) mocker.patch( # test-quality-ok: endpoint collaborators are module-level, not injectable into create_group "litellm.proxy.management_endpoints.scim.scim_v2.ScimTransformations.transform_litellm_team_to_scim_group", @@ -4927,9 +4936,7 @@ async def test_process_group_patch_remove_by_the_id_the_directory_added_with( @pytest.mark.asyncio -async def test_process_group_patch_remove_still_drops_a_placeholder_by_its_literal_id( - mocker, scim_upsert_user_enabled -): +async def test_process_group_patch_remove_still_drops_a_placeholder_by_its_literal_id(mocker, scim_upsert_user_enabled): """An earlier release put unmatched ids on the roster verbatim, so a remove has to keep clearing the id as written even once it also resolves.""" patch_ops = SCIMPatchOp( @@ -4940,7 +4947,10 @@ async def test_process_group_patch_remove_still_drops_a_placeholder_by_its_liter team_id="parent-group", team_alias="Parent Group", members=[], - members_with_roles=[Member(user_id="legacy@example.com", role="user"), Member(user_id="keep-user", role="user")], + members_with_roles=[ + Member(user_id="legacy@example.com", role="user"), + Member(user_id="keep-user", role="user"), + ], ) _, final_members, _ = await _process_group_patch_operations( @@ -5105,11 +5115,8 @@ async def test_process_group_patch_remove_refuses_when_two_members_share_the_id( assert "more than one member of this group" in str(exc_info.value.detail) - @pytest.mark.asyncio -async def test_resolve_group_member_ids_exact_user_id_wins_when_it_names_nobody_else( - mocker, scim_upsert_user_enabled -): +async def test_resolve_group_member_ids_exact_user_id_wins_when_it_names_nobody_else(mocker, scim_upsert_user_enabled): """The canonical user id stays authoritative, including when the same account also holds that value as its email, which is how a SCIM-provisioned account is keyed.""" prisma_client = _member_resolution_prisma( @@ -5171,9 +5178,7 @@ async def test_resolve_group_member_ids_refuses_a_user_id_that_names_another_acc assert exc_info.value.status_code == 400 assert "member-id" in str(exc_info.value.detail) create_user_mock.assert_not_called() - assert any( - record.levelno >= logging.WARNING and "someone-else" in record.getMessage() for record in caplog.records - ) + assert any(record.levelno >= logging.WARNING and "someone-else" in record.getMessage() for record in caplog.records) @pytest.mark.asyncio @@ -5711,3 +5716,196 @@ async def test_patch_group_404s_when_team_deleted_mid_request(mocker): assert exc_info.value.code == "404" assert f"Group not found with ID: {group_id}" in exc_info.value.message + + +_SHADOW_MEMBER_VALUE: Final = "00u1shadow" +_SHADOWED_ACCOUNT: Final = "real-1" +_SHADOWED_GROUP: Final = "grp-eng" + + +def _shadowed_tenant_rows() -> tuple[LiteLLM_UserTable, ...]: + """A placeholder keyed by the raw member value, and the real account that value names by SSO id.""" + return ( + LiteLLM_UserTable(user_id=_SHADOW_MEMBER_VALUE, user_email=_SHADOW_MEMBER_VALUE, teams=[_SHADOWED_GROUP]), + LiteLLM_UserTable(user_id=_SHADOWED_ACCOUNT, user_email="alice@example.com", sso_user_id=_SHADOW_MEMBER_VALUE), + ) + + +def _shadow_tenant_prisma( + mocker: MockerFixture, + *, + rows: Sequence[LiteLLM_UserTable], + keys_owned_by: Mapping[str, int] = MappingProxyType({}), +) -> MagicMock: + """Prisma fake whose user rows are live: deleting one removes it from every later lookup.""" + users: Final[dict[str, LiteLLM_UserTable]] = {row.user_id: row for row in rows} + team: Final = LiteLLM_TeamTable( + team_id=_SHADOWED_GROUP, + members=[_SHADOW_MEMBER_VALUE], + members_with_roles=[Member(user_id=_SHADOW_MEMBER_VALUE, role="user")], + metadata={SCIM_MANAGED_TEAM_METADATA_KEY: True}, + ) + + async def find_unique(where: Mapping[str, str]) -> LiteLLM_UserTable | None: + return users.get(where["user_id"]) + + def clause_matches(row: LiteLLM_UserTable, clause: Mapping[str, object]) -> bool: + if "user_id" in clause: + return row.user_id == clause["user_id"] + if "sso_user_id" in clause: + return row.sso_user_id == clause["sso_user_id"] + email_filter: Final = clause["user_email"] + assert isinstance(email_filter, dict) + return (row.user_email or "").casefold() == str(email_filter["equals"]).casefold() + + async def identity_rows(where: Mapping[str, object], take: int | None = None) -> tuple[LiteLLM_UserTable, ...]: + clauses: Final = where["OR"] + assert isinstance(clauses, list) + matched: Final = tuple(row for row in users.values() if any(clause_matches(row, clause) for clause in clauses)) + return matched[:take] if take else matched + + async def delete(where: Mapping[str, str]) -> LiteLLM_UserTable | None: + return users.pop(where["user_id"], None) + + async def keys_for(where: Mapping[str, object]) -> tuple[MagicMock, ...]: + return tuple(mocker.MagicMock() for _ in range(keys_owned_by.get(str(where["user_id"]), 0))) + + async def team_lookup(where: Mapping[str, str]) -> LiteLLM_TeamTable | None: + return team if where["team_id"] == team.team_id else None + + prisma_client = mocker.MagicMock() + prisma_client.db = mocker.MagicMock() + prisma_client.db.litellm_usertable = mocker.MagicMock() + prisma_client.db.litellm_usertable.find_unique = AsyncMock(side_effect=find_unique) + prisma_client.db.litellm_usertable.find_many = AsyncMock(side_effect=identity_rows) + prisma_client.db.litellm_usertable.delete = AsyncMock(side_effect=delete) + prisma_client.db.litellm_teamtable = mocker.MagicMock() + prisma_client.db.litellm_teamtable.find_unique = AsyncMock(side_effect=team_lookup) + prisma_client.db.litellm_teamtable.update = AsyncMock(return_value=team) + prisma_client.db.litellm_verificationtoken = mocker.MagicMock() + prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(side_effect=keys_for) + prisma_client.db.litellm_invitationlink = mocker.MagicMock(delete_many=AsyncMock(return_value=0)) + prisma_client.db.litellm_organizationmembership = mocker.MagicMock(delete_many=AsyncMock(return_value=0)) + prisma_client.db.litellm_teammembership = mocker.MagicMock(delete_many=AsyncMock(return_value=0)) + return prisma_client + + +@pytest.fixture +def shadowed_tenant(mocker, monkeypatch, scim_upsert_user_enabled) -> MagicMock: + from litellm.proxy import proxy_server + + prisma_client: Final = _shadow_tenant_prisma(mocker, rows=_shadowed_tenant_rows()) + monkeypatch.setattr(proxy_server, "prisma_client", prisma_client) + return prisma_client + + +async def _push_shadow_member(prisma_client: MagicMock): + return await _resolve_group_member_ids( + members=[SCIMMember(value=_SHADOW_MEMBER_VALUE)], + created_via="scim_group_membership", + prisma_client=prisma_client, + ) + + +@pytest.mark.asyncio +async def test_merge_placeholder_hands_the_group_to_the_shadowed_account(mocker, shadowed_tenant): + """Every group push of the shadowing value is refused until the placeholder is folded into + the real account; after the merge the same push resolves to that account.""" + team_member_add_mock = ( + mocker.patch( # test-quality-ok: roster helpers are module-level, not injectable into the endpoint + "litellm.proxy.management_endpoints.scim.scim_v2.team_member_add", AsyncMock() + ) + ) + team_member_delete_mock = ( + mocker.patch( # test-quality-ok: roster helpers are module-level, not injectable into the endpoint + "litellm.proxy.management_endpoints.scim.scim_v2.team_member_delete", AsyncMock() + ) + ) + + with pytest.raises(HTTPException) as before: + await _push_shadow_member(shadowed_tenant) + assert before.value.status_code == 400 + + result: Final = await merge_placeholder(user_id=_SHADOW_MEMBER_VALUE) + + assert result == SCIMPlaceholderMergeResult( + placeholder_user_id=_SHADOW_MEMBER_VALUE, + merged_into_user_id=_SHADOWED_ACCOUNT, + team_ids=(_SHADOWED_GROUP,), + ) + added: Final = team_member_add_mock.call_args.kwargs["data"] + assert (added.team_id, added.member.user_id) == (_SHADOWED_GROUP, _SHADOWED_ACCOUNT) + dropped: Final = team_member_delete_mock.call_args.kwargs["data"] + assert (dropped.team_id, dropped.user_id) == (_SHADOWED_GROUP, _SHADOW_MEMBER_VALUE) + shadowed_tenant.db.litellm_teammembership.delete_many.assert_awaited_once_with( + where={"user_id": _SHADOW_MEMBER_VALUE} + ) + shadowed_tenant.db.litellm_usertable.delete.assert_awaited_once_with(where={"user_id": _SHADOW_MEMBER_VALUE}) + + after: Final = await _push_shadow_member(shadowed_tenant) + assert after.all_member_ids == [_SHADOWED_ACCOUNT] + assert after.created_users == [] + + +@pytest.mark.asyncio +async def test_merge_placeholder_keeps_the_placeholder_when_the_roster_write_fails(mocker, shadowed_tenant): + """If the real account cannot join the team, the placeholder stays on it, or the membership is gone + from both accounts.""" + mocker.patch( # test-quality-ok: roster helpers are module-level, not injectable into the endpoint + "litellm.proxy.management_endpoints.scim.scim_v2.team_member_add", + AsyncMock(side_effect=Exception("database connection lost")), + ) + team_member_delete_mock = ( + mocker.patch( # test-quality-ok: roster helpers are module-level, not injectable into the endpoint + "litellm.proxy.management_endpoints.scim.scim_v2.team_member_delete", AsyncMock() + ) + ) + + with pytest.raises(ProxyException): + await merge_placeholder(user_id=_SHADOW_MEMBER_VALUE) + + team_member_delete_mock.assert_not_awaited() + shadowed_tenant.db.litellm_usertable.delete.assert_not_awaited() + assert await shadowed_tenant.db.litellm_usertable.find_unique(where={"user_id": _SHADOW_MEMBER_VALUE}) is not None + + +@pytest.mark.parametrize( + ("rows", "keys_owned_by", "merged", "reason"), + [ + pytest.param(_shadowed_tenant_rows(), {}, _SHADOWED_ACCOUNT, "SSO identity of its own", id="real-account"), + pytest.param( + _shadowed_tenant_rows(), {_SHADOW_MEMBER_VALUE: 2}, _SHADOW_MEMBER_VALUE, "2 virtual keys", id="owns-keys" + ), + pytest.param(_shadowed_tenant_rows()[:1], {}, _SHADOW_MEMBER_VALUE, "shadows no account", id="names-nobody"), + pytest.param( + (*_shadowed_tenant_rows(), LiteLLM_UserTable(user_id="real-2", user_email=_SHADOW_MEMBER_VALUE.upper())), + {}, + _SHADOW_MEMBER_VALUE, + "names 2 accounts (real-1, real-2)", + id="names-two-accounts", + ), + ], +) +@pytest.mark.asyncio +async def test_merge_placeholder_refuses_rows_that_are_not_a_lone_placeholder( + mocker, monkeypatch, scim_upsert_user_enabled, rows, keys_owned_by, merged, reason +): + """Only a row with no SSO identity and no keys whose id names exactly one other account is folded; + anything else could move memberships to the wrong person, so nothing is written.""" + from litellm.proxy import proxy_server + + prisma_client: Final = _shadow_tenant_prisma(mocker, rows=rows, keys_owned_by=keys_owned_by) + monkeypatch.setattr(proxy_server, "prisma_client", prisma_client) + team_member_add_mock = ( + mocker.patch( # test-quality-ok: roster helpers are module-level, not injectable into the endpoint + "litellm.proxy.management_endpoints.scim.scim_v2.team_member_add", AsyncMock() + ) + ) + + with pytest.raises(ProxyException) as exc_info: + await merge_placeholder(user_id=merged) + + assert int(exc_info.value.code) == 409 + assert reason in str(exc_info.value.message) + team_member_add_mock.assert_not_awaited() + prisma_client.db.litellm_usertable.delete.assert_not_awaited() diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 6e38f3fa15e..6f044fec3f3 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -13289,6 +13289,58 @@ export interface paths { patch: operations["patch_user_scim_v2_Users__user_id__patch"]; trace?: never; }; + "/scim/v2/placeholders": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * List Placeholders + * @description List user rows whose id is another account's SSO identity or email. + * + * An earlier release provisioned a group member it could not match as a user keyed + * by the raw member value, and that row now shadows the account the value really + * names, so every push of that member is refused. This lists those rows so an + * operator can fold each one into the account it shadows with + * ``POST /scim/v2/placeholders/{user_id}/merge``. A row that has an SSO identity of + * its own or owns virtual keys is left out: someone uses that account. + */ + get: operations["list_placeholders_scim_v2_placeholders_get"]; + put?: never; + post?: never; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; + "/scim/v2/placeholders/{user_id}/merge": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + get?: never; + put?: never; + /** + * Merge Placeholder + * @description Fold a placeholder user into the one account its id names by SSO identity or email. + * + * The account is added to every team the placeholder is on, then the placeholder is + * deleted the way ``DELETE /scim/v2/Users/{id}`` deletes a user, so the next group + * push resolves the member value to the real account. Refused with 409 when the row + * has an SSO identity of its own, owns virtual keys, or names no account or several. + */ + post: operations["merge_placeholder_scim_v2_placeholders__user_id__merge_post"]; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; "/search": { parameters: { query?: never; @@ -34927,6 +34979,27 @@ export interface components { /** Value */ value?: unknown | null; }; + /** + * SCIMPlaceholder + * @description A user row keyed by a value that names another account by SSO identity or email. + */ + SCIMPlaceholder: { + /** Placeholder User Id */ + placeholder_user_id: string; + /** Resolved User Ids */ + resolved_user_ids: string[]; + /** Team Ids */ + team_ids: string[]; + }; + /** SCIMPlaceholderMergeResult */ + SCIMPlaceholderMergeResult: { + /** Merged Into User Id */ + merged_into_user_id: string; + /** Placeholder User Id */ + placeholder_user_id: string; + /** Team Ids */ + team_ids: string[]; + }; /** SCIMServiceProviderConfig */ SCIMServiceProviderConfig: { /** Authenticationschemes */ @@ -55858,6 +55931,70 @@ export interface operations { }; }; }; + list_placeholders_scim_v2_placeholders_get: { + parameters: { + query?: { + feature?: string | null; + }; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["SCIMPlaceholder"][]; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; + merge_placeholder_scim_v2_placeholders__user_id__merge_post: { + parameters: { + query?: { + feature?: string | null; + }; + header?: never; + path: { + user_id: string; + }; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["SCIMPlaceholderMergeResult"]; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; search_search_post: { parameters: { query?: { From 4b87fd5718ed96e7080c68695750b90234194587 Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Tue, 1 Sep 2026 18:06:35 -0700 Subject: [PATCH 48/97] fix: normalize provider-specific cache token fields in OTel v2 usage (#39202) * fix: normalize provider-specific cache token fields in OTel v2 usage * fix: use an immutable empty mapping for the cache token details fallback * fix: ignore malformed cache token values instead of emitting or raising --- litellm/integrations/otel/model/payloads.py | 43 +++++++++++- .../otel/test_otel_v2_sources_of_truth.py | 68 +++++++++++++++++++ 2 files changed, 109 insertions(+), 2 deletions(-) diff --git a/litellm/integrations/otel/model/payloads.py b/litellm/integrations/otel/model/payloads.py index d35405538f6..e8ed269f6cb 100644 --- a/litellm/integrations/otel/model/payloads.py +++ b/litellm/integrations/otel/model/payloads.py @@ -6,6 +6,7 @@ import json from collections.abc import Mapping from dataclasses import dataclass, field from enum import Enum +from types import MappingProxyType from typing import TYPE_CHECKING, ClassVar, Final, cast from urllib.parse import urlsplit @@ -62,6 +63,31 @@ if TYPE_CHECKING: # --- typed sub-structures ---------------------------------------------------- # +def _cache_token_value(*values: object) -> int | None: + explicit_zero = False + invalid_before_zero = False + for raw_value in values: + if raw_value is None: + continue + if isinstance(raw_value, bool): + parsed = None + else: + try: + parsed = as_int(raw_value) + except (OverflowError, ValueError): + parsed = None + if parsed is None: + if not explicit_zero: + invalid_before_zero = True + elif parsed > 0: + return parsed + elif parsed == 0: + explicit_zero = True + elif not explicit_zero: + invalid_before_zero = True + return 0 if explicit_zero and not invalid_before_zero else None + + @dataclass(frozen=True) class LLMRequestParams: temperature: float | None = None @@ -104,12 +130,25 @@ class LLMUsage: metadata: Final[Mapping[str, object]] = payload.get("metadata") or {} raw_usage: Final = metadata.get("usage_object") usage_object: Final[Mapping[str, object]] = raw_usage if isinstance(raw_usage, Mapping) else {} + raw_details: Final = usage_object.get("prompt_tokens_details") + prompt_details: Final[Mapping[str, object]] = ( + raw_details if isinstance(raw_details, Mapping) else MappingProxyType({}) + ) return cls( input_tokens=as_int(payload.get("prompt_tokens")), output_tokens=as_int(payload.get("completion_tokens")), total_tokens=as_int(payload.get("total_tokens")), - cache_creation_input_tokens=as_int(usage_object.get("cache_creation_input_tokens")), - cache_read_input_tokens=as_int(usage_object.get("cache_read_input_tokens")), + cache_creation_input_tokens=_cache_token_value( + usage_object.get("cache_creation_input_tokens"), + prompt_details.get("cache_write_tokens"), + prompt_details.get("cache_creation_tokens"), + prompt_details.get("cache_creation_input_tokens"), + ), + cache_read_input_tokens=_cache_token_value( + usage_object.get("cache_read_input_tokens"), + prompt_details.get("cached_tokens"), + usage_object.get("prompt_cache_hit_tokens"), + ), ) diff --git a/tests/test_litellm/integrations/otel/test_otel_v2_sources_of_truth.py b/tests/test_litellm/integrations/otel/test_otel_v2_sources_of_truth.py index ca628aa3405..99d706a9c44 100644 --- a/tests/test_litellm/integrations/otel/test_otel_v2_sources_of_truth.py +++ b/tests/test_litellm/integrations/otel/test_otel_v2_sources_of_truth.py @@ -541,6 +541,74 @@ def test_llm_call_adapter_extracts_cache_tokens_from_usage_object(): assert data.usage.cache_read_input_tokens == 3 +def test_llm_call_adapter_normalizes_nested_cache_tokens(): + cases: Final = ( + ({"prompt_tokens_details": {"cached_tokens": 3}}, 3, None), + ({"prompt_cache_hit_tokens": 11}, 11, None), + ({"prompt_tokens_details": {"cache_write_tokens": 7}}, None, 7), + ({"prompt_tokens_details": {"cache_creation_tokens": 13}}, None, 13), + ({"prompt_tokens_details": {"cache_creation_input_tokens": 17}}, None, 17), + ) + for usage_object, expected_read, expected_creation in cases: + case_payload = _sample_payload(metadata={"usage_object": usage_object}) + data = LLMCallSpanData.from_standard_logging_payload(case_payload) + assert data.usage.cache_read_input_tokens == expected_read + assert data.usage.cache_creation_input_tokens == expected_creation + + +def test_llm_call_adapter_prefers_nested_count_over_zero_top_level(): + payload = _sample_payload( + metadata={ + "usage_object": { + "cache_read_input_tokens": 0, + "cache_creation_input_tokens": 0, + "prompt_tokens_details": {"cached_tokens": 5, "cache_write_tokens": 7}, + } + } + ) + data = LLMCallSpanData.from_standard_logging_payload(payload) + assert data.usage.cache_read_input_tokens == 5 + assert data.usage.cache_creation_input_tokens == 7 + + +def test_llm_call_adapter_ignores_invalid_cache_values_before_valid_fallbacks(): + payload = _sample_payload( + metadata={ + "usage_object": { + "cache_read_input_tokens": -1, + "cache_creation_input_tokens": "5.0", + "prompt_tokens_details": {"cached_tokens": 5, "cache_write_tokens": 7}, + } + } + ) + data = LLMCallSpanData.from_standard_logging_payload(payload) + assert data.usage.cache_read_input_tokens == 5 + assert data.usage.cache_creation_input_tokens == 7 + + +def test_llm_call_adapter_ignores_non_finite_cache_values(): + payload = _sample_payload( + metadata={ + "usage_object": { + "prompt_tokens_details": {"cached_tokens": float("nan")}, + } + } + ) + data = LLMCallSpanData.from_standard_logging_payload(payload) + assert data.usage.cache_read_input_tokens is None + + +def test_llm_call_adapter_preserves_explicit_zero_and_omits_missing_cache_tokens(): + for usage_object, expected_read, expected_creation in ( + ({"prompt_tokens_details": {"cached_tokens": 0}}, 0, None), + ({}, None, None), + ): + case_payload = _sample_payload(metadata={"usage_object": usage_object}) + data = LLMCallSpanData.from_standard_logging_payload(case_payload) + assert data.usage.cache_read_input_tokens == expected_read + assert data.usage.cache_creation_input_tokens == expected_creation + + def test_llm_call_adapter_cache_tokens_none_without_usage_object(): data = LLMCallSpanData.from_standard_logging_payload(_sample_payload()) assert data.usage.cache_creation_input_tokens is None From 62f032cca53c16967af6a424facc5dd88205d4dd Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Tue, 1 Sep 2026 18:07:04 -0700 Subject: [PATCH 49/97] fix(proxy): keep passthrough logging metadata and model_info dicts when team callbacks are wired (#39216) * fix(proxy): keep passthrough logging metadata and model_info dicts when team callbacks are wired Passing team callback vars into Logging(kwargs=...) makes get_litellm_params materialize a full litellm_params, where metadata and model_info default to None instead of being absent. Readers that resolve them as .get(key, {}).get(...) then raise, so any passthrough request from a team with logging callbacks 500s once a pre-call guardrail is on, and the router strategy loggers log a traceback per request. * test(proxy): annotate the closure dicts the passthrough logging tests record into --- .../pass_through_endpoints.py | 2 + .../test_pass_through_endpoints.py | 111 +++++++++++++++++- 2 files changed, 111 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index ff306eb65f1..79d5d0a016f 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -812,6 +812,8 @@ def _resolve_team_callback_wiring( else { # mutable-ok: Logging arg **callback_vars, TRUSTED_CALLBACK_VARS_FIELD: callback_vars, + "metadata": {}, # mutable-ok: Logging arg + "model_info": {}, # mutable-ok: Logging arg } ) return _TeamCallbackWiring( diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py index f5ae0fe5977..d3f17c73499 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -2,6 +2,7 @@ import asyncio import json import logging import os +from collections.abc import Callable from contextlib import ExitStack, contextmanager from io import BytesIO from types import SimpleNamespace @@ -29,6 +30,7 @@ from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( websocket_passthrough_request, ) from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.proxy._types import ProxyException, UserAPIKeyAuth from litellm.types.passthrough_endpoints.pass_through_endpoints import ( LITELLM_PASS_THROUGH_RAW_BODY_STATE_KEY, @@ -5464,7 +5466,10 @@ def test_the_marker_check_distinguishes_the_two_route_kinds(): assert request_dispatched_to_pass_through_endpoint(builtin) is False -async def _drive_passthrough_request_and_capture_logging(user_api_key_dict: UserAPIKeyAuth) -> tuple[int, object]: +async def _drive_passthrough_request_and_capture_logging( + user_api_key_dict: UserAPIKeyAuth, + on_pre_call: Callable[[LiteLLMLoggingObj | None], None] | None = None, +) -> tuple[int, LiteLLMLoggingObj | None]: import litellm from litellm.llms.custom_httpx.http_handler import get_async_httpx_client from litellm.types.llms.custom_http import httpxSpecialProvider @@ -5487,10 +5492,12 @@ async def _drive_passthrough_request_and_capture_logging(user_api_key_dict: User mock_request.query_params = QueryParams({}) mock_request.body = AsyncMock(return_value=b'{"model": "gemini-2.0-flash"}') - captured_data: dict = {} + captured_data: dict = {} # mutable-ok: the pre-call hook records the request data into it async def capture_pre_call_hook(user_api_key_dict, data, call_type): captured_data.update(data) + if on_pre_call is not None: + on_pre_call(data.get("litellm_logging_obj")) return data mock_proxy_logging = MagicMock() @@ -5623,3 +5630,103 @@ async def test_resolve_team_callback_wiring_fails_open_on_operational_error(): assert wiring.success_callbacks is None assert wiring.failure_callbacks is None assert wiring.logging_kwargs is None + + +@pytest.mark.asyncio +async def test_pass_through_request_leaves_guardrail_readable_metadata(): + """A pre-call guardrail reads the request headers off the passthrough logging + params without raising.""" + from litellm.proxy.guardrails.guardrail_hooks.hiddenlayer.hiddenlayer import ( + _logged_request_headers, + ) + + user_api_key_dict = UserAPIKeyAuth( + api_key="test-key", + team_id="test-team", + team_metadata={ + "logging": [ + { + "callback_name": "langfuse", + "callback_type": "success_and_failure", + "callback_vars": { + "langfuse_public_key": "pk_test", + "langfuse_secret_key": "sk_test", + }, + } + ] + }, + ) + + observed: dict[str, dict[str, str] | BaseException] = {} # mutable-ok: the pre-call hook records into it + + def read_headers_the_way_a_guardrail_does(logging_obj: LiteLLMLoggingObj | None) -> None: + assert logging_obj is not None + try: + observed["headers"] = _logged_request_headers(logging_obj) + except Exception as exc: # noqa: BLE001 - the regression is that this used to raise + observed["headers"] = exc + + status_code, logging_obj = await _drive_passthrough_request_and_capture_logging( + user_api_key_dict, on_pre_call=read_headers_the_way_a_guardrail_does + ) + + assert "headers" in observed, "the pre-call hook never ran, so nothing was observed" + assert observed["headers"] == {}, f"guardrail header read failed: {observed['headers']!r}" + assert status_code == 200 + assert logging_obj is not None + assert logging_obj.dynamic_success_callbacks, "team success callbacks must stay wired" + assert logging_obj.standard_callback_dynamic_params.get("langfuse_public_key") == "pk_test" + + +@pytest.mark.asyncio +async def test_pass_through_request_leaves_cost_router_logger_working(): + """The cost router's logger reads the deployment id off the passthrough logging + params without raising. least_busy shares the read but swallows the exception, + so this is the strategy where the break is observable.""" + from litellm._logging import verbose_logger + from litellm.caching.caching import DualCache + from litellm.router_strategy.lowest_cost import LowestCostLoggingHandler + + handler = LowestCostLoggingHandler(router_cache=DualCache()) + + user_api_key_dict = UserAPIKeyAuth( + api_key="test-key", + team_id="test-team", + team_metadata={ + "logging": [ + { + "callback_name": "langfuse", + "callback_type": "success_and_failure", + "callback_vars": { + "langfuse_public_key": "pk_test", + "langfuse_secret_key": "sk_test", + }, + } + ] + }, + ) + + status_code, logging_obj = await _drive_passthrough_request_and_capture_logging(user_api_key_dict) + assert status_code == 200 + assert logging_obj is not None + + raised: list[logging.LogRecord] = [] # mutable-ok: logging.Handler records into it + + class _RecordTracebacks(logging.Handler): + def emit(self, record: logging.LogRecord) -> None: + if record.exc_info is not None: + raised.append(record) + + recorder = _RecordTracebacks() + verbose_logger.addHandler(recorder) + try: + await handler.async_log_success_event( + kwargs=logging_obj.model_call_details, + response_obj=None, + start_time=None, + end_time=None, + ) + finally: + verbose_logger.removeHandler(recorder) + + assert not raised, f"cost router logger raised on the passthrough logging params: {raised[0].exc_info}" 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 50/97] 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 51/97] 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 52/97] 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 53/97] 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 54/97] 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 55/97] 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 56/97] 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 57/97] 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 58/97] 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 59/97] 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 e2c3f51c46aa2a11a930b6c19bc28e4144a38a1e Mon Sep 17 00:00:00 2001 From: James Liounis Date: Wed, 2 Sep 2026 00:46:46 -0400 Subject: [PATCH 60/97] fix(search): forward search-tool params through the router, complete Parallel AI v1 param mapping (#37883) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix(search): forward search-tool params through the router, complete Parallel AI v1 param mapping SearchAPIRouter dropped every parameter configured on a search tool, forwarding only per-request kwargs. Any tool-level setting (mode, max_results, ...) was silently lost on the way to the adapter, for every search provider. Also completes the Parallel AI v1 search surface: after_date, fetch_policy, location and include_domains now nest under advanced_settings instead of being sent as unknown top-level fields, responses preserve search_id / session_id / warnings / raw excerpts, and search cost is derived from the request mode and the provider's reported usage rather than a single flat rate. * fix(parallel_ai): stop a caller from pricing its own search request `_parallel_ai_usage` carries the provider's reported usage into cost calculation. It was only written when the response contained a usage block, so a caller could pass `_parallel_ai_usage=[{"name": "sku_search", "count": 0}]` and, whenever the provider omitted usage, bill $0.00 instead of $0.005 — the value also reached the upstream request body as an unknown field. The key is now stripped from inbound params and written unconditionally from the parsed response, so only the provider can populate it. * fix(parallel_ai): price fast search mode correctly * test(parallel_ai): fake search at HTTP boundary * fix(parallel_ai): tolerate null search result fields --------- Co-authored-by: khushishelat --- .../parallel_ai/search/cost_calculator.py | 90 +++++++++ .../llms/parallel_ai/search/transformation.py | 143 +++++++++---- ...odel_prices_and_context_window_backup.json | 14 +- litellm/router_utils/search_api_router.py | 12 +- litellm/search/cost_calculator.py | 35 +++- model_prices_and_context_window.json | 14 +- .../test_router_helper_utils.py | 3 + .../parallel_ai/test_parallel_ai_search.py | 190 +++++++++++++++-- .../test_parallel_ai_search_gateway.py | 191 ++++++++++++++++++ 9 files changed, 637 insertions(+), 55 deletions(-) create mode 100644 litellm/llms/parallel_ai/search/cost_calculator.py create mode 100644 tests/test_litellm/llms/parallel_ai/test_parallel_ai_search_gateway.py diff --git a/litellm/llms/parallel_ai/search/cost_calculator.py b/litellm/llms/parallel_ai/search/cost_calculator.py new file mode 100644 index 00000000000..809cd280cc8 --- /dev/null +++ b/litellm/llms/parallel_ai/search/cost_calculator.py @@ -0,0 +1,90 @@ +from collections.abc import Mapping, Sequence +from types import MappingProxyType +from typing import Final + +from pydantic import TypeAdapter, ValidationError + +from litellm.utils import get_model_info + +PARALLEL_AI_DEFAULT_RESULTS: Final = 10 +PARALLEL_AI_ADDITIONAL_RESULT_COST: Final = 0.001 +PARALLEL_AI_USAGE_PARAM: Final = "_parallel_ai_usage" +PARALLEL_AI_STANDARD_SEARCH_MODEL: Final = "parallel_ai/search" +PARALLEL_AI_FAST_SEARCH_MODEL: Final = "parallel_ai/search-fast" +PARALLEL_AI_TURBO_SEARCH_MODEL: Final = "parallel_ai/search-turbo" +PARALLEL_AI_PRICING_MODEL_BY_MODE: Final[Mapping[str, str]] = MappingProxyType( + { + "fast": PARALLEL_AI_FAST_SEARCH_MODEL, + "turbo": PARALLEL_AI_TURBO_SEARCH_MODEL, + } +) +ADVANCED_SETTINGS_ADAPTER: Final[TypeAdapter[Mapping[str, object]]] = TypeAdapter(Mapping[str, object]) + + +def _non_negative_int(value: object) -> int | None: + if isinstance(value, bool) or not isinstance(value, int) or value < 0: + return None + return value + + +def _usage_count(usage: Sequence[Mapping[str, object]], sku: str) -> int | None: + counts: Final = tuple( + count + for item in usage + if item.get("name") == sku + if (count := _non_negative_int(item.get("count"))) is not None + ) + return sum(counts) if counts else None + + +def _effective_mode(optional_params: Mapping[str, object]) -> str: + mode: Final = optional_params.get("mode") + if isinstance(mode, str): + return mode + + processor: Final = optional_params.get("processor") + if processor == "pro": + return "advanced" + return "basic" + + +def _effective_max_results(optional_params: Mapping[str, object]) -> int: + try: + advanced_settings: Final = ADVANCED_SETTINGS_ADAPTER.validate_python(optional_params.get("advanced_settings")) + advanced_max_results: Final = _non_negative_int(advanced_settings.get("max_results")) + if advanced_max_results is not None: + return advanced_max_results + except ValidationError: + pass + + max_results: Final = _non_negative_int(optional_params.get("max_results")) + return max_results if max_results is not None else PARALLEL_AI_DEFAULT_RESULTS + + +def _request_cost(mode: str) -> float: + pricing_model: Final = PARALLEL_AI_PRICING_MODEL_BY_MODE.get(mode, PARALLEL_AI_STANDARD_SEARCH_MODEL) + model_info: Final = get_model_info(model=pricing_model, custom_llm_provider="parallel_ai") + return float(model_info.get("input_cost_per_query") or 0.0) + + +def _additional_results( + optional_params: Mapping[str, object], + usage: Sequence[Mapping[str, object]] | None, +) -> int: + usage_count: Final = _usage_count(usage, "sku_search_additional_results") if usage is not None else None + if usage_count is not None: + return usage_count + if usage is not None: + return 0 + return max(_effective_max_results(optional_params) - PARALLEL_AI_DEFAULT_RESULTS, 0) + + +def parallel_ai_search_cost( + optional_params: Mapping[str, object], + usage: Sequence[Mapping[str, object]] | None, +) -> float: + request_cost: Final = _request_cost(_effective_mode(optional_params)) + request_count_from_usage: Final = _usage_count(usage, "sku_search") if usage is not None else None + request_count: Final = request_count_from_usage if request_count_from_usage is not None else 1 + additional_results: Final = _additional_results(optional_params, usage) + return request_count * request_cost + additional_results * PARALLEL_AI_ADDITIONAL_RESULT_COST diff --git a/litellm/llms/parallel_ai/search/transformation.py b/litellm/llms/parallel_ai/search/transformation.py index ea21d1153fe..bde7b7b86db 100644 --- a/litellm/llms/parallel_ai/search/transformation.py +++ b/litellm/llms/parallel_ai/search/transformation.py @@ -4,9 +4,13 @@ Calls Parallel AI's /v1/search endpoint to search the web. Parallel AI API Reference: https://docs.parallel.ai/api-reference/search/search """ +from collections.abc import Mapping, Sequence +from types import MappingProxyType from typing import Final, TypedDict import httpx +from pydantic import BaseModel, ConfigDict +from typing_extensions import ReadOnly from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.base_llm.search.transformation import ( @@ -14,9 +18,29 @@ from litellm.llms.base_llm.search.transformation import ( SearchResponse, SearchResult, ) +from litellm.llms.parallel_ai.search.cost_calculator import PARALLEL_AI_USAGE_PARAM from litellm.secret_managers.main import get_secret_str +class _ParallelAIV1SearchResult(BaseModel): + model_config = ConfigDict(extra="ignore") + + url: str | None = None + title: str | None = None + publish_date: str | None = None + excerpts: Sequence[str] | None = None + + +class _ParallelAIV1SearchResponse(BaseModel): + model_config = ConfigDict(extra="ignore") + + search_id: str | None = None + session_id: str | None = None + results: Sequence[_ParallelAIV1SearchResult] = () + usage: Sequence[Mapping[str, object]] | None = None + warnings: Sequence[Mapping[str, object]] | None = None + + class _ParallelAISourcePolicy(TypedDict, total=False): include_domains: list[str] exclude_domains: list[str] @@ -27,10 +51,16 @@ class _ParallelAIExcerptSettings(TypedDict, total=False): max_chars_per_result: int +class _ParallelAIFetchPolicy(TypedDict, total=False): + max_age_seconds: ReadOnly[int] + timeout_seconds: ReadOnly[float] + disable_cache_fallback: ReadOnly[bool] + + class _ParallelAIAdvancedSettings(TypedDict, total=False): source_policy: _ParallelAISourcePolicy excerpt_settings: _ParallelAIExcerptSettings - fetch_policy: dict + fetch_policy: _ParallelAIFetchPolicy location: str max_results: int @@ -43,14 +73,14 @@ class ParallelAISearchRequest(TypedDict, total=False): search_queries: list[str] # Required - at least one keyword search query objective: str # Optional - natural-language description of search goal - mode: str # Optional - 'turbo', 'basic', or 'advanced' (default 'advanced') + mode: str # Optional - 'turbo', 'fast', 'basic', or 'advanced' (default 'advanced') max_chars_total: int # Optional - upper bound on total excerpt characters session_id: str # Optional - tracks calls across search/extract requests client_model: str # Optional - model consuming the results advanced_settings: _ParallelAIAdvancedSettings -LEGACY_PROCESSOR_TO_MODE: Final = {"base": "basic", "pro": "advanced"} +LEGACY_PROCESSOR_TO_MODE: Final = MappingProxyType({"base": "basic", "pro": "advanced"}) class ParallelAISearchConfig(BaseSearchConfig): @@ -67,16 +97,16 @@ class ParallelAISearchConfig(BaseSearchConfig): api_base: str | None = None, **kwargs, ) -> dict: - api_key = self.resolve_server_api_key( + resolved_api_key: Final = self.resolve_server_api_key( caller_api_key=api_key, caller_api_base=api_base, key_env_vars=("PARALLEL_AI_API_KEY", "PARALLEL_API_KEY"), base_env_var="PARALLEL_AI_API_BASE", default_api_base=self.PARALLEL_AI_API_BASE, ) - if not api_key: + if not resolved_api_key: raise ValueError("PARALLEL_API_KEY is not set. Set `PARALLEL_API_KEY` environment variable.") - headers["x-api-key"] = api_key + headers["x-api-key"] = resolved_api_key headers["Content-Type"] = "application/json" return headers @@ -87,13 +117,12 @@ class ParallelAISearchConfig(BaseSearchConfig): data: dict | list[dict] | None = None, **kwargs, ) -> str: - api_base = api_base or get_secret_str("PARALLEL_AI_API_BASE") or self.PARALLEL_AI_API_BASE + resolved_api_base: Final = api_base or get_secret_str("PARALLEL_AI_API_BASE") or self.PARALLEL_AI_API_BASE - api_base = api_base.rstrip("/") - if not api_base.endswith("/v1/search"): - api_base = f"{api_base.removesuffix('/v1')}/v1/search" - - return api_base + trimmed: Final = resolved_api_base.rstrip("/") + if trimmed.endswith("/v1/search"): + return trimmed + return f"{trimmed.removesuffix('/v1')}/v1/search" def transform_search_request( self, @@ -109,14 +138,17 @@ class ParallelAISearchConfig(BaseSearchConfig): - If string: maps to `search_queries` (single item) and `objective` - If list: maps to `search_queries` (keyword queries) optional_params: Optional parameters for the request - - mode: Search mode ('turbo', 'basic', 'advanced'); defaults to 'basic' + - mode: Search mode ('turbo', 'fast', 'basic', 'advanced'); defaults to 'basic' - processor: Legacy v1beta param; 'base' maps to mode 'basic', 'pro' to 'advanced' - max_results: Maximum number of search results -> `advanced_settings.max_results` - - search_domain_filter: Domains to include -> `advanced_settings.source_policy.include_domains` + - search_domain_filter / include_domains: Domains to include -> `advanced_settings.source_policy.include_domains` - exclude_domains: Domains to exclude -> `advanced_settings.source_policy.exclude_domains` - - country: ISO 3166-1 alpha-2 code -> `advanced_settings.location` + - after_date: RFC 3339 date (YYYY-MM-DD) -> `advanced_settings.source_policy.after_date` + - country / location: ISO 3166-1 alpha-2 code -> `advanced_settings.location` - max_chars_per_result: -> `advanced_settings.excerpt_settings.max_chars_per_result` - - Any other params are passed through to the request body as-is + - fetch_policy: Cache vs live-fetch policy -> `advanced_settings.fetch_policy` + - Any other params (objective, max_chars_total, session_id, client_model, ...) + are passed through to the request body as-is Returns: Dict with request data following the v1 search request spec @@ -137,7 +169,7 @@ class ParallelAISearchConfig(BaseSearchConfig): mode = LEGACY_PROCESSOR_TO_MODE.get(processor, processor) # the v1 API defaults to 'advanced' when mode is omitted; default to 'basic' # instead to keep v1beta's default tier (processor 'base') and litellm's - # $0.004/query cost map entry for `parallel_ai/search` accurate + # cost map entry for `parallel_ai/search` accurate request_data["mode"] = mode or "basic" advanced_settings: Final[_ParallelAIAdvancedSettings] = {} @@ -148,17 +180,29 @@ class ParallelAISearchConfig(BaseSearchConfig): if "country" in params: advanced_settings["location"] = params.pop("country") + if "location" in params: + advanced_settings["location"] = params.pop("location") + if "max_chars_per_result" in params: advanced_settings["excerpt_settings"] = {"max_chars_per_result": params.pop("max_chars_per_result")} + if "fetch_policy" in params: + advanced_settings["fetch_policy"] = params.pop("fetch_policy") + source_policy: Final[_ParallelAISourcePolicy] = {} if "search_domain_filter" in params: source_policy["include_domains"] = params.pop("search_domain_filter") + if "include_domains" in params: + source_policy["include_domains"] = params.pop("include_domains") + if "exclude_domains" in params: source_policy["exclude_domains"] = params.pop("exclude_domains") + if "after_date" in params: + source_policy["after_date"] = params.pop("after_date") + if source_policy: advanced_settings["source_policy"] = source_policy @@ -170,9 +214,11 @@ class ParallelAISearchConfig(BaseSearchConfig): # unified-spec param with no v1 equivalent params.pop("max_tokens_per_page", None) - result_data: Final[dict] = dict(request_data) - result_data.update(params) - return result_data + # reserved for the provider's own reported usage, which prices the request; + # a caller-supplied value would otherwise set its own cost + params.pop(PARALLEL_AI_USAGE_PARAM, None) + + return {**request_data, **params} def transform_search_response( self, @@ -186,26 +232,49 @@ class ParallelAISearchConfig(BaseSearchConfig): Parallel AI -> LiteLLM mappings: - results[].title -> SearchResult.title - results[].url -> SearchResult.url - - results[].excerpts (array) -> SearchResult.snippet (joined string) + - results[].excerpts (array) -> SearchResult.snippet (joined string); the raw + array is preserved as an extra `excerpts` field on each result - results[].publish_date -> SearchResult.date + - search_id / session_id / warnings are preserved as extra fields on the + response; usage is preserved as `parallel_usage` (the `usage` name is + reserved for LiteLLM's token-usage object) """ - response_json: Final = raw_response.json() + parsed: Final = _ParallelAIV1SearchResponse.model_validate(raw_response.json()) - results: Final = [] - for result in response_json.get("results", []): - excerpts = result.get("excerpts") or [] - snippet = " ... ".join(excerpts) if excerpts else "" + # written unconditionally: leaving a caller-supplied value in place when the + # provider reports no usage would let the caller price its own request + logging_obj.optional_params = { + **logging_obj.optional_params, + PARALLEL_AI_USAGE_PARAM: parsed.usage, + } - search_result = SearchResult( - title=result.get("title") or "", - url=result.get("url") or "", - snippet=snippet, - date=result.get("publish_date"), - last_updated=None, + results: Final = tuple( + SearchResult.model_validate( + MappingProxyType( + { + "title": result.title or "", + "url": result.url or "", + "snippet": " ... ".join(result.excerpts or ()), + "date": result.publish_date, + "last_updated": None, + "excerpts": result.excerpts or (), + } + ) ) - results.append(search_result) - - return SearchResponse( - results=results, - object="search", + for result in parsed.results ) + + extra_fields: Final = MappingProxyType( + { + key: value + for key, value in ( + ("search_id", parsed.search_id), + ("session_id", parsed.session_id), + ("parallel_usage", parsed.usage), + ("warnings", parsed.warnings), + ) + if value is not None + } + ) + + return SearchResponse.model_validate(MappingProxyType({"results": results, "object": "search", **extra_fields})) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 55618d9f772..a3cfb300ea6 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -38556,12 +38556,22 @@ "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models" }, "parallel_ai/search": { - "input_cost_per_query": 0.004, + "input_cost_per_query": 0.005, + "litellm_provider": "parallel_ai", + "mode": "search" + }, + "parallel_ai/search-fast": { + "input_cost_per_query": 0.001, "litellm_provider": "parallel_ai", "mode": "search" }, "parallel_ai/search-pro": { - "input_cost_per_query": 0.009, + "input_cost_per_query": 0.005, + "litellm_provider": "parallel_ai", + "mode": "search" + }, + "parallel_ai/search-turbo": { + "input_cost_per_query": 0.001, "litellm_provider": "parallel_ai", "mode": "search" }, diff --git a/litellm/router_utils/search_api_router.py b/litellm/router_utils/search_api_router.py index d96defbbcd6..ab5ef5853c9 100644 --- a/litellm/router_utils/search_api_router.py +++ b/litellm/router_utils/search_api_router.py @@ -9,6 +9,7 @@ import random import traceback from collections.abc import Callable from functools import partial +from types import MappingProxyType from typing import Any, Final from litellm._logging import verbose_router_logger @@ -214,6 +215,15 @@ class SearchAPIRouter: api_key, api_base = SearchAPIRouter._resolve_search_provider_credentials( tool_litellm_params=litellm_params, ) + protected_params: Final = frozenset(("search_provider", "api_key", "api_base")) + search_params: Final = MappingProxyType( + { + key: value + for params in (litellm_params, kwargs) + for key, value in params.items() + if key not in protected_params and value is not None + } + ) verbose_router_logger.debug("Selected search tool with provider: %s", search_provider) @@ -222,7 +232,7 @@ class SearchAPIRouter: search_provider=search_provider, api_key=api_key, api_base=api_base, - **kwargs, + **search_params, ) return response diff --git a/litellm/search/cost_calculator.py b/litellm/search/cost_calculator.py index 84461115e8e..21f27075e0f 100644 --- a/litellm/search/cost_calculator.py +++ b/litellm/search/cost_calculator.py @@ -2,16 +2,37 @@ Cost calculation for search providers. """ +from collections.abc import Mapping +from types import MappingProxyType from typing import Final +from pydantic import TypeAdapter, ValidationError + from litellm.utils import get_model_info +PROVIDER_USAGE_ADAPTER: Final[TypeAdapter[tuple[Mapping[str, object], ...]]] = TypeAdapter( + tuple[Mapping[str, object], ...] +) +EMPTY_OPTIONAL_PARAMS: Final[Mapping[str, object]] = MappingProxyType({}) + + +def _provider_usage( + optional_params: Mapping[str, object] | None, + usage_param: str, +) -> tuple[Mapping[str, object], ...] | None: + params: Final = optional_params if optional_params is not None else EMPTY_OPTIONAL_PARAMS + raw_usage: Final[object] = params.get(usage_param) + try: + return PROVIDER_USAGE_ADAPTER.validate_python(raw_usage) + except ValidationError: + return None + def search_provider_cost_per_query( model: str, custom_llm_provider: str | None = None, number_of_queries: int = 1, - optional_params: dict | None = None, + optional_params: Mapping[str, object] | None = None, ) -> tuple[float, float]: """ Calculate cost for search-only providers. @@ -28,6 +49,18 @@ def search_provider_cost_per_query( Returns: Tuple of (input_cost, output_cost) where output_cost is always 0.0 """ + if custom_llm_provider == "parallel_ai": + from litellm.llms.parallel_ai.search.cost_calculator import ( + PARALLEL_AI_USAGE_PARAM, + parallel_ai_search_cost, + ) + + input_cost: Final = parallel_ai_search_cost( + optional_params=optional_params if optional_params is not None else EMPTY_OPTIONAL_PARAMS, + usage=_provider_usage(optional_params, PARALLEL_AI_USAGE_PARAM), + ) + return (input_cost, 0.0) + model_info: Final = get_model_info(model=model, custom_llm_provider=custom_llm_provider) # Check for tiered pricing (e.g., Exa AI based on max_results) diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 55618d9f772..a3cfb300ea6 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -38556,12 +38556,22 @@ "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models" }, "parallel_ai/search": { - "input_cost_per_query": 0.004, + "input_cost_per_query": 0.005, + "litellm_provider": "parallel_ai", + "mode": "search" + }, + "parallel_ai/search-fast": { + "input_cost_per_query": 0.001, "litellm_provider": "parallel_ai", "mode": "search" }, "parallel_ai/search-pro": { - "input_cost_per_query": 0.009, + "input_cost_per_query": 0.005, + "litellm_provider": "parallel_ai", + "mode": "search" + }, + "parallel_ai/search-turbo": { + "input_cost_per_query": 0.001, "litellm_provider": "parallel_ai", "mode": "search" }, diff --git a/tests/router_unit_tests/test_router_helper_utils.py b/tests/router_unit_tests/test_router_helper_utils.py index dcd2e9edf7b..7dbac243d55 100644 --- a/tests/router_unit_tests/test_router_helper_utils.py +++ b/tests/router_unit_tests/test_router_helper_utils.py @@ -2294,6 +2294,7 @@ def search_tools(): "search_provider": "perplexity", "api_key": "test-api-key", "api_base": "https://api.perplexity.ai", + "mode": "turbo", }, }, { @@ -2302,6 +2303,7 @@ def search_tools(): "search_provider": "perplexity", "api_key": "test-api-key-2", "api_base": "https://api.perplexity.ai", + "mode": "turbo", }, }, ] @@ -2393,6 +2395,7 @@ async def test_asearch_with_fallbacks_helper(search_tools): assert "search_provider" in kwargs assert kwargs["search_provider"] == "perplexity" assert "api_key" in kwargs + assert kwargs["mode"] == "turbo" assert kwargs["query"] == "helper test query" return mock_response diff --git a/tests/test_litellm/llms/parallel_ai/test_parallel_ai_search.py b/tests/test_litellm/llms/parallel_ai/test_parallel_ai_search.py index 8a9ae4dae6d..62b4d003b45 100644 --- a/tests/test_litellm/llms/parallel_ai/test_parallel_ai_search.py +++ b/tests/test_litellm/llms/parallel_ai/test_parallel_ai_search.py @@ -2,6 +2,7 @@ Tests for Parallel AI Search API integration (v1 endpoint). """ +import json from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -30,13 +31,41 @@ MOCK_V1_RESPONSE = { } -def _mock_response(): +def _mock_response(payload=None): mock_response = MagicMock() mock_response.status_code = 200 - mock_response.json.return_value = MOCK_V1_RESPONSE + mock_response.json.return_value = payload if payload is not None else MOCK_V1_RESPONSE return mock_response +@pytest.fixture +def httpx_transport(monkeypatch): + monkeypatch.setattr( # test-quality-ok: respx needs HTTPX enabled to fake the provider HTTP boundary. + litellm, + "disable_aiohttp_transport", + True, + ) + litellm.in_memory_llm_clients_cache.flush_cache() + yield + litellm.in_memory_llm_clients_cache.flush_cache() + + +@pytest.fixture +def bundled_cost_map(monkeypatch): + """Price lookups against the bundled cost map. + + litellm caches model-info lookups, so swapping ``model_cost`` only takes + effect once those caches are invalidated -- on the way in and back out. + """ + from litellm.utils import _invalidate_model_cost_lowercase_map + + monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) + _invalidate_model_cost_lowercase_map() + yield + monkeypatch.undo() + _invalidate_model_cost_lowercase_map() + + class TestParallelAISearch: @pytest.fixture(autouse=True) def _set_api_key(self, monkeypatch): @@ -135,9 +164,7 @@ class TestParallelAISearch: json_data = mock_post.call_args.kwargs.get("json") assert json_data["mode"] == "basic" - @pytest.mark.parametrize( - "processor,expected_mode", [("base", "basic"), ("pro", "advanced")] - ) + @pytest.mark.parametrize("processor,expected_mode", [("base", "basic"), ("pro", "advanced")]) @pytest.mark.asyncio async def test_legacy_processor_maps_to_mode(self, processor, expected_mode): with patch( @@ -222,9 +249,7 @@ class TestParallelAISearch: "arxiv.org", "nature.com", ] - assert advanced_settings["source_policy"]["exclude_domains"] == [ - "reddit.com" - ] + assert advanced_settings["source_policy"]["exclude_domains"] == ["reddit.com"] assert advanced_settings["excerpt_settings"]["max_chars_per_result"] == 1500 assert "max_results" not in json_data @@ -306,10 +331,7 @@ class TestParallelAISearch: ) call_args = mock_post.call_args - assert ( - call_args.kwargs["url"] - == "https://proxy.internal.example.com/v1/search" - ) + assert call_args.kwargs["url"] == "https://proxy.internal.example.com/v1/search" @pytest.mark.asyncio async def test_caller_api_base_without_key_is_refused(self, monkeypatch): @@ -338,3 +360,147 @@ class TestParallelAISearch: query="AI developments", search_provider="parallel_ai", ) + + @pytest.mark.asyncio + async def test_flat_source_and_fetch_params_nest_under_advanced_settings(self, respx_mock, httpx_transport): + route = respx_mock.post("https://api.parallel.ai/v1/search").respond(json=MOCK_V1_RESPONSE) + + await litellm.asearch( + query="AI developments", + search_provider="parallel_ai", + objective="find peer-reviewed AI research", + include_domains=["arxiv.org"], + after_date="2026-01-01", + location="gb", + fetch_policy={"max_age_seconds": 600, "disable_cache_fallback": True}, + client_model="claude-fable-5", + ) + + json_data = json.loads(route.calls[0].request.content) + assert json_data["objective"] == "find peer-reviewed AI research" + assert json_data["client_model"] == "claude-fable-5" + + advanced_settings = json_data["advanced_settings"] + assert advanced_settings["location"] == "gb" + assert advanced_settings["fetch_policy"] == { + "max_age_seconds": 600, + "disable_cache_fallback": True, + } + assert advanced_settings["source_policy"]["include_domains"] == ["arxiv.org"] + assert advanced_settings["source_policy"]["after_date"] == "2026-01-01" + + assert "include_domains" not in json_data + assert "after_date" not in json_data + assert "location" not in json_data + assert "fetch_policy" not in json_data + + @pytest.mark.asyncio + async def test_response_preserves_raw_parallel_fields(self, respx_mock, httpx_transport): + respx_mock.post("https://api.parallel.ai/v1/search").respond(json=MOCK_V1_RESPONSE) + + response = await litellm.asearch( + query="AI developments", + search_provider="parallel_ai", + ) + + dumped = response.model_dump() + assert dumped["search_id"] == "search_abc123" + assert dumped["session_id"] == "session_xyz" + assert dumped["parallel_usage"] == [{"name": "search_advanced", "count": 1}] + + first = response.results[0].model_dump() + assert first["excerpts"] == ["First excerpt.", "Second excerpt."] + + @pytest.mark.asyncio + async def test_response_normalizes_null_result_fields(self, respx_mock, httpx_transport): + response_payload = { + **MOCK_V1_RESPONSE, + "results": [{"url": None, "title": None, "publish_date": None, "excerpts": None}], + } + respx_mock.post("https://api.parallel.ai/v1/search").respond(json=response_payload) + + response = await litellm.asearch( + query="AI developments", + search_provider="parallel_ai", + ) + + assert len(response.results) == 1 + result = response.results[0] + assert result.url == "" + assert result.title == "" + assert result.snippet == "" + assert result.date is None + assert result.model_dump()["excerpts"] == () + + @pytest.mark.parametrize( + "mode,usage,max_results,expected_cost", + [ + ("turbo", [{"name": "sku_search", "count": 1}], None, 0.001), + ("fast", [{"name": "sku_search", "count": 1}], None, 0.001), + ("basic", [{"name": "sku_search", "count": 1}], None, 0.005), + ("advanced", [{"name": "sku_search", "count": 1}], None, 0.005), + ( + "basic", + [ + {"name": "sku_search", "count": 1}, + {"name": "sku_search_additional_results", "count": 2}, + ], + 20, + 0.007, + ), + ("basic", None, 20, 0.015), + ], + ) + @pytest.mark.asyncio + async def test_search_cost_uses_mode_and_provider_usage( + self, mode, usage, max_results, expected_cost, bundled_cost_map, respx_mock, httpx_transport + ): + response_payload = {**MOCK_V1_RESPONSE, "usage": usage} + respx_mock.post("https://api.parallel.ai/v1/search").respond(json=response_payload) + + response = await litellm.asearch( + query="AI developments", + search_provider="parallel_ai", + mode=mode, + max_results=max_results, + ) + + assert response._hidden_params["response_cost"] == pytest.approx(expected_cost) + + @pytest.mark.asyncio + async def test_search_cost_treats_keyword_queries_as_one_request( + self, bundled_cost_map, respx_mock, httpx_transport + ): + response_payload = { + **MOCK_V1_RESPONSE, + "usage": [{"name": "sku_search", "count": 1}], + } + respx_mock.post("https://api.parallel.ai/v1/search").respond(json=response_payload) + + response = await litellm.asearch( + query=["AI developments", "machine learning trends"], + search_provider="parallel_ai", + mode="basic", + ) + + assert response._hidden_params["response_cost"] == pytest.approx(0.005) + + @pytest.mark.asyncio + async def test_caller_cannot_supply_provider_usage(self, bundled_cost_map, respx_mock, httpx_transport): + """`_parallel_ai_usage` prices the request, so a caller must not be able to set it. + + The provider reports no usage here, which is the case where a caller-supplied + value would otherwise survive into the cost calculation. + """ + response_payload = {k: v for k, v in MOCK_V1_RESPONSE.items() if k != "usage"} + route = respx_mock.post("https://api.parallel.ai/v1/search").respond(json=response_payload) + + response = await litellm.asearch( + query="AI developments", + search_provider="parallel_ai", + mode="basic", + _parallel_ai_usage=[{"name": "sku_search", "count": 0}], + ) + + assert response._hidden_params["response_cost"] == pytest.approx(0.005) + assert "_parallel_ai_usage" not in json.loads(route.calls[0].request.content) diff --git a/tests/test_litellm/llms/parallel_ai/test_parallel_ai_search_gateway.py b/tests/test_litellm/llms/parallel_ai/test_parallel_ai_search_gateway.py new file mode 100644 index 00000000000..72c69fc622c --- /dev/null +++ b/tests/test_litellm/llms/parallel_ai/test_parallel_ai_search_gateway.py @@ -0,0 +1,191 @@ +"""Gateway coverage for Parallel AI Search.""" + +from __future__ import annotations + +from collections.abc import Iterator +from typing import Final +from unittest.mock import AsyncMock + +import httpx +import pytest +from fastapi.testclient import TestClient + +import litellm +from litellm import Router +from litellm.integrations.websearch_interception.handler import ( + WebSearchInterceptionLogger, +) +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler +from litellm.proxy import proxy_server +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.types.utils import LlmProviders + +PARALLEL_SEARCH_URL: Final = "https://api.parallel.ai/v1/search" + + +@pytest.fixture +def client() -> TestClient: + return TestClient(proxy_server.app, raise_server_exceptions=False) + + +@pytest.fixture +def auth_as() -> Iterator[None]: + async def _authorized_request() -> UserAPIKeyAuth: + return UserAPIKeyAuth( + api_key="hashed-sk-test", + user_id="parallel-test-user", + ) + + previous: Final = proxy_server.app.dependency_overrides.get(user_api_key_auth) + proxy_server.app.dependency_overrides[user_api_key_auth] = _authorized_request + try: + yield + finally: + if previous is None: + proxy_server.app.dependency_overrides.pop(user_api_key_auth, None) + else: + proxy_server.app.dependency_overrides[user_api_key_auth] = previous + + +def _parallel_search_body() -> dict[str, object]: + return { + "search_id": "search_parallel_gateway", + "results": [ + { + "url": "https://example.com/parallel", + "title": "Parallel result", + "publish_date": "2026-08-13", + "excerpts": ["First excerpt", "Second excerpt"], + } + ], + "usage": [{"name": "sku_search", "count": 1}], + } + + +def _parallel_router(mode: str = "turbo") -> Router: + return Router( + model_list=[], + search_tools=[ + { + "search_tool_name": "parallel-search", + "litellm_params": { + "search_provider": "parallel_ai", + "api_key": "parallel-search-key", + "mode": mode, + }, + } + ], + num_retries=0, + ) + + +def _mock_async_post( + monkeypatch, + *, + url: str, + response_body: dict[str, object], +) -> AsyncMock: + response = httpx.Response( + status_code=200, + json=response_body, + request=httpx.Request("POST", url), + ) + mock_post = AsyncMock(return_value=response) + monkeypatch.setattr(AsyncHTTPHandler, "post", mock_post) + return mock_post + + +def test_parallel_search_gateway_route(client, auth_as, monkeypatch): + """The named search route selects its configured Parallel Search tool. + + The tool-level `mode` must survive the router hop, so the upstream request + is sent as `turbo` rather than falling back to the adapter default. + """ + monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) + monkeypatch.setattr(proxy_server, "llm_router", _parallel_router()) + mock_post = _mock_async_post( + monkeypatch, + url=PARALLEL_SEARCH_URL, + response_body=_parallel_search_body(), + ) + + response = client.post( + "/v1/search/parallel-search", + json={"query": "Parallel AI news", "max_results": 3}, + ) + + assert response.status_code == 200, response.text + assert response.json()["results"] == [ + { + "title": "Parallel result", + "url": "https://example.com/parallel", + "snippet": "First excerpt ... Second excerpt", + "date": "2026-08-13", + "last_updated": None, + "excerpts": ["First excerpt", "Second excerpt"], + } + ] + + request_kwargs = mock_post.await_args.kwargs + assert request_kwargs["url"] == PARALLEL_SEARCH_URL + assert request_kwargs["headers"]["x-api-key"] == "parallel-search-key" + assert request_kwargs["json"] == { + "objective": "Parallel AI news", + "search_queries": ["Parallel AI news"], + "mode": "turbo", + "advanced_settings": {"max_results": 3}, + } + + +@pytest.mark.asyncio +async def test_web_search_interception_executes_parallel_search(monkeypatch): + """An intercepted web-search call uses the configured Parallel Search tool.""" + monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) + monkeypatch.setattr(proxy_server, "llm_router", _parallel_router(mode="fast")) + mock_post = _mock_async_post( + monkeypatch, + url=PARALLEL_SEARCH_URL, + response_body=_parallel_search_body(), + ) + logger = WebSearchInterceptionLogger( + enabled_providers=[LlmProviders.OPENAI], + search_tool_name="parallel-search", + ) + + plan = await logger.async_build_responses_agentic_loop_plan( + tools={ + "tool_calls": [ + { + "id": "fc_parallel", + "call_id": "fc_parallel", + "type": "function_call", + "name": "litellm_web_search", + "arguments": '{"query":"Parallel AI news"}', + "input": {"query": "Parallel AI news"}, + } + ] + }, + model="gpt-5", + messages=[{"role": "user", "content": "Research Parallel"}], + response=None, + optional_params={"tools": [{"type": "function", "name": "litellm_web_search"}]}, + logging_obj=None, + stream=False, + kwargs={"custom_llm_provider": "openai"}, + ) + + assert plan.run_agentic_loop is True + assert plan.request_patch is not None + assert plan.request_patch.messages[-1] == { + "type": "function_call_output", + "call_id": "fc_parallel", + "output": ( + "Title: Parallel result\nURL: https://example.com/parallel\nSnippet: First excerpt ... Second excerpt" + ), + } + + request_kwargs = mock_post.await_args.kwargs + assert request_kwargs["url"] == PARALLEL_SEARCH_URL + assert request_kwargs["headers"]["x-api-key"] == "parallel-search-key" + assert request_kwargs["json"]["mode"] == "fast" From 974b331a4d09e2883d6fe84bb87ce57cc49ab5bb Mon Sep 17 00:00:00 2001 From: mateo Date: Wed, 2 Sep 2026 06:15:52 +0000 Subject: [PATCH 61/97] 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 62/97] 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 63/97] 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 64/97] 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 65/97] 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 66/97] 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 67/97] 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 68/97] 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 69/97] 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 70/97] 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 71/97] 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 72/97] 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 73/97] 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 74/97] 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 75/97] 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 76/97] 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 77/97] 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 78/97] 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 79/97] 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 80/97] 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 25987cb961567541922cef80e1433137f42c7b77 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 31 Aug 2026 14:42:48 -0700 Subject: [PATCH 81/97] test(build): validate release wheel contracts --- .github/scripts/smoke_test_native_wheel.py | 68 ++++++++ .github/scripts/verify_linux_native_wheel.py | 161 +++++++++++++++++++ .github/workflows/test-rust.yml | 111 +++++++++++++ litellm-rust/crates/python-bridge/Cargo.toml | 1 + litellm-rust/crates/python-bridge/src/lib.rs | 8 + 5 files changed, 349 insertions(+) create mode 100644 .github/scripts/smoke_test_native_wheel.py create mode 100644 .github/scripts/verify_linux_native_wheel.py diff --git a/.github/scripts/smoke_test_native_wheel.py b/.github/scripts/smoke_test_native_wheel.py new file mode 100644 index 00000000000..577bb32fcf0 --- /dev/null +++ b/.github/scripts/smoke_test_native_wheel.py @@ -0,0 +1,68 @@ +from __future__ import annotations + +import subprocess +import sys +import tempfile +import zipfile +from pathlib import Path +from typing import Final + +CHILD_SCRIPT: Final = """ +from importlib.util import module_from_spec, spec_from_file_location +from pathlib import Path +import sys + +native_path = Path(sys.argv[1]) +spec = spec_from_file_location("litellm.rust_bridge._native", native_path) +if spec is None or spec.loader is None: + raise RuntimeError("cannot create native extension import specification") +module = module_from_spec(spec) +spec.loader.exec_module(module) + +before = module.gil_stats() +if not isinstance(before.get("releases"), int): + raise AssertionError(f"unexpected gil_stats result: {before!r}") + +try: + module._panic_for_test() +except BaseException as error: + if type(error).__name__ != "PanicException": + raise AssertionError(f"expected PanicException, got {type(error).__name__}") from error +else: + raise AssertionError("Rust panic returned without raising") + +after = module.gil_stats() +if not isinstance(after.get("releases"), int): + raise AssertionError(f"native module unusable after panic: {after!r}") +""" + + +def main() -> int: + if len(sys.argv) != 2: + sys.stderr.write(f"usage: {Path(sys.argv[0]).name} WHEEL\n") + return 2 + + wheel: Final = Path(sys.argv[1]) + with tempfile.TemporaryDirectory() as temporary_directory, zipfile.ZipFile(wheel) as archive: + native_members: Final = tuple( + member + for member in archive.infolist() + if member.filename.startswith("litellm/rust_bridge/_native.") and member.filename.endswith(".so") + ) + if len(native_members) != 1: + sys.stderr.write(f"expected one native extension, found {len(native_members)}\n") + return 1 + + native_path: Final = Path(temporary_directory) / Path(native_members[0].filename).name + native_path.write_bytes(archive.read(native_members[0])) + result: Final = subprocess.run((sys.executable, "-c", CHILD_SCRIPT, str(native_path)), check=False) + + if result.returncode != 0: + sys.stderr.write(f"native wheel smoke test exited with status {result.returncode}\n") + return 1 + + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/.github/scripts/verify_linux_native_wheel.py b/.github/scripts/verify_linux_native_wheel.py new file mode 100644 index 00000000000..1de98ae0168 --- /dev/null +++ b/.github/scripts/verify_linux_native_wheel.py @@ -0,0 +1,161 @@ +from __future__ import annotations + +import importlib.util +import os +import re +import subprocess +import sys +import zipfile +from pathlib import Path, PurePosixPath +from typing import Final + + +def _loads_native_module(native_path: Path) -> bool: + module_spec: Final = importlib.util.spec_from_file_location("litellm.rust_bridge._native", native_path) + if module_spec is None or module_spec.loader is None: + return False + try: + native_module: Final = importlib.util.module_from_spec(module_spec) + module_spec.loader.exec_module(native_module) + except Exception as error: + sys.stderr.write(f"native module load failed: {error}\n") + return False + return True + + +def main() -> int: + if len(sys.argv) != 2: + sys.stderr.write(f"usage: {Path(sys.argv[0]).name} WHEEL\n") + return 2 + + wheel: Final = Path(sys.argv[1]) + with zipfile.ZipFile(wheel) as archive: + wheel_members: Final = archive.infolist() + native_members: Final = tuple( + member + for member in wheel_members + if member.filename.startswith("litellm/rust_bridge/_native.") and member.filename.endswith(".so") + ) + if len(native_members) != 1: + sys.stderr.write(f"expected one native extension, found {len(native_members)}\n") + return 1 + + unexpected_members: Final = tuple( + member.filename + for member in wheel_members + if member.filename.endswith((".pdb", ".dwp", ".rlib", ".rmeta", "Cargo.toml", "Cargo.lock")) + or any(part.endswith(".dSYM") for part in PurePosixPath(member.filename).parts) + ) + native_member: Final = native_members[0] + uncompressed_wheel_size: Final = sum(member.file_size for member in wheel_members) + native_path: Final = wheel.parent / "native" / Path(native_member.filename).name + native_path.parent.mkdir(parents=True, exist_ok=True) + native_path.write_bytes(archive.read(native_member)) + + wheel_tags: Final = wheel.stem.rsplit("-", maxsplit=3) + if len(wheel_tags) != 4: + sys.stderr.write(f"cannot parse wheel tags from {wheel.name}\n") + return 1 + + python_tag: Final = wheel_tags[1] + abi_tag: Final = wheel_tags[2] + platform_tag: Final = wheel_tags[3] + commit_sha: Final = os.environ.get("RELEASE_WHEEL_COMMIT_SHA", os.environ.get("GITHUB_SHA", "unknown")) + rustc_version: Final = subprocess.run( + ("rustc", "--version"), + check=True, + capture_output=True, + text=True, + ).stdout.strip() + pyproject: Final = (Path(__file__).parents[2] / "pyproject.toml").read_text() + maturin_match: Final = re.search(r'"maturin==([^";]+)', pyproject) + if maturin_match is None: + sys.stderr.write("build-system does not pin an exact Maturin version\n") + return 1 + + maturin_version: Final = maturin_match.group(1) + native_percentage: Final = native_member.file_size / uncompressed_wheel_size * 100 + size_report: Final = "\n".join( + ( + "## Native wheel build report", + "", + "| Build | Value |", + "| --- | --- |", + f"| Commit | `{commit_sha}` |", + f"| Platform | `{platform_tag}` |", + f"| Python ABI | `{python_tag}-{abi_tag}` |", + f"| Rust compiler | `{rustc_version}` |", + f"| Maturin | `{maturin_version}` |", + "| Cargo profile | `release` |", + "", + "| Artifact | Size |", + "| --- | ---: |", + f"| Compressed wheel | {wheel.stat().st_size / 1_000_000:.2f} MB |", + f"| Uncompressed wheel | {uncompressed_wheel_size / 1_000_000:.2f} MB |", + f"| Native extension | {native_member.file_size / 1_000_000:.2f} MB |", + f"| Native share | {native_percentage:.2f}% |", + "", + ) + ) + summary_path: Final = os.environ.get("GITHUB_STEP_SUMMARY") + if summary_path is None: + sys.stdout.write(size_report) + else: + Path(summary_path).write_text(size_report) + + sections: Final = subprocess.run( + ("readelf", "--sections", "--wide", native_path), + check=True, + capture_output=True, + text=True, + ).stdout + debug_sections: Final = tuple(section for section in (".debug_", ".zdebug_") if section in sections) + debug_sections_absent: Final = not debug_sections + static_symbol_table_absent: Final = ".symtab" not in sections + + dynamic_symbols: Final = subprocess.run( + ("readelf", "--dyn-syms", "--wide", native_path), + check=True, + capture_output=True, + text=True, + ).stdout + extension_entry_point_present: Final = "PyInit__native" in dynamic_symbols + native_module_loads: Final = _loads_native_module(native_path) + native_size_limit: Final = 20_000_000 + native_size_within_limit: Final = native_member.file_size <= native_size_limit + validations: Final = ( + ("Debug sections are absent", debug_sections_absent), + ("Static symbol table is absent", static_symbol_table_absent), + ("Python extension entry point is present", extension_entry_point_present), + ("Native module loads", native_module_loads), + ("Native extension does not exceed 20 MB", native_size_within_limit), + ("Wheel contents are valid", not unexpected_members), + ) + + verified_report: Final = size_report + "\n".join( + ("", "| Validation | Expected | Result |", "| --- | --- | :---: |") + + tuple(f"| {label} | Yes | {'O' if passed else 'X'} |" for label, passed in validations) + + ("",) + ) + report_path: Final = os.environ.get("RELEASE_WHEEL_REPORT") + if report_path is not None: + Path(report_path).write_text(size_report) + if summary_path is not None: + Path(summary_path).write_text(verified_report) + + if debug_sections: + sys.stderr.write(f"{native_member.filename} contains debug sections: {', '.join(debug_sections)}\n") + if not static_symbol_table_absent: + sys.stderr.write(f"{native_member.filename} contains a static symbol table\n") + if not extension_entry_point_present: + sys.stderr.write("native extension does not export PyInit__native\n") + if not native_size_within_limit: + sys.stderr.write(f"native extension exceeds 20 MB: {native_member.file_size / 1_000_000:.2f} MB\n") + if unexpected_members: + sys.stderr.write(f"wheel contains unexpected build artifacts: {', '.join(unexpected_members)}\n") + + return 0 if all(passed for _, passed in validations) else 1 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/.github/workflows/test-rust.yml b/.github/workflows/test-rust.yml index 21e1bcb90c6..b05ddb81740 100644 --- a/.github/workflows/test-rust.yml +++ b/.github/workflows/test-rust.yml @@ -4,6 +4,10 @@ on: push: paths: - "litellm-rust/**" + - ".cargo/**" + - "pyproject.toml" + - ".github/scripts/smoke_test_native_wheel.py" + - ".github/scripts/verify_linux_native_wheel.py" - ".github/workflows/test-rust.yml" pull_request: branches: @@ -13,6 +17,10 @@ on: - "litellm_**" paths: - "litellm-rust/**" + - ".cargo/**" + - "pyproject.toml" + - ".github/scripts/smoke_test_native_wheel.py" + - ".github/scripts/verify_linux_native_wheel.py" - ".github/workflows/test-rust.yml" permissions: @@ -69,3 +77,106 @@ jobs: - name: Run core tests with Bedrock auth run: cargo test -p litellm-core --features bedrock-auth --locked + + release-wheel: + name: release wheel + runs-on: ubuntu-latest + timeout-minutes: 20 + permissions: + contents: read + pull-requests: write + env: + CARGO_TERM_COLOR: always + + steps: + - name: Checkout repository + uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 + with: + persist-credentials: false + + - name: Set up Python + uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0 + with: + python-version: "3.12" + + - name: Set up uv + uses: ./.github/actions/setup-uv-with-retries + with: + version: "0.10.9" + + - name: Set up Rust + run: | + rustup toolchain install stable --profile minimal + rustup default stable + + - name: Cache release build + uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0 + with: + path: | + ~/.cargo/registry + ~/.cargo/git + litellm-rust/target + key: ${{ runner.os }}-cargo-release-${{ hashFiles('litellm-rust/Cargo.lock') }} + restore-keys: | + ${{ runner.os }}-cargo-release- + + - name: Build release wheel + run: uv build --wheel --out-dir dist + + - name: Verify stripped native extension + env: + RELEASE_WHEEL_COMMIT_SHA: ${{ github.event.pull_request.head.sha || github.sha }} + RELEASE_WHEEL_REPORT: dist/release-wheel-report.md + run: python .github/scripts/verify_linux_native_wheel.py dist/*.whl + + - name: Build panic contract wheel + run: >- + uv build --wheel --out-dir panic-dist + --config-setting "maturin.build-args=--features panic-test" + + - name: Smoke-test native panic unwinding + run: python .github/scripts/smoke_test_native_wheel.py panic-dist/*.whl + + - name: Report release wheel size on PR + if: >- + github.event_name == 'pull_request' && + github.event.pull_request.head.repo.full_name == github.repository + uses: actions/github-script@f28e40c7f34bde8b3046d885e986cb6290c5673b # v7.1.0 + env: + COMMENT_MARKER: "" + REPORT_PATH: dist/release-wheel-report.md + with: + script: | + const fs = require("fs"); + const marker = process.env.COMMENT_MARKER; + const report = fs.readFileSync(process.env.REPORT_PATH, "utf8"); + const body = `${marker}\n${report}`; + const comments = await github.paginate(github.rest.issues.listComments, { + owner: context.repo.owner, + repo: context.repo.repo, + issue_number: context.issue.number, + per_page: 100, + }); + const existing = comments.find((comment) => comment.body?.includes(marker)); + if (existing) { + await github.rest.issues.updateComment({ + owner: context.repo.owner, + repo: context.repo.repo, + comment_id: existing.id, + body, + }); + } else { + await github.rest.issues.createComment({ + owner: context.repo.owner, + repo: context.repo.repo, + issue_number: context.issue.number, + body, + }); + } + + - name: Upload release wheel + uses: actions/upload-artifact@4cec3d8aa04e39d1a68397de0c4cd6fb9dce8ec1 # v4.6.1 + with: + name: litellm-release-wheel-linux-x86_64 + path: dist/*.whl + if-no-files-found: error diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index d461a483ae0..b1fdfd7677a 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -13,6 +13,7 @@ crate-type = ["cdylib"] default = ["abi3"] abi3 = ["pyo3/abi3-py310"] extension-module = ["pyo3/extension-module"] +panic-test = [] [dependencies] litellm-core = { workspace = true, features = ["bedrock-auth"] } diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index f9e75f45f75..18e0b05cbb5 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -614,6 +614,12 @@ fn gil_stats(py: Python<'_>) -> PyResult> { Ok(stats.into_any().unbind()) } +#[cfg(feature = "panic-test")] +#[pyfunction] +fn _panic_for_test() { + panic!("intentional PyO3 panic smoke test"); +} + #[pymodule] fn _native(module: &Bound<'_, PyModule>) -> PyResult<()> { let py = module.py(); @@ -630,5 +636,7 @@ fn _native(module: &Bound<'_, PyModule>) -> PyResult<()> { module.add_function(wrap_pyfunction!(achat_completions, module)?)?; module.add_class::()?; module.add_function(wrap_pyfunction!(gil_stats, module)?)?; + #[cfg(feature = "panic-test")] + module.add_function(wrap_pyfunction!(_panic_for_test, module)?)?; Ok(()) } From 2f362cfec2d122d4d2c14b73a988a20fe61e6629 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 31 Aug 2026 14:59:03 -0700 Subject: [PATCH 82/97] fix(ci): preserve release wheel contract parity --- .github/scripts/verify_linux_native_wheel.py | 2 +- .github/workflows/test-rust.yml | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/scripts/verify_linux_native_wheel.py b/.github/scripts/verify_linux_native_wheel.py index 1de98ae0168..b0ab34df98d 100644 --- a/.github/scripts/verify_linux_native_wheel.py +++ b/.github/scripts/verify_linux_native_wheel.py @@ -139,7 +139,7 @@ def main() -> int: ) report_path: Final = os.environ.get("RELEASE_WHEEL_REPORT") if report_path is not None: - Path(report_path).write_text(size_report) + Path(report_path).write_text(verified_report) if summary_path is not None: Path(summary_path).write_text(verified_report) diff --git a/.github/workflows/test-rust.yml b/.github/workflows/test-rust.yml index b05ddb81740..1a9f36f24d2 100644 --- a/.github/workflows/test-rust.yml +++ b/.github/workflows/test-rust.yml @@ -132,7 +132,7 @@ jobs: - name: Build panic contract wheel run: >- uv build --wheel --out-dir panic-dist - --config-setting "maturin.build-args=--features panic-test" + --config-setting "maturin.build-args=--features panic-test,extension-module" - name: Smoke-test native panic unwinding run: python .github/scripts/smoke_test_native_wheel.py panic-dist/*.whl From 9dc9cd325cf5adca79d910a30c420381a28d1923 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 31 Aug 2026 20:21:37 -0700 Subject: [PATCH 83/97] fix(ci): isolate release wheel reporting permissions --- .github/scripts/verify_linux_native_wheel.py | 3 - .../workflows/report-rust-release-wheel.yml | 95 +++++++++++++++++++ .github/workflows/test-rust.yml | 39 -------- 3 files changed, 95 insertions(+), 42 deletions(-) create mode 100644 .github/workflows/report-rust-release-wheel.yml diff --git a/.github/scripts/verify_linux_native_wheel.py b/.github/scripts/verify_linux_native_wheel.py index b0ab34df98d..96264beb632 100644 --- a/.github/scripts/verify_linux_native_wheel.py +++ b/.github/scripts/verify_linux_native_wheel.py @@ -137,9 +137,6 @@ def main() -> int: + tuple(f"| {label} | Yes | {'O' if passed else 'X'} |" for label, passed in validations) + ("",) ) - report_path: Final = os.environ.get("RELEASE_WHEEL_REPORT") - if report_path is not None: - Path(report_path).write_text(verified_report) if summary_path is not None: Path(summary_path).write_text(verified_report) diff --git a/.github/workflows/report-rust-release-wheel.yml b/.github/workflows/report-rust-release-wheel.yml new file mode 100644 index 00000000000..76aefd70ef6 --- /dev/null +++ b/.github/workflows/report-rust-release-wheel.yml @@ -0,0 +1,95 @@ +name: Report LiteLLM Rust release wheel + +on: # zizmor: ignore[dangerous-triggers] reporter executes no PR code and consumes no PR artifacts or outputs + workflow_run: + workflows: + - LiteLLM Rust + types: + - completed + +permissions: {} + +jobs: + report-release-wheel: + name: report release wheel + if: >- + github.event.workflow_run.conclusion == 'success' && + github.event.workflow_run.event == 'pull_request' && + github.event.workflow_run.path == '.github/workflows/test-rust.yml' && + github.event.workflow_run.head_repository.full_name == github.repository && + github.event.workflow_run.pull_requests[0].number != null + runs-on: ubuntu-latest + timeout-minutes: 5 + permissions: + issues: write + + steps: + - name: Link release wheel report on PR + uses: actions/github-script@f28e40c7f34bde8b3046d885e986cb6290c5673b # v7.1.0 + env: + COMMENT_MARKER: "" + with: + script: | + const marker = process.env.COMMENT_MARKER; + const workflowRun = context.payload.workflow_run; + if ( + workflowRun.conclusion !== "success" || + workflowRun.event !== "pull_request" || + workflowRun.path !== ".github/workflows/test-rust.yml" || + workflowRun.head_repository?.full_name !== + `${context.repo.owner}/${context.repo.repo}` || + workflowRun.pull_requests?.length !== 1 + ) { + throw new Error("unexpected source workflow"); + } + const pullRequest = workflowRun.pull_requests[0]; + const pullRequestNumber = pullRequest.number; + const headSha = workflowRun.head_sha; + const runId = workflowRun.id; + if ( + !Number.isSafeInteger(pullRequestNumber) || + pullRequestNumber <= 0 || + !Number.isSafeInteger(runId) || + runId <= 0 || + !/^[0-9a-f]{40}$/.test(headSha) || + pullRequest.head?.sha !== headSha + ) { + throw new Error("invalid source workflow metadata"); + } + const runUrl = + `${context.serverUrl}/${context.repo.owner}/${context.repo.repo}` + + `/actions/runs/${runId}`; + const body = [ + marker, + "## LiteLLM Rust workflow", + "", + `Workflow completed successfully for \`${headSha}\``, + "", + `[View workflow run](${runUrl})`, + ].join("\n"); + const comments = await github.paginate(github.rest.issues.listComments, { + owner: context.repo.owner, + repo: context.repo.repo, + issue_number: pullRequestNumber, + per_page: 100, + }); + const existing = comments.find( + (comment) => + comment.user?.login === "github-actions[bot]" && + comment.body?.startsWith(marker), + ); + if (existing) { + await github.rest.issues.updateComment({ + owner: context.repo.owner, + repo: context.repo.repo, + comment_id: existing.id, + body, + }); + } else { + await github.rest.issues.createComment({ + owner: context.repo.owner, + repo: context.repo.repo, + issue_number: pullRequestNumber, + body, + }); + } diff --git a/.github/workflows/test-rust.yml b/.github/workflows/test-rust.yml index 1a9f36f24d2..b01e7fabe4a 100644 --- a/.github/workflows/test-rust.yml +++ b/.github/workflows/test-rust.yml @@ -84,7 +84,6 @@ jobs: timeout-minutes: 20 permissions: contents: read - pull-requests: write env: CARGO_TERM_COLOR: always @@ -126,7 +125,6 @@ jobs: - name: Verify stripped native extension env: RELEASE_WHEEL_COMMIT_SHA: ${{ github.event.pull_request.head.sha || github.sha }} - RELEASE_WHEEL_REPORT: dist/release-wheel-report.md run: python .github/scripts/verify_linux_native_wheel.py dist/*.whl - name: Build panic contract wheel @@ -137,43 +135,6 @@ jobs: - name: Smoke-test native panic unwinding run: python .github/scripts/smoke_test_native_wheel.py panic-dist/*.whl - - name: Report release wheel size on PR - if: >- - github.event_name == 'pull_request' && - github.event.pull_request.head.repo.full_name == github.repository - uses: actions/github-script@f28e40c7f34bde8b3046d885e986cb6290c5673b # v7.1.0 - env: - COMMENT_MARKER: "" - REPORT_PATH: dist/release-wheel-report.md - with: - script: | - const fs = require("fs"); - const marker = process.env.COMMENT_MARKER; - const report = fs.readFileSync(process.env.REPORT_PATH, "utf8"); - const body = `${marker}\n${report}`; - const comments = await github.paginate(github.rest.issues.listComments, { - owner: context.repo.owner, - repo: context.repo.repo, - issue_number: context.issue.number, - per_page: 100, - }); - const existing = comments.find((comment) => comment.body?.includes(marker)); - if (existing) { - await github.rest.issues.updateComment({ - owner: context.repo.owner, - repo: context.repo.repo, - comment_id: existing.id, - body, - }); - } else { - await github.rest.issues.createComment({ - owner: context.repo.owner, - repo: context.repo.repo, - issue_number: context.issue.number, - body, - }); - } - - name: Upload release wheel uses: actions/upload-artifact@4cec3d8aa04e39d1a68397de0c4cd6fb9dce8ec1 # v4.6.1 with: From e6a317e0790c575d6452f59c90e861f110b54a0c Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 31 Aug 2026 20:53:17 -0700 Subject: [PATCH 84/97] fix(ci): enforce release wheel metadata contract --- .github/scripts/verify_linux_native_wheel.py | 110 ++++++++++-- .../test_verify_linux_native_wheel.py | 170 ++++++++++++++++++ 2 files changed, 266 insertions(+), 14 deletions(-) create mode 100644 tests/test_litellm/test_verify_linux_native_wheel.py diff --git a/.github/scripts/verify_linux_native_wheel.py b/.github/scripts/verify_linux_native_wheel.py index 96264beb632..61f176ab183 100644 --- a/.github/scripts/verify_linux_native_wheel.py +++ b/.github/scripts/verify_linux_native_wheel.py @@ -6,21 +6,44 @@ import re import subprocess import sys import zipfile +from email import policy +from email.parser import BytesParser +from itertools import product from pathlib import Path, PurePosixPath -from typing import Final +from types import ModuleType +from typing import Final, cast + +EXPECTED_PYTHON_TAG: Final = "cp310" +EXPECTED_ABI_TAG: Final = "abi3" +EXPECTED_PLATFORM_TAG: Final = "linux_x86_64" -def _loads_native_module(native_path: Path) -> bool: +def _dist_info_directory(member: zipfile.ZipInfo) -> str | None: + parts: Final = PurePosixPath(member.filename).parts + if not parts or not parts[0].endswith(".dist-info"): + return None + return parts[0] + + +def _wheel_metadata_tags(archive: zipfile.ZipFile, members: tuple[zipfile.ZipInfo, ...]) -> tuple[str, ...]: + if len(members) != 1: + return () + metadata: Final = BytesParser(policy=policy.default).parsebytes(archive.read(members[0])) + tags: Final = cast(list[str], metadata.get_all("Tag", [])) + return tuple(tag.strip() for tag in tags) + + +def _load_native_module(native_path: Path) -> ModuleType | None: module_spec: Final = importlib.util.spec_from_file_location("litellm.rust_bridge._native", native_path) if module_spec is None or module_spec.loader is None: - return False + return None try: native_module: Final = importlib.util.module_from_spec(module_spec) module_spec.loader.exec_module(native_module) except Exception as error: sys.stderr.write(f"native module load failed: {error}\n") - return False - return True + return None + return native_module def main() -> int: @@ -29,8 +52,39 @@ def main() -> int: return 2 wheel: Final = Path(sys.argv[1]) + wheel_tags: Final = wheel.stem.rsplit("-", maxsplit=3) + if len(wheel_tags) != 4: + sys.stderr.write(f"cannot parse wheel tags from {wheel.name}\n") + return 1 + + wheel_identity: Final = wheel_tags[0].split("-") + if len(wheel_identity) != 2 or wheel_identity[0] != "litellm" or not wheel_identity[1]: + sys.stderr.write(f"unexpected wheel identity: {wheel_tags[0]}\n") + return 1 + + expected_dist_info_directory: Final = f"{wheel_tags[0]}.dist-info" + expected_dist_info_directories: Final = frozenset((expected_dist_info_directory,)) + python_tag: Final = wheel_tags[1] + abi_tag: Final = wheel_tags[2] + platform_tag: Final = wheel_tags[3] + expanded_filename_tags: Final = frozenset( + "-".join(tag) for tag in product(python_tag.split("."), abi_tag.split("."), platform_tag.split(".")) + ) + with zipfile.ZipFile(wheel) as archive: wheel_members: Final = archive.infolist() + dist_info_directories: Final = frozenset( + directory for member in wheel_members if (directory := _dist_info_directory(member)) is not None + ) + required_dist_info_files: Final = ("METADATA", "RECORD", "WHEEL") + dist_info_file_counts: Final = { + filename: sum(member.filename == f"{expected_dist_info_directory}/{filename}" for member in wheel_members) + for filename in required_dist_info_files + } + wheel_metadata_members: Final = tuple( + member for member in wheel_members if member.filename == f"{expected_dist_info_directory}/WHEEL" + ) + wheel_metadata_tags: Final = _wheel_metadata_tags(archive, wheel_metadata_members) native_members: Final = tuple( member for member in wheel_members @@ -52,14 +106,10 @@ def main() -> int: native_path.parent.mkdir(parents=True, exist_ok=True) native_path.write_bytes(archive.read(native_member)) - wheel_tags: Final = wheel.stem.rsplit("-", maxsplit=3) - if len(wheel_tags) != 4: - sys.stderr.write(f"cannot parse wheel tags from {wheel.name}\n") - return 1 - - python_tag: Final = wheel_tags[1] - abi_tag: Final = wheel_tags[2] - platform_tag: Final = wheel_tags[3] + wheel_metadata_tags_match: Final = ( + len(wheel_metadata_tags) == len(expanded_filename_tags) + and frozenset(wheel_metadata_tags) == expanded_filename_tags + ) commit_sha: Final = os.environ.get("RELEASE_WHEEL_COMMIT_SHA", os.environ.get("GITHUB_SHA", "unknown")) rustc_version: Final = subprocess.run( ("rustc", "--version"), @@ -120,14 +170,26 @@ def main() -> int: text=True, ).stdout extension_entry_point_present: Final = "PyInit__native" in dynamic_symbols - native_module_loads: Final = _loads_native_module(native_path) + native_module: Final = _load_native_module(native_path) + native_module_loads: Final = native_module is not None + panic_test_hook_absent: Final = native_module is not None and not hasattr(native_module, "_panic_for_test") native_size_limit: Final = 20_000_000 native_size_within_limit: Final = native_member.file_size <= native_size_limit validations: Final = ( + (f"Python tag is {EXPECTED_PYTHON_TAG}", python_tag == EXPECTED_PYTHON_TAG), + (f"ABI tag is {EXPECTED_ABI_TAG}", abi_tag == EXPECTED_ABI_TAG), + (f"Platform tag is {EXPECTED_PLATFORM_TAG}", platform_tag == EXPECTED_PLATFORM_TAG), + ("Wheel dist-info directory matches the filename", dist_info_directories == expected_dist_info_directories), + ( + "Required dist-info files are present exactly once", + all(count == 1 for count in dist_info_file_counts.values()), + ), + ("Wheel metadata tags match the filename", wheel_metadata_tags_match), ("Debug sections are absent", debug_sections_absent), ("Static symbol table is absent", static_symbol_table_absent), ("Python extension entry point is present", extension_entry_point_present), ("Native module loads", native_module_loads), + ("Production module omits the panic test hook", panic_test_hook_absent), ("Native extension does not exceed 20 MB", native_size_within_limit), ("Wheel contents are valid", not unexpected_members), ) @@ -146,6 +208,26 @@ def main() -> int: sys.stderr.write(f"{native_member.filename} contains a static symbol table\n") if not extension_entry_point_present: sys.stderr.write("native extension does not export PyInit__native\n") + if python_tag != EXPECTED_PYTHON_TAG: + sys.stderr.write(f"unexpected Python tag: expected {EXPECTED_PYTHON_TAG}, found {python_tag}\n") + if abi_tag != EXPECTED_ABI_TAG: + sys.stderr.write(f"unexpected ABI tag: expected {EXPECTED_ABI_TAG}, found {abi_tag}\n") + if platform_tag != EXPECTED_PLATFORM_TAG: + sys.stderr.write(f"unexpected platform tag: expected {EXPECTED_PLATFORM_TAG}, found {platform_tag}\n") + if dist_info_directories != expected_dist_info_directories: + sys.stderr.write( + f"unexpected dist-info directories: expected {[expected_dist_info_directory]}, " + f"found {sorted(dist_info_directories)}\n" + ) + if any(count != 1 for count in dist_info_file_counts.values()): + sys.stderr.write(f"required dist-info file counts are invalid: {dist_info_file_counts}\n") + elif not wheel_metadata_tags_match: + sys.stderr.write( + f"WHEEL tags do not match filename: expected {sorted(expanded_filename_tags)}, " + f"found {sorted(wheel_metadata_tags)}\n" + ) + if native_module is not None and not panic_test_hook_absent: + sys.stderr.write("production native module exposes _panic_for_test\n") if not native_size_within_limit: sys.stderr.write(f"native extension exceeds 20 MB: {native_member.file_size / 1_000_000:.2f} MB\n") if unexpected_members: diff --git a/tests/test_litellm/test_verify_linux_native_wheel.py b/tests/test_litellm/test_verify_linux_native_wheel.py new file mode 100644 index 00000000000..86f5debfe3b --- /dev/null +++ b/tests/test_litellm/test_verify_linux_native_wheel.py @@ -0,0 +1,170 @@ +from __future__ import annotations + +import importlib.util +import subprocess +import sys +import zipfile +from collections.abc import Callable +from pathlib import Path +from types import ModuleType +from typing import Final, Protocol, cast + +import pytest + +_REPO_ROOT: Final = Path(__file__).resolve().parents[2] +_MODULE_PATH: Final = _REPO_ROOT / ".github" / "scripts" / "verify_linux_native_wheel.py" + + +class _VerifierModule(Protocol): + subprocess: ModuleType + _load_native_module: Callable[[Path], ModuleType | None] + main: Callable[[], int] + + +_SPEC: Final = importlib.util.spec_from_file_location("verify_linux_native_wheel", _MODULE_PATH) +assert _SPEC is not None and _SPEC.loader is not None +_LOADED_VERIFIER: Final = importlib.util.module_from_spec(_SPEC) +sys.modules[_SPEC.name] = _LOADED_VERIFIER +_SPEC.loader.exec_module(_LOADED_VERIFIER) +verifier: Final = cast(_VerifierModule, _LOADED_VERIFIER) + +_EXPECTED_TAG: Final = "cp310-abi3-linux_x86_64" +_NATIVE_MEMBER: Final = "litellm/rust_bridge/_native.abi3.so" +_DIST_INFO: Final = "litellm-1.100.0.dist-info" + + +def _write_wheel( + tmp_path: Path, + *, + filename_tag: str, + metadata_tags: tuple[str, ...] | None = (_EXPECTED_TAG,), + dist_info: str = _DIST_INFO, + duplicate_wheel: bool = False, +) -> Path: + wheel: Final = tmp_path / f"litellm-1.100.0-{filename_tag}.whl" + with zipfile.ZipFile(wheel, "w", compression=zipfile.ZIP_DEFLATED) as archive: + archive.writestr(_NATIVE_MEMBER, b"synthetic native extension") + archive.writestr( + f"{dist_info}/METADATA", + "Metadata-Version: 2.1\nName: litellm\nVersion: 1.100.0\n", + ) + archive.writestr( + f"{dist_info}/RECORD", + f"{_NATIVE_MEMBER},,\n{dist_info}/WHEEL,,\n", + ) + if metadata_tags is not None: + wheel_metadata: Final = ( + "Wheel-Version: 1.0\nGenerator: regression-test\nRoot-Is-Purelib: false\n" + + "".join(f"Tag: {tag}\n" for tag in metadata_tags) + ) + archive.writestr(f"{dist_info}/WHEEL", wheel_metadata) + if duplicate_wheel: + archive.writestr(f"{dist_info}/WHEEL", wheel_metadata) + return wheel + + +def _fake_subprocess_run(command: tuple[str, ...], **_: object) -> subprocess.CompletedProcess[str]: + if command == ("rustc", "--version"): + stdout = "rustc 1.98.0 (regression-test)\n" + elif "--sections" in command: + stdout = "[ 1] .text PROGBITS\n" + elif "--dyn-syms" in command: + stdout = "PyInit__native\n" + else: + raise AssertionError(f"unexpected subprocess command: {command}") + return subprocess.CompletedProcess(command, 0, stdout=stdout, stderr="") + + +def _run_verifier( + monkeypatch: pytest.MonkeyPatch, + wheel: Path, + *, + exposes_panic: bool = False, +) -> int: + native_module: Final = ModuleType("litellm.rust_bridge._native") + if exposes_panic: + setattr(native_module, "_panic_for_test", lambda: None) + + def _fake_load_native_module(_: Path) -> ModuleType: + return native_module + + monkeypatch.setattr(verifier, "_load_native_module", _fake_load_native_module) + monkeypatch.setattr(verifier.subprocess, "run", _fake_subprocess_run) + monkeypatch.setattr(sys, "argv", [str(_MODULE_PATH), str(wheel)]) + monkeypatch.setenv("GITHUB_STEP_SUMMARY", str(wheel.parent / "summary.md")) + return verifier.main() + + +def test_accepts_expected_release_wheel_tags(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + wheel: Final = _write_wheel(tmp_path, filename_tag=_EXPECTED_TAG) + + assert _run_verifier(monkeypatch, wheel) == 0 + + +def test_rejects_cp312_version_specific_wheel(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + tag: Final = "cp312-cp312-linux_x86_64" + wheel: Final = _write_wheel(tmp_path, filename_tag=tag, metadata_tags=(tag,)) + + assert _run_verifier(monkeypatch, wheel) == 1 + + +def test_rejects_non_linux_platform_tag(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + tag: Final = "cp310-abi3-win_amd64" + wheel: Final = _write_wheel(tmp_path, filename_tag=tag, metadata_tags=(tag,)) + + assert _run_verifier(monkeypatch, wheel) == 1 + + +@pytest.mark.parametrize( + "metadata_tags", + [None, ("cp312-cp312-linux_x86_64",)], + ids=["missing", "mismatched"], +) +def test_rejects_missing_or_mismatched_wheel_metadata_tag( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + metadata_tags: tuple[str, ...] | None, +) -> None: + wheel: Final = _write_wheel(tmp_path, filename_tag=_EXPECTED_TAG, metadata_tags=metadata_tags) + + assert _run_verifier(monkeypatch, wheel) == 1 + + +def test_rejects_wheel_metadata_from_wrong_dist_info_directory( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + wheel: Final = _write_wheel( + tmp_path, + filename_tag=_EXPECTED_TAG, + dist_info="decoy-1.0.0.dist-info", + ) + + assert _run_verifier(monkeypatch, wheel) == 1 + + +def test_rejects_duplicate_wheel_metadata_tags(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + wheel: Final = _write_wheel( + tmp_path, + filename_tag=_EXPECTED_TAG, + metadata_tags=(_EXPECTED_TAG, _EXPECTED_TAG), + ) + + assert _run_verifier(monkeypatch, wheel) == 1 + + +def test_rejects_duplicate_wheel_metadata_file(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + with pytest.warns(UserWarning, match="Duplicate name"): + wheel: Final = _write_wheel( + tmp_path, + filename_tag=_EXPECTED_TAG, + duplicate_wheel=True, + ) + + assert _run_verifier(monkeypatch, wheel) == 1 + + +def test_rejects_production_module_exposing_panic_hook(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + wheel: Final = _write_wheel(tmp_path, filename_tag=_EXPECTED_TAG) + + assert _run_verifier(monkeypatch, wheel, exposes_panic=True) == 1 From ef9a207ed545e740863926647af148e5195285c9 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Tue, 1 Sep 2026 05:44:35 -0700 Subject: [PATCH 85/97] fix(ci): harden release wheel reporting --- .../workflows/report-rust-release-wheel.yml | 43 +++++++++++++++++-- .github/workflows/test-rust.yml | 26 ++--------- 2 files changed, 43 insertions(+), 26 deletions(-) diff --git a/.github/workflows/report-rust-release-wheel.yml b/.github/workflows/report-rust-release-wheel.yml index 76aefd70ef6..1d93b56f77f 100644 --- a/.github/workflows/report-rust-release-wheel.yml +++ b/.github/workflows/report-rust-release-wheel.yml @@ -9,11 +9,14 @@ on: # zizmor: ignore[dangerous-triggers] reporter executes no PR code and consum permissions: {} +concurrency: + group: ${{ github.workflow }}-${{ github.event.workflow_run.pull_requests[0].number || github.event.workflow_run.id }} + cancel-in-progress: false + jobs: report-release-wheel: name: report release wheel if: >- - github.event.workflow_run.conclusion == 'success' && github.event.workflow_run.event == 'pull_request' && github.event.workflow_run.path == '.github/workflows/test-rust.yml' && github.event.workflow_run.head_repository.full_name == github.repository && @@ -21,7 +24,8 @@ jobs: runs-on: ubuntu-latest timeout-minutes: 5 permissions: - issues: write + issues: write # PR comments use the issues API + pull-requests: read # Current-head validation rejects stale workflow runs steps: - name: Link release wheel report on PR @@ -32,8 +36,19 @@ jobs: script: | const marker = process.env.COMMENT_MARKER; const workflowRun = context.payload.workflow_run; + const allowedConclusions = new Set([ + "action_required", + "cancelled", + "failure", + "neutral", + "skipped", + "stale", + "startup_failure", + "success", + "timed_out", + ]); if ( - workflowRun.conclusion !== "success" || + !allowedConclusions.has(workflowRun.conclusion) || workflowRun.event !== "pull_request" || workflowRun.path !== ".github/workflows/test-rust.yml" || workflowRun.head_repository?.full_name !== @@ -59,11 +74,15 @@ jobs: const runUrl = `${context.serverUrl}/${context.repo.owner}/${context.repo.repo}` + `/actions/runs/${runId}`; + const result = + workflowRun.conclusion === "success" + ? "successfully" + : `with \`${workflowRun.conclusion}\``; const body = [ marker, "## LiteLLM Rust workflow", "", - `Workflow completed successfully for \`${headSha}\``, + `Workflow completed ${result} for \`${headSha}\``, "", `[View workflow run](${runUrl})`, ].join("\n"); @@ -78,6 +97,22 @@ jobs: comment.user?.login === "github-actions[bot]" && comment.body?.startsWith(marker), ); + const currentPullRequest = ( + await github.rest.pulls.get({ + owner: context.repo.owner, + repo: context.repo.repo, + pull_number: pullRequestNumber, + }) + ).data; + if ( + currentPullRequest.state !== "open" || + currentPullRequest.head.repo?.full_name !== + `${context.repo.owner}/${context.repo.repo}` || + currentPullRequest.head.sha !== headSha + ) { + core.info("source workflow no longer matches the current pull request head"); + return; + } if (existing) { await github.rest.issues.updateComment({ owner: context.repo.owner, diff --git a/.github/workflows/test-rust.yml b/.github/workflows/test-rust.yml index b01e7fabe4a..77da358a2fc 100644 --- a/.github/workflows/test-rust.yml +++ b/.github/workflows/test-rust.yml @@ -108,25 +108,9 @@ jobs: rustup toolchain install stable --profile minimal rustup default stable - - name: Cache release build - uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0 - with: - path: | - ~/.cargo/registry - ~/.cargo/git - litellm-rust/target - key: ${{ runner.os }}-cargo-release-${{ hashFiles('litellm-rust/Cargo.lock') }} - restore-keys: | - ${{ runner.os }}-cargo-release- - - name: Build release wheel run: uv build --wheel --out-dir dist - - name: Verify stripped native extension - env: - RELEASE_WHEEL_COMMIT_SHA: ${{ github.event.pull_request.head.sha || github.sha }} - run: python .github/scripts/verify_linux_native_wheel.py dist/*.whl - - name: Build panic contract wheel run: >- uv build --wheel --out-dir panic-dist @@ -135,9 +119,7 @@ jobs: - name: Smoke-test native panic unwinding run: python .github/scripts/smoke_test_native_wheel.py panic-dist/*.whl - - name: Upload release wheel - uses: actions/upload-artifact@4cec3d8aa04e39d1a68397de0c4cd6fb9dce8ec1 # v4.6.1 - with: - name: litellm-release-wheel-linux-x86_64 - path: dist/*.whl - if-no-files-found: error + - name: Verify stripped native extension + env: + RELEASE_WHEEL_COMMIT_SHA: ${{ github.event.pull_request.head.sha || github.sha }} + run: python .github/scripts/verify_linux_native_wheel.py dist/*.whl From 38150dfc2c50526664125c25c6eabf07d06e07ed Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Tue, 1 Sep 2026 06:37:39 -0700 Subject: [PATCH 86/97] fix(ci): pin workflow toolchain dependencies --- .circleci/config.yml | 8 ++--- .../actions/setup-uv-with-retries/action.yml | 14 +++----- .github/workflows/test-rust.yml | 12 +++---- rust-toolchain.toml | 4 +++ .../test_circleci_rust_toolchain.py | 33 ++++++++++++++----- 5 files changed, 42 insertions(+), 29 deletions(-) create mode 100644 rust-toolchain.toml diff --git a/.circleci/config.yml b/.circleci/config.yml index 55fa9410845..dfc539fb80e 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -112,10 +112,10 @@ commands: node --version npm --version install_rust: - description: "Install pinned rustup (1.28.2) and Rust toolchain (1.97.1) with checksum verification. Adds ~/.cargo/bin to PATH. Run this before any `uv sync` or `uv build` of the workspace: the root package builds litellm-rust through maturin, and on an image without cargo maturin fetches an unpinned rustup and a floating toolchain by itself." + description: "Install pinned rustup (1.28.2) and Rust toolchain (1.98.0) with checksum verification. Adds ~/.cargo/bin to PATH. Run this before any `uv sync` or `uv build` of the workspace: the root package builds litellm-rust through maturin, and on an image without cargo maturin fetches an unpinned rustup and a floating toolchain by itself." steps: - run: - name: Install Rust (rustup 1.28.2, toolchain 1.97.1) + name: Install Rust (rustup 1.28.2, toolchain 1.98.0) command: | case "$(uname -m)" in x86_64) @@ -135,7 +135,7 @@ commands: "https://static.rust-lang.org/rustup/archive/1.28.2/${RUSTUP_TRIPLE}/rustup-init" echo "${RUSTUP_SHA256} /tmp/rustup-init" | sha256sum -c - chmod +x /tmp/rustup-init - /tmp/rustup-init -y --no-modify-path --profile minimal --default-toolchain 1.97.1 + /tmp/rustup-init -y --no-modify-path --profile minimal --default-toolchain 1.98.0 rm -f /tmp/rustup-init echo 'export PATH="$HOME/.cargo/bin:$PATH"' >> "$BASH_ENV" export PATH="$HOME/.cargo/bin:$PATH" @@ -300,7 +300,7 @@ jobs: if ($rustupActual -ne $rustupExpected) { throw "rustup installer hash mismatch: expected $rustupExpected got $rustupActual" } - & $rustupInit -y --profile minimal --default-toolchain stable + & $rustupInit -y --profile minimal --default-toolchain 1.98.0 if ($LASTEXITCODE -ne 0) { exit $LASTEXITCODE } diff --git a/.github/actions/setup-uv-with-retries/action.yml b/.github/actions/setup-uv-with-retries/action.yml index 1627038dc3d..98ff91f0283 100644 --- a/.github/actions/setup-uv-with-retries/action.yml +++ b/.github/actions/setup-uv-with-retries/action.yml @@ -1,11 +1,7 @@ name: "Set up uv with retries" description: >- - Install uv via astral-sh/setup-uv, retrying on transient failures. Even with - an exact pinned version, the action resolves the artifact URL by fetching - https://raw.githubusercontent.com/astral-sh/versions/main/v1/uv.ndjson in a - single request with no retry, timeout, or fallback, so one connection-level - network error ("fetch failed") fails the whole job before any test runs. - Retrying the full step covers the manifest fetch and the binary download. + Install uv via astral-sh/setup-uv, retrying the full setup step so manifest + resolution and binary downloads get fresh attempts after transient failures. inputs: version: @@ -18,7 +14,7 @@ runs: - name: Set up uv (attempt 1) id: attempt-1 continue-on-error: true - uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7.6.0 + uses: astral-sh/setup-uv@20cfd1bf945f4377ade1205e4dbc17946fc9a30d # v10.0.1 with: version: ${{ inputs.version }} @@ -31,7 +27,7 @@ runs: id: attempt-2 if: steps.attempt-1.outcome == 'failure' continue-on-error: true - uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7.6.0 + uses: astral-sh/setup-uv@20cfd1bf945f4377ade1205e4dbc17946fc9a30d # v10.0.1 with: version: ${{ inputs.version }} @@ -42,6 +38,6 @@ runs: - name: Set up uv (attempt 3) if: steps.attempt-2.outcome == 'failure' - uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7.6.0 + uses: astral-sh/setup-uv@20cfd1bf945f4377ade1205e4dbc17946fc9a30d # v10.0.1 with: version: ${{ inputs.version }} diff --git a/.github/workflows/test-rust.yml b/.github/workflows/test-rust.yml index 77da358a2fc..aada0fcf239 100644 --- a/.github/workflows/test-rust.yml +++ b/.github/workflows/test-rust.yml @@ -6,6 +6,7 @@ on: - "litellm-rust/**" - ".cargo/**" - "pyproject.toml" + - "rust-toolchain.toml" - ".github/scripts/smoke_test_native_wheel.py" - ".github/scripts/verify_linux_native_wheel.py" - ".github/workflows/test-rust.yml" @@ -19,6 +20,7 @@ on: - "litellm-rust/**" - ".cargo/**" - "pyproject.toml" + - "rust-toolchain.toml" - ".github/scripts/smoke_test_native_wheel.py" - ".github/scripts/verify_linux_native_wheel.py" - ".github/workflows/test-rust.yml" @@ -48,9 +50,7 @@ jobs: persist-credentials: false - name: Set up Rust - run: | - rustup toolchain install stable --profile minimal --component clippy,rustfmt - rustup default stable + run: rustup toolchain install - name: Cache Cargo registry and target uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0 @@ -59,7 +59,7 @@ jobs: ~/.cargo/registry ~/.cargo/git litellm-rust/target - key: ${{ runner.os }}-cargo-${{ hashFiles('litellm-rust/Cargo.lock') }} + key: ${{ runner.os }}-cargo-${{ hashFiles('rust-toolchain.toml', 'litellm-rust/Cargo.lock') }} restore-keys: | ${{ runner.os }}-cargo- @@ -104,9 +104,7 @@ jobs: version: "0.10.9" - name: Set up Rust - run: | - rustup toolchain install stable --profile minimal - rustup default stable + run: rustup toolchain install - name: Build release wheel run: uv build --wheel --out-dir dist diff --git a/rust-toolchain.toml b/rust-toolchain.toml new file mode 100644 index 00000000000..a1598ccbb34 --- /dev/null +++ b/rust-toolchain.toml @@ -0,0 +1,4 @@ +[toolchain] +channel = "1.98.0" +profile = "minimal" +components = ["rustfmt", "clippy"] diff --git a/tests/test_litellm/test_circleci_rust_toolchain.py b/tests/test_litellm/test_circleci_rust_toolchain.py index c35ced51e16..800ca21b95d 100644 --- a/tests/test_litellm/test_circleci_rust_toolchain.py +++ b/tests/test_litellm/test_circleci_rust_toolchain.py @@ -17,28 +17,27 @@ Two invariants are pinned here: Windows job, so the check accepts either. A new job that syncs without one falls back to the unpinned path, which is exactly the regression a static check catches at PR time and a green CI run does not. - 2. `install_rust` itself pins what it downloads: an explicit rustup version in - the URL, a verified SHA-256, and an exact toolchain version rather than a - channel name. - -The Windows job predates `install_rust` and provisions its toolchain inline, so -invariant 2 is scoped to `install_rust`; invariant 1 covers both. + 2. Both installers pin what they download: an explicit rustup version, a + verified SHA-256, and the exact toolchain in `rust-toolchain.toml`. """ from __future__ import annotations import re from pathlib import Path +from typing import Final import pytest import yaml REPO_ROOT = Path(__file__).resolve().parents[2] CONFIG = REPO_ROOT / ".circleci" / "config.yml" +TOOLCHAIN: Final = REPO_ROOT / "rust-toolchain.toml" BUILDS_WORKSPACE = re.compile(r"\buv\s+(?:sync|build)\b") RUSTUP_ARCHIVE_URL = re.compile(r"https://static\.rust-lang\.org/rustup/archive/\d+\.\d+\.\d+/") -EXACT_TOOLCHAIN = re.compile(r"--default-toolchain\s+\"?\d+\.\d+\.\d+\"?") +EXACT_TOOLCHAIN = re.compile(r"--default-toolchain\s+\"?(\d+\.\d+\.\d+)\"?") +TOOLCHAIN_CHANNEL: Final = re.compile(r'^channel = "(\d+\.\d+\.\d+)"$', re.MULTILINE) def _config() -> dict[str, object]: @@ -57,6 +56,12 @@ def _step_text(step: object) -> str: return "" +def _pinned_toolchain() -> str: + match: Final = TOOLCHAIN_CHANNEL.search(TOOLCHAIN.read_text()) + assert match is not None, "rust-toolchain.toml must pin an exact channel" + return match.group(1) + + def _without_comments(text: str) -> str: return "\n".join(line for line in text.splitlines() if not line.lstrip().startswith("#")) @@ -142,7 +147,17 @@ def test_install_rust_verifies_the_installer_checksum(install_rust_command: str) def test_install_rust_pins_an_exact_toolchain_version(install_rust_command: str) -> None: - assert EXACT_TOOLCHAIN.search(install_rust_command), ( - "install_rust must pin an exact toolchain version (e.g. 1.97.1); a channel name like " + match: Final = EXACT_TOOLCHAIN.search(install_rust_command) + assert match is not None, ( + "install_rust must pin an exact toolchain version (e.g. 1.98.0); a channel name like " "stable/beta/nightly makes the compiler drift with whatever upstream published that day" ) + assert match.group(1) == _pinned_toolchain() + + +def test_windows_installer_matches_the_repo_toolchain() -> None: + windows_steps: Final = _step_lists()["job using_litellm_on_windows"] + windows_command: Final = "\n".join(_step_text(step) for step in windows_steps) + match: Final = EXACT_TOOLCHAIN.search(windows_command) + assert match is not None + assert match.group(1) == _pinned_toolchain() From cbb8a1784db64b8433c07d09af2a114993b717ff Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Tue, 1 Sep 2026 06:49:43 -0700 Subject: [PATCH 87/97] chore(ci): extract setup-uv pin --- .github/actions/setup-uv-with-retries/action.yml | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) diff --git a/.github/actions/setup-uv-with-retries/action.yml b/.github/actions/setup-uv-with-retries/action.yml index 98ff91f0283..1627038dc3d 100644 --- a/.github/actions/setup-uv-with-retries/action.yml +++ b/.github/actions/setup-uv-with-retries/action.yml @@ -1,7 +1,11 @@ name: "Set up uv with retries" description: >- - Install uv via astral-sh/setup-uv, retrying the full setup step so manifest - resolution and binary downloads get fresh attempts after transient failures. + Install uv via astral-sh/setup-uv, retrying on transient failures. Even with + an exact pinned version, the action resolves the artifact URL by fetching + https://raw.githubusercontent.com/astral-sh/versions/main/v1/uv.ndjson in a + single request with no retry, timeout, or fallback, so one connection-level + network error ("fetch failed") fails the whole job before any test runs. + Retrying the full step covers the manifest fetch and the binary download. inputs: version: @@ -14,7 +18,7 @@ runs: - name: Set up uv (attempt 1) id: attempt-1 continue-on-error: true - uses: astral-sh/setup-uv@20cfd1bf945f4377ade1205e4dbc17946fc9a30d # v10.0.1 + uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7.6.0 with: version: ${{ inputs.version }} @@ -27,7 +31,7 @@ runs: id: attempt-2 if: steps.attempt-1.outcome == 'failure' continue-on-error: true - uses: astral-sh/setup-uv@20cfd1bf945f4377ade1205e4dbc17946fc9a30d # v10.0.1 + uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7.6.0 with: version: ${{ inputs.version }} @@ -38,6 +42,6 @@ runs: - name: Set up uv (attempt 3) if: steps.attempt-2.outcome == 'failure' - uses: astral-sh/setup-uv@20cfd1bf945f4377ade1205e4dbc17946fc9a30d # v10.0.1 + uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7.6.0 with: version: ${{ inputs.version }} From ce0c85ea691922d1e2f06bec93df10b7dc7f43ad Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Tue, 1 Sep 2026 12:06:00 -0700 Subject: [PATCH 88/97] refactor(rust): colocate native wheel contract checks --- .github/workflows/test-rust.yml | 12 ++-- .../rust_bridge}/smoke_test_native_wheel.py | 0 .../rust_bridge}/verify_linux_native_wheel.py | 57 ++++++++++++---- .../test_verify_linux_native_wheel.py | 65 +++++++------------ 4 files changed, 74 insertions(+), 60 deletions(-) rename {.github/scripts => litellm/rust_bridge}/smoke_test_native_wheel.py (100%) rename {.github/scripts => litellm/rust_bridge}/verify_linux_native_wheel.py (86%) rename tests/test_litellm/{ => rust_bridge}/test_verify_linux_native_wheel.py (62%) diff --git a/.github/workflows/test-rust.yml b/.github/workflows/test-rust.yml index aada0fcf239..271c73733a3 100644 --- a/.github/workflows/test-rust.yml +++ b/.github/workflows/test-rust.yml @@ -7,8 +7,8 @@ on: - ".cargo/**" - "pyproject.toml" - "rust-toolchain.toml" - - ".github/scripts/smoke_test_native_wheel.py" - - ".github/scripts/verify_linux_native_wheel.py" + - "litellm/rust_bridge/smoke_test_native_wheel.py" + - "litellm/rust_bridge/verify_linux_native_wheel.py" - ".github/workflows/test-rust.yml" pull_request: branches: @@ -21,8 +21,8 @@ on: - ".cargo/**" - "pyproject.toml" - "rust-toolchain.toml" - - ".github/scripts/smoke_test_native_wheel.py" - - ".github/scripts/verify_linux_native_wheel.py" + - "litellm/rust_bridge/smoke_test_native_wheel.py" + - "litellm/rust_bridge/verify_linux_native_wheel.py" - ".github/workflows/test-rust.yml" permissions: @@ -115,9 +115,9 @@ jobs: --config-setting "maturin.build-args=--features panic-test,extension-module" - name: Smoke-test native panic unwinding - run: python .github/scripts/smoke_test_native_wheel.py panic-dist/*.whl + run: python litellm/rust_bridge/smoke_test_native_wheel.py panic-dist/*.whl - name: Verify stripped native extension env: RELEASE_WHEEL_COMMIT_SHA: ${{ github.event.pull_request.head.sha || github.sha }} - run: python .github/scripts/verify_linux_native_wheel.py dist/*.whl + run: python litellm/rust_bridge/verify_linux_native_wheel.py dist/*.whl diff --git a/.github/scripts/smoke_test_native_wheel.py b/litellm/rust_bridge/smoke_test_native_wheel.py similarity index 100% rename from .github/scripts/smoke_test_native_wheel.py rename to litellm/rust_bridge/smoke_test_native_wheel.py diff --git a/.github/scripts/verify_linux_native_wheel.py b/litellm/rust_bridge/verify_linux_native_wheel.py similarity index 86% rename from .github/scripts/verify_linux_native_wheel.py rename to litellm/rust_bridge/verify_linux_native_wheel.py index 61f176ab183..783de87a8b2 100644 --- a/.github/scripts/verify_linux_native_wheel.py +++ b/litellm/rust_bridge/verify_linux_native_wheel.py @@ -6,18 +6,40 @@ import re import subprocess import sys import zipfile +from collections.abc import Callable, Mapping, Sequence from email import policy from email.parser import BytesParser from itertools import product from pathlib import Path, PurePosixPath from types import ModuleType -from typing import Final, cast +from typing import Final, Protocol, cast EXPECTED_PYTHON_TAG: Final = "cp310" EXPECTED_ABI_TAG: Final = "abi3" EXPECTED_PLATFORM_TAG: Final = "linux_x86_64" +class CommandRunner(Protocol): + def __call__( + self, + command: tuple[str, ...], + *, + check: bool, + capture_output: bool, + text: bool, + ) -> subprocess.CompletedProcess[str]: ... + + +def _run_command( + command: tuple[str, ...], + *, + check: bool, + capture_output: bool, + text: bool, +) -> subprocess.CompletedProcess[str]: + return subprocess.run(command, check=check, capture_output=capture_output, text=text) + + def _dist_info_directory(member: zipfile.ZipInfo) -> str | None: parts: Final = PurePosixPath(member.filename).parts if not parts or not parts[0].endswith(".dist-info"): @@ -46,12 +68,19 @@ def _load_native_module(native_path: Path) -> ModuleType | None: return native_module -def main() -> int: - if len(sys.argv) != 2: - sys.stderr.write(f"usage: {Path(sys.argv[0]).name} WHEEL\n") +def main( + argv: Sequence[str] | None = None, + environment: Mapping[str, str] | None = None, + load_native_module: Callable[[Path], ModuleType | None] = _load_native_module, + run_command: CommandRunner = _run_command, +) -> int: + arguments: Final = tuple(sys.argv if argv is None else argv) + resolved_environment: Final = os.environ if environment is None else environment + if len(arguments) != 2: + sys.stderr.write(f"usage: {Path(arguments[0]).name} WHEEL\n") return 2 - wheel: Final = Path(sys.argv[1]) + wheel: Final = Path(arguments[1]) wheel_tags: Final = wheel.stem.rsplit("-", maxsplit=3) if len(wheel_tags) != 4: sys.stderr.write(f"cannot parse wheel tags from {wheel.name}\n") @@ -110,8 +139,10 @@ def main() -> int: len(wheel_metadata_tags) == len(expanded_filename_tags) and frozenset(wheel_metadata_tags) == expanded_filename_tags ) - commit_sha: Final = os.environ.get("RELEASE_WHEEL_COMMIT_SHA", os.environ.get("GITHUB_SHA", "unknown")) - rustc_version: Final = subprocess.run( + commit_sha: Final = resolved_environment.get( + "RELEASE_WHEEL_COMMIT_SHA", resolved_environment.get("GITHUB_SHA", "unknown") + ) + rustc_version: Final = run_command( ("rustc", "--version"), check=True, capture_output=True, @@ -147,14 +178,14 @@ def main() -> int: "", ) ) - summary_path: Final = os.environ.get("GITHUB_STEP_SUMMARY") + summary_path: Final = resolved_environment.get("GITHUB_STEP_SUMMARY") if summary_path is None: sys.stdout.write(size_report) else: Path(summary_path).write_text(size_report) - sections: Final = subprocess.run( - ("readelf", "--sections", "--wide", native_path), + sections: Final = run_command( + ("readelf", "--sections", "--wide", str(native_path)), check=True, capture_output=True, text=True, @@ -163,14 +194,14 @@ def main() -> int: debug_sections_absent: Final = not debug_sections static_symbol_table_absent: Final = ".symtab" not in sections - dynamic_symbols: Final = subprocess.run( - ("readelf", "--dyn-syms", "--wide", native_path), + dynamic_symbols: Final = run_command( + ("readelf", "--dyn-syms", "--wide", str(native_path)), check=True, capture_output=True, text=True, ).stdout extension_entry_point_present: Final = "PyInit__native" in dynamic_symbols - native_module: Final = _load_native_module(native_path) + native_module: Final = load_native_module(native_path) native_module_loads: Final = native_module is not None panic_test_hook_absent: Final = native_module is not None and not hasattr(native_module, "_panic_for_test") native_size_limit: Final = 20_000_000 diff --git a/tests/test_litellm/test_verify_linux_native_wheel.py b/tests/test_litellm/rust_bridge/test_verify_linux_native_wheel.py similarity index 62% rename from tests/test_litellm/test_verify_linux_native_wheel.py rename to tests/test_litellm/rust_bridge/test_verify_linux_native_wheel.py index 86f5debfe3b..a4291ce0a65 100644 --- a/tests/test_litellm/test_verify_linux_native_wheel.py +++ b/tests/test_litellm/rust_bridge/test_verify_linux_native_wheel.py @@ -1,32 +1,16 @@ from __future__ import annotations -import importlib.util import subprocess -import sys import zipfile -from collections.abc import Callable from pathlib import Path from types import ModuleType -from typing import Final, Protocol, cast +from typing import Final import pytest -_REPO_ROOT: Final = Path(__file__).resolve().parents[2] -_MODULE_PATH: Final = _REPO_ROOT / ".github" / "scripts" / "verify_linux_native_wheel.py" +from litellm.rust_bridge import verify_linux_native_wheel as verifier - -class _VerifierModule(Protocol): - subprocess: ModuleType - _load_native_module: Callable[[Path], ModuleType | None] - main: Callable[[], int] - - -_SPEC: Final = importlib.util.spec_from_file_location("verify_linux_native_wheel", _MODULE_PATH) -assert _SPEC is not None and _SPEC.loader is not None -_LOADED_VERIFIER: Final = importlib.util.module_from_spec(_SPEC) -sys.modules[_SPEC.name] = _LOADED_VERIFIER -_SPEC.loader.exec_module(_LOADED_VERIFIER) -verifier: Final = cast(_VerifierModule, _LOADED_VERIFIER) +_MODULE_PATH: Final = Path(verifier.__file__) _EXPECTED_TAG: Final = "cp310-abi3-linux_x86_64" _NATIVE_MEMBER: Final = "litellm/rust_bridge/_native.abi3.so" @@ -76,7 +60,6 @@ def _fake_subprocess_run(command: tuple[str, ...], **_: object) -> subprocess.Co def _run_verifier( - monkeypatch: pytest.MonkeyPatch, wheel: Path, *, exposes_panic: bool = False, @@ -88,31 +71,33 @@ def _run_verifier( def _fake_load_native_module(_: Path) -> ModuleType: return native_module - monkeypatch.setattr(verifier, "_load_native_module", _fake_load_native_module) - monkeypatch.setattr(verifier.subprocess, "run", _fake_subprocess_run) - monkeypatch.setattr(sys, "argv", [str(_MODULE_PATH), str(wheel)]) - monkeypatch.setenv("GITHUB_STEP_SUMMARY", str(wheel.parent / "summary.md")) - return verifier.main() + environment: Final = {"GITHUB_STEP_SUMMARY": str(wheel.parent / "summary.md")} + return verifier.main( + (str(_MODULE_PATH), str(wheel)), + environment, + _fake_load_native_module, + _fake_subprocess_run, + ) -def test_accepts_expected_release_wheel_tags(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: +def test_accepts_expected_release_wheel_tags(tmp_path: Path) -> None: wheel: Final = _write_wheel(tmp_path, filename_tag=_EXPECTED_TAG) - assert _run_verifier(monkeypatch, wheel) == 0 + assert _run_verifier(wheel) == 0 -def test_rejects_cp312_version_specific_wheel(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: +def test_rejects_cp312_version_specific_wheel(tmp_path: Path) -> None: tag: Final = "cp312-cp312-linux_x86_64" wheel: Final = _write_wheel(tmp_path, filename_tag=tag, metadata_tags=(tag,)) - assert _run_verifier(monkeypatch, wheel) == 1 + assert _run_verifier(wheel) == 1 -def test_rejects_non_linux_platform_tag(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: +def test_rejects_non_linux_platform_tag(tmp_path: Path) -> None: tag: Final = "cp310-abi3-win_amd64" wheel: Final = _write_wheel(tmp_path, filename_tag=tag, metadata_tags=(tag,)) - assert _run_verifier(monkeypatch, wheel) == 1 + assert _run_verifier(wheel) == 1 @pytest.mark.parametrize( @@ -122,17 +107,15 @@ def test_rejects_non_linux_platform_tag(tmp_path: Path, monkeypatch: pytest.Monk ) def test_rejects_missing_or_mismatched_wheel_metadata_tag( tmp_path: Path, - monkeypatch: pytest.MonkeyPatch, metadata_tags: tuple[str, ...] | None, ) -> None: wheel: Final = _write_wheel(tmp_path, filename_tag=_EXPECTED_TAG, metadata_tags=metadata_tags) - assert _run_verifier(monkeypatch, wheel) == 1 + assert _run_verifier(wheel) == 1 def test_rejects_wheel_metadata_from_wrong_dist_info_directory( tmp_path: Path, - monkeypatch: pytest.MonkeyPatch, ) -> None: wheel: Final = _write_wheel( tmp_path, @@ -140,20 +123,20 @@ def test_rejects_wheel_metadata_from_wrong_dist_info_directory( dist_info="decoy-1.0.0.dist-info", ) - assert _run_verifier(monkeypatch, wheel) == 1 + assert _run_verifier(wheel) == 1 -def test_rejects_duplicate_wheel_metadata_tags(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: +def test_rejects_duplicate_wheel_metadata_tags(tmp_path: Path) -> None: wheel: Final = _write_wheel( tmp_path, filename_tag=_EXPECTED_TAG, metadata_tags=(_EXPECTED_TAG, _EXPECTED_TAG), ) - assert _run_verifier(monkeypatch, wheel) == 1 + assert _run_verifier(wheel) == 1 -def test_rejects_duplicate_wheel_metadata_file(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: +def test_rejects_duplicate_wheel_metadata_file(tmp_path: Path) -> None: with pytest.warns(UserWarning, match="Duplicate name"): wheel: Final = _write_wheel( tmp_path, @@ -161,10 +144,10 @@ def test_rejects_duplicate_wheel_metadata_file(tmp_path: Path, monkeypatch: pyte duplicate_wheel=True, ) - assert _run_verifier(monkeypatch, wheel) == 1 + assert _run_verifier(wheel) == 1 -def test_rejects_production_module_exposing_panic_hook(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: +def test_rejects_production_module_exposing_panic_hook(tmp_path: Path) -> None: wheel: Final = _write_wheel(tmp_path, filename_tag=_EXPECTED_TAG) - assert _run_verifier(monkeypatch, wheel, exposes_panic=True) == 1 + assert _run_verifier(wheel, exposes_panic=True) == 1 From 814204e21f87c65234d5167cff92c92d357a5587 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Tue, 1 Sep 2026 16:06:39 -0700 Subject: [PATCH 89/97] fix(rust): satisfy native wheel verifier lint --- .../rust_bridge/verify_linux_native_wheel.py | 93 +++++++++++-------- .../test_verify_linux_native_wheel.py | 44 +++++---- 2 files changed, 80 insertions(+), 57 deletions(-) diff --git a/litellm/rust_bridge/verify_linux_native_wheel.py b/litellm/rust_bridge/verify_linux_native_wheel.py index 783de87a8b2..899e2a211c0 100644 --- a/litellm/rust_bridge/verify_linux_native_wheel.py +++ b/litellm/rust_bridge/verify_linux_native_wheel.py @@ -7,12 +7,10 @@ import subprocess import sys import zipfile from collections.abc import Callable, Mapping, Sequence -from email import policy -from email.parser import BytesParser from itertools import product from pathlib import Path, PurePosixPath -from types import ModuleType -from typing import Final, Protocol, cast +from types import MappingProxyType, ModuleType +from typing import Final, Protocol EXPECTED_PYTHON_TAG: Final = "cp310" EXPECTED_ABI_TAG: Final = "abi3" @@ -50,9 +48,8 @@ def _dist_info_directory(member: zipfile.ZipInfo) -> str | None: def _wheel_metadata_tags(archive: zipfile.ZipFile, members: tuple[zipfile.ZipInfo, ...]) -> tuple[str, ...]: if len(members) != 1: return () - metadata: Final = BytesParser(policy=policy.default).parsebytes(archive.read(members[0])) - tags: Final = cast(list[str], metadata.get_all("Tag", [])) - return tuple(tag.strip() for tag in tags) + lines: Final = archive.read(members[0]).splitlines() + return tuple(line.removeprefix(b"Tag:").strip().decode("ascii") for line in lines if line.startswith(b"Tag:")) def _load_native_module(native_path: Path) -> ModuleType | None: @@ -62,7 +59,7 @@ def _load_native_module(native_path: Path) -> ModuleType | None: try: native_module: Final = importlib.util.module_from_spec(module_spec) module_spec.loader.exec_module(native_module) - except Exception as error: + except Exception as error: # noqa: BLE001 # native module initialization can raise arbitrary exceptions sys.stderr.write(f"native module load failed: {error}\n") return None return native_module @@ -106,10 +103,14 @@ def main( directory for member in wheel_members if (directory := _dist_info_directory(member)) is not None ) required_dist_info_files: Final = ("METADATA", "RECORD", "WHEEL") - dist_info_file_counts: Final = { - filename: sum(member.filename == f"{expected_dist_info_directory}/{filename}" for member in wheel_members) - for filename in required_dist_info_files - } + dist_info_file_counts: Final = MappingProxyType( + { + filename: sum( + member.filename == f"{expected_dist_info_directory}/{filename}" for member in wheel_members + ) + for filename in required_dist_info_files + } + ) wheel_metadata_members: Final = tuple( member for member in wheel_members if member.filename == f"{expected_dist_info_directory}/WHEEL" ) @@ -233,36 +234,46 @@ def main( if summary_path is not None: Path(summary_path).write_text(verified_report) - if debug_sections: - sys.stderr.write(f"{native_member.filename} contains debug sections: {', '.join(debug_sections)}\n") - if not static_symbol_table_absent: - sys.stderr.write(f"{native_member.filename} contains a static symbol table\n") - if not extension_entry_point_present: - sys.stderr.write("native extension does not export PyInit__native\n") - if python_tag != EXPECTED_PYTHON_TAG: - sys.stderr.write(f"unexpected Python tag: expected {EXPECTED_PYTHON_TAG}, found {python_tag}\n") - if abi_tag != EXPECTED_ABI_TAG: - sys.stderr.write(f"unexpected ABI tag: expected {EXPECTED_ABI_TAG}, found {abi_tag}\n") - if platform_tag != EXPECTED_PLATFORM_TAG: - sys.stderr.write(f"unexpected platform tag: expected {EXPECTED_PLATFORM_TAG}, found {platform_tag}\n") - if dist_info_directories != expected_dist_info_directories: - sys.stderr.write( - f"unexpected dist-info directories: expected {[expected_dist_info_directory]}, " - f"found {sorted(dist_info_directories)}\n" + invalid_dist_info_files: Final = any(count != 1 for count in dist_info_file_counts.values()) + validation_errors: Final = tuple( + message + for failed, message in ( + (bool(debug_sections), f"{native_member.filename} contains debug sections: {', '.join(debug_sections)}"), + (not static_symbol_table_absent, f"{native_member.filename} contains a static symbol table"), + (not extension_entry_point_present, "native extension does not export PyInit__native"), + ( + python_tag != EXPECTED_PYTHON_TAG, + f"unexpected Python tag: expected {EXPECTED_PYTHON_TAG}, found {python_tag}", + ), + (abi_tag != EXPECTED_ABI_TAG, f"unexpected ABI tag: expected {EXPECTED_ABI_TAG}, found {abi_tag}"), + ( + platform_tag != EXPECTED_PLATFORM_TAG, + f"unexpected platform tag: expected {EXPECTED_PLATFORM_TAG}, found {platform_tag}", + ), + ( + dist_info_directories != expected_dist_info_directories, + f"unexpected dist-info directories: expected {expected_dist_info_directory}, " + f"found {', '.join(sorted(dist_info_directories))}", + ), + (invalid_dist_info_files, f"required dist-info file counts are invalid: {dist_info_file_counts}"), + ( + not invalid_dist_info_files and not wheel_metadata_tags_match, + f"WHEEL tags do not match filename: expected {', '.join(sorted(expanded_filename_tags))}, " + f"found {', '.join(sorted(wheel_metadata_tags))}", + ), + ( + native_module is not None and not panic_test_hook_absent, + "production native module exposes _panic_for_test", + ), + ( + not native_size_within_limit, + f"native extension exceeds 20 MB: {native_member.file_size / 1_000_000:.2f} MB", + ), + (bool(unexpected_members), f"wheel contains unexpected build artifacts: {', '.join(unexpected_members)}"), ) - if any(count != 1 for count in dist_info_file_counts.values()): - sys.stderr.write(f"required dist-info file counts are invalid: {dist_info_file_counts}\n") - elif not wheel_metadata_tags_match: - sys.stderr.write( - f"WHEEL tags do not match filename: expected {sorted(expanded_filename_tags)}, " - f"found {sorted(wheel_metadata_tags)}\n" - ) - if native_module is not None and not panic_test_hook_absent: - sys.stderr.write("production native module exposes _panic_for_test\n") - if not native_size_within_limit: - sys.stderr.write(f"native extension exceeds 20 MB: {native_member.file_size / 1_000_000:.2f} MB\n") - if unexpected_members: - sys.stderr.write(f"wheel contains unexpected build artifacts: {', '.join(unexpected_members)}\n") + if failed + ) + sys.stderr.write("".join(f"{message}\n" for message in validation_errors)) return 0 if all(passed for _, passed in validations) else 1 diff --git a/tests/test_litellm/rust_bridge/test_verify_linux_native_wheel.py b/tests/test_litellm/rust_bridge/test_verify_linux_native_wheel.py index a4291ce0a65..8d0082dddc1 100644 --- a/tests/test_litellm/rust_bridge/test_verify_linux_native_wheel.py +++ b/tests/test_litellm/rust_bridge/test_verify_linux_native_wheel.py @@ -3,7 +3,7 @@ from __future__ import annotations import subprocess import zipfile from pathlib import Path -from types import ModuleType +from types import MappingProxyType, ModuleType from typing import Final import pytest @@ -47,16 +47,26 @@ def _write_wheel( return wheel -def _fake_subprocess_run(command: tuple[str, ...], **_: object) -> subprocess.CompletedProcess[str]: +def _fake_subprocess_run( + command: tuple[str, ...], + *, + check: bool, + capture_output: bool, + text: bool, +) -> subprocess.CompletedProcess[str]: + assert check and capture_output and text if command == ("rustc", "--version"): - stdout = "rustc 1.98.0 (regression-test)\n" - elif "--sections" in command: - stdout = "[ 1] .text PROGBITS\n" - elif "--dyn-syms" in command: - stdout = "PyInit__native\n" - else: - raise AssertionError(f"unexpected subprocess command: {command}") - return subprocess.CompletedProcess(command, 0, stdout=stdout, stderr="") + return subprocess.CompletedProcess(command, 0, stdout="rustc 1.98.0 (regression-test)\n", stderr="") + if "--sections" in command: + return subprocess.CompletedProcess(command, 0, stdout="[ 1] .text PROGBITS\n", stderr="") + if "--dyn-syms" in command: + return subprocess.CompletedProcess(command, 0, stdout="PyInit__native\n", stderr="") + raise AssertionError(f"unexpected subprocess command: {command}") + + +class _NativeModuleWithPanicHook(ModuleType): + def _panic_for_test(self) -> None: + return None def _run_verifier( @@ -64,14 +74,16 @@ def _run_verifier( *, exposes_panic: bool = False, ) -> int: - native_module: Final = ModuleType("litellm.rust_bridge._native") - if exposes_panic: - setattr(native_module, "_panic_for_test", lambda: None) + native_module: Final = ( + _NativeModuleWithPanicHook("litellm.rust_bridge._native") + if exposes_panic + else ModuleType("litellm.rust_bridge._native") + ) def _fake_load_native_module(_: Path) -> ModuleType: return native_module - environment: Final = {"GITHUB_STEP_SUMMARY": str(wheel.parent / "summary.md")} + environment: Final = MappingProxyType({"GITHUB_STEP_SUMMARY": str(wheel.parent / "summary.md")}) return verifier.main( (str(_MODULE_PATH), str(wheel)), environment, @@ -102,8 +114,8 @@ def test_rejects_non_linux_platform_tag(tmp_path: Path) -> None: @pytest.mark.parametrize( "metadata_tags", - [None, ("cp312-cp312-linux_x86_64",)], - ids=["missing", "mismatched"], + (None, ("cp312-cp312-linux_x86_64",)), + ids=("missing", "mismatched"), ) def test_rejects_missing_or_mismatched_wheel_metadata_tag( tmp_path: Path, From 90eadac40927315831bcc7d7ae6de57f155acbd8 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Wed, 2 Sep 2026 06:55:19 -0700 Subject: [PATCH 90/97] test(build): keep wheel checks outside package --- .../scripts}/smoke_test_native_wheel.py | 0 .../scripts}/verify_linux_native_wheel.py | 0 .github/workflows/test-rust.yml | 12 +++--- .../test_verify_linux_native_wheel.py | 38 +++++++++++++++++-- 4 files changed, 41 insertions(+), 9 deletions(-) rename {litellm/rust_bridge => .github/scripts}/smoke_test_native_wheel.py (100%) rename {litellm/rust_bridge => .github/scripts}/verify_linux_native_wheel.py (100%) diff --git a/litellm/rust_bridge/smoke_test_native_wheel.py b/.github/scripts/smoke_test_native_wheel.py similarity index 100% rename from litellm/rust_bridge/smoke_test_native_wheel.py rename to .github/scripts/smoke_test_native_wheel.py diff --git a/litellm/rust_bridge/verify_linux_native_wheel.py b/.github/scripts/verify_linux_native_wheel.py similarity index 100% rename from litellm/rust_bridge/verify_linux_native_wheel.py rename to .github/scripts/verify_linux_native_wheel.py diff --git a/.github/workflows/test-rust.yml b/.github/workflows/test-rust.yml index 271c73733a3..aada0fcf239 100644 --- a/.github/workflows/test-rust.yml +++ b/.github/workflows/test-rust.yml @@ -7,8 +7,8 @@ on: - ".cargo/**" - "pyproject.toml" - "rust-toolchain.toml" - - "litellm/rust_bridge/smoke_test_native_wheel.py" - - "litellm/rust_bridge/verify_linux_native_wheel.py" + - ".github/scripts/smoke_test_native_wheel.py" + - ".github/scripts/verify_linux_native_wheel.py" - ".github/workflows/test-rust.yml" pull_request: branches: @@ -21,8 +21,8 @@ on: - ".cargo/**" - "pyproject.toml" - "rust-toolchain.toml" - - "litellm/rust_bridge/smoke_test_native_wheel.py" - - "litellm/rust_bridge/verify_linux_native_wheel.py" + - ".github/scripts/smoke_test_native_wheel.py" + - ".github/scripts/verify_linux_native_wheel.py" - ".github/workflows/test-rust.yml" permissions: @@ -115,9 +115,9 @@ jobs: --config-setting "maturin.build-args=--features panic-test,extension-module" - name: Smoke-test native panic unwinding - run: python litellm/rust_bridge/smoke_test_native_wheel.py panic-dist/*.whl + run: python .github/scripts/smoke_test_native_wheel.py panic-dist/*.whl - name: Verify stripped native extension env: RELEASE_WHEEL_COMMIT_SHA: ${{ github.event.pull_request.head.sha || github.sha }} - run: python litellm/rust_bridge/verify_linux_native_wheel.py dist/*.whl + run: python .github/scripts/verify_linux_native_wheel.py dist/*.whl diff --git a/tests/test_litellm/rust_bridge/test_verify_linux_native_wheel.py b/tests/test_litellm/rust_bridge/test_verify_linux_native_wheel.py index 8d0082dddc1..e449d4392d8 100644 --- a/tests/test_litellm/rust_bridge/test_verify_linux_native_wheel.py +++ b/tests/test_litellm/rust_bridge/test_verify_linux_native_wheel.py @@ -1,16 +1,48 @@ from __future__ import annotations +import importlib.util import subprocess +import sys import zipfile +from collections.abc import Callable, Mapping, Sequence from pathlib import Path from types import MappingProxyType, ModuleType -from typing import Final +from typing import Final, Protocol, cast import pytest -from litellm.rust_bridge import verify_linux_native_wheel as verifier -_MODULE_PATH: Final = Path(verifier.__file__) +class _CommandRunner(Protocol): + def __call__( + self, + command: tuple[str, ...], + *, + check: bool, + capture_output: bool, + text: bool, + ) -> subprocess.CompletedProcess[str]: ... + + +class _VerifierModule(Protocol): + main: Callable[ + [ + Sequence[str] | None, + Mapping[str, str] | None, + Callable[[Path], ModuleType | None], + _CommandRunner, + ], + int, + ] + + +_REPO_ROOT: Final = Path(__file__).resolve().parents[3] +_MODULE_PATH: Final = _REPO_ROOT / ".github" / "scripts" / "verify_linux_native_wheel.py" +_SPEC: Final = importlib.util.spec_from_file_location("verify_linux_native_wheel", _MODULE_PATH) +assert _SPEC is not None and _SPEC.loader is not None +_LOADED_VERIFIER: Final = importlib.util.module_from_spec(_SPEC) +sys.modules[_SPEC.name] = _LOADED_VERIFIER +_SPEC.loader.exec_module(_LOADED_VERIFIER) +verifier: Final = cast(_VerifierModule, _LOADED_VERIFIER) _EXPECTED_TAG: Final = "cp310-abi3-linux_x86_64" _NATIVE_MEMBER: Final = "litellm/rust_bridge/_native.abi3.so" From 9de4e84feb128f2248253d9f52ae13049f2c88a5 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 31 Aug 2026 15:07:57 -0700 Subject: [PATCH 91/97] refactor(rust): extract domain-neutral Python interop --- litellm-rust/AGENTS.md | 9 +- litellm-rust/CLAUDE.md | 9 +- litellm-rust/Cargo.lock | 109 +++++++++++++++++- litellm-rust/Cargo.toml | 3 + litellm-rust/README.md | 8 +- .../PROVIDER_CODING_STANDARDS.md | 2 +- litellm-rust/crates/ai-gateway/README.md | 7 +- .../core/tests/workspace_crate_allowlist.rs | 16 ++- litellm-rust/crates/python-bridge/AGENTS.md | 2 +- litellm-rust/crates/python-bridge/CLAUDE.md | 5 +- litellm-rust/crates/python-bridge/Cargo.toml | 3 +- .../python-bridge/benches/serialization.rs | 11 +- litellm-rust/crates/python-bridge/src/gil.rs | 32 ----- litellm-rust/crates/python-bridge/src/lib.rs | 16 +-- .../python-bridge/tests/marshal_boundary.rs | 11 +- litellm-rust/crates/python-interop/AGENTS.md | 1 + litellm-rust/crates/python-interop/Cargo.toml | 15 +++ litellm-rust/crates/python-interop/src/gil.rs | 21 ++++ litellm-rust/crates/python-interop/src/lib.rs | 5 + .../src/marshal.rs | 0 .../crates/python-interop/tests/interop.rs | 44 +++++++ 21 files changed, 247 insertions(+), 82 deletions(-) delete mode 100644 litellm-rust/crates/python-bridge/src/gil.rs create mode 100644 litellm-rust/crates/python-interop/AGENTS.md create mode 100644 litellm-rust/crates/python-interop/Cargo.toml create mode 100644 litellm-rust/crates/python-interop/src/gil.rs create mode 100644 litellm-rust/crates/python-interop/src/lib.rs rename litellm-rust/crates/{python-bridge => python-interop}/src/marshal.rs (100%) create mode 100644 litellm-rust/crates/python-interop/tests/interop.rs diff --git a/litellm-rust/AGENTS.md b/litellm-rust/AGENTS.md index 36a5ad5a8f4..b8b6291283d 100644 --- a/litellm-rust/AGENTS.md +++ b/litellm-rust/AGENTS.md @@ -1,6 +1,6 @@ # AGENTS.md -litellm-rust has exactly THREE crates. A crate is a LAYER, not a route. Routes (ocr, realtime, chat) and providers (mistral, openai) are MODULES inside the layers. +litellm-rust has four crates. A crate is a layer or shared foundation, not a route. Routes (ocr, realtime, chat) and providers (mistral, openai) are modules inside the layers. ## Crates @@ -8,9 +8,10 @@ litellm-rust has exactly THREE crates. A crate is a LAYER, not a route. Routes ( |-------|------| | litellm-core | The LiteLLM SDK in Rust. One public entrypoint per top-level call (`messages::messages()`), owning types, transforms, provider resolution, auth, and the provider HTTP call. Call it, get a typed response. | | litellm-ai-gateway | The axum server (behind the `server` feature) plus the WebSocket hosts. Translates HTTP/WS to core entrypoints; owns no provider logic and no handlers. | -| litellm-python-bridge | PyO3 cdylib exposing Rust to the litellm Python SDK — marshals Python objects and calls core entrypoints. | +| litellm-python-interop | Domain-neutral PyO3 foundation for GIL handling and typed Python/Serde conversion. | +| litellm-python-bridge | PyO3 cdylib exposing LiteLLM Rust APIs to the Python SDK. Owns API registration, domain wiring, and Python exception mapping. | -Dependency direction (acyclic): litellm-core ← litellm-ai-gateway ← litellm-python-bridge. +Dependency direction is acyclic: `litellm-python-bridge` depends on the domain layers and `litellm-python-interop`; the interop foundation depends on no LiteLLM domain crate. ## Where a route lives @@ -28,7 +29,7 @@ core/src/messages/ Handlers never live in `ai-gateway`. `ocr`, `audio_transcription`, and `realtime` are still hosted there from before this rule; they move to `core` as they are touched. -Adding a crate: default to a MODULE. New crate ONLY on a real trigger — separate artifact (binary/cdylib), proc-macro, shared foundation, or publishable standalone. A new provider or route is none of these. +Adding a crate: default to a module. A new crate requires a real trigger: separate artifact (binary/cdylib), proc-macro, shared foundation, or publishable standalone. A new provider or route is none of these. Adding a crate fails crates/core/tests/workspace_crate_allowlist.rs until you update its allowlist and this file — intentional. diff --git a/litellm-rust/CLAUDE.md b/litellm-rust/CLAUDE.md index fe6ceedbb86..3dcf1853efc 100644 --- a/litellm-rust/CLAUDE.md +++ b/litellm-rust/CLAUDE.md @@ -21,12 +21,13 @@ variants of it. The test for a good abstraction is that adding the next provider is a few declarative lines, not a new file of duplicated flow. Only diverge from the base when behavior is genuinely different, and say so explicitly in the PR. -## Crates (exactly three — see AGENTS.md) +## Crates (see AGENTS.md) `litellm-core` **is** the LiteLLM SDK in Rust: it makes the LLM call. `litellm-ai-gateway` is an HTTP/WebSocket server in front of it, and -`litellm-python-bridge` exposes it to the Python SDK. A crate is a **layer**, not -a route — add modules, not crates. +`litellm-python-bridge` exposes it to the Python SDK. `litellm-python-interop` +holds domain-neutral PyO3 primitives shared by Python-facing Rust code. A crate +is a layer or shared foundation, not a route; add modules, not crates. ## Core Boundary @@ -175,7 +176,7 @@ cd litellm-rust cargo fmt --check # the ai-gateway binary + server code is behind the `server` feature cargo clippy -p litellm-ai-gateway --all-targets --features server -- -D warnings -cargo clippy -p litellm-core -p litellm-python-bridge --all-targets -- -D warnings +cargo clippy -p litellm-core -p litellm-python-interop -p litellm-python-bridge --all-targets -- -D warnings cargo test --workspace ``` diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 4388e561026..dd41cf0e84b 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -919,6 +919,12 @@ version = "0.3.33" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b231ed28831efb4a61a08580c4bc233ec56bc009f4cd8f52da2c3cb97df0c109" +[[package]] +name = "futures-timer" +version = "3.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "af43fadb8a98512d547e37b4e92e0ced13e205c061b87b4623eff01d918d6968" + [[package]] name = "futures-util" version = "0.3.33" @@ -972,6 +978,12 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "glob" +version = "0.3.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e4eba85ea1d0a966a983acd07deee566e67395d2d96b6fb39e62b5a833f1eb0b" + [[package]] name = "h2" version = "0.3.27" @@ -1432,14 +1444,24 @@ dependencies = [ "criterion", "litellm-ai-gateway", "litellm-core", + "litellm-python-interop", "pyo3", "pyo3-async-runtimes", - "pythonize", - "serde", "serde_json", "tokio", ] +[[package]] +name = "litellm-python-interop" +version = "0.1.0" +dependencies = [ + "pyo3", + "pythonize", + "rstest", + "serde", + "serde_json", +] + [[package]] name = "litemap" version = "0.8.2" @@ -1627,6 +1649,15 @@ dependencies = [ "zerocopy", ] +[[package]] +name = "proc-macro-crate" +version = "3.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e67ba7e9b2b56446f1d419b1d807906278ffa1a658a8a5d8a39dcb1f5a78614f" +dependencies = [ + "toml_edit", +] + [[package]] name = "proc-macro2" version = "1.0.107" @@ -1899,6 +1930,12 @@ version = "0.8.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d6f6ff9a378485b298a5286656da665ba74413d36db0979633275d2e708145d4" +[[package]] +name = "relative-path" +version = "1.9.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ba39f3699c378cd8970968dcbff9c43159ea4cfbd88d43c00b22f2ef10a435d2" + [[package]] name = "reqwest" version = "0.12.28" @@ -1956,6 +1993,35 @@ dependencies = [ "windows-sys 0.52.0", ] +[[package]] +name = "rstest" +version = "0.26.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f5a3193c063baaa2a95a33f03035c8a72b83d97a54916055ba22d35ed3839d49" +dependencies = [ + "futures-timer", + "futures-util", + "rstest_macros", +] + +[[package]] +name = "rstest_macros" +version = "0.26.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9c845311f0ff7951c5506121a9ad75aec44d083c31583b2ea5a30bcb0b0abba0" +dependencies = [ + "cfg-if", + "glob", + "proc-macro-crate", + "proc-macro2", + "quote", + "regex", + "relative-path", + "rustc_version", + "syn 2.0.119", + "unicode-ident", +] + [[package]] name = "rustc-hash" version = "2.1.3" @@ -2488,6 +2554,36 @@ dependencies = [ "tokio", ] +[[package]] +name = "toml_datetime" +version = "1.1.1+spec-1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3165f65f62e28e0115a00b2ebdd37eb6f3b641855f9d636d3cd4103767159ad7" +dependencies = [ + "serde_core", +] + +[[package]] +name = "toml_edit" +version = "0.25.13+spec-1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6975367e4d2ef766d86af01ffad14b622fecc8d4357a998fbc4deb6e9bacaf9b" +dependencies = [ + "indexmap", + "toml_datetime", + "toml_parser", + "winnow", +] + +[[package]] +name = "toml_parser" +version = "1.1.3+spec-1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1d38ac1cf9b95face32296c0a3ede1fdc270627c9d9c02a7274dd6d960dc4d56" +dependencies = [ + "winnow", +] + [[package]] name = "tower" version = "0.5.3" @@ -2903,6 +2999,15 @@ version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec" +[[package]] +name = "winnow" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "23b97319f7b8343df12cc98938e5c3eb436064524c8d2b4e30a1d3a36eecdf81" +dependencies = [ + "memchr", +] + [[package]] name = "writeable" version = "0.6.3" diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index c17a0605fc7..c447d915abe 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -2,6 +2,7 @@ members = [ "crates/core", "crates/ai-gateway", + "crates/python-interop", "crates/python-bridge", ] resolver = "2" @@ -15,12 +16,14 @@ repository = "https://github.com/BerriAI/litellm" [workspace.dependencies] litellm-core = { path = "crates/core" } litellm-ai-gateway = { path = "crates/ai-gateway", default-features = false } +litellm-python-interop = { path = "crates/python-interop" } axum = "0.7" pyo3 = "0.29.0" pyo3-async-runtimes = { version = "0.29.0", features = ["tokio-runtime"] } pythonize = "0.29.0" rand = "0.8" reqwest = { version = "0.12", default-features = false, features = ["blocking", "json", "rustls-tls", "http2", "stream"] } +rstest = "0.26.1" serde = { version = "1.0", features = ["derive"] } serde_json = "1.0" sha2 = "0.10" diff --git a/litellm-rust/README.md b/litellm-rust/README.md index bcccf93300b..a0d79c6f0a5 100644 --- a/litellm-rust/README.md +++ b/litellm-rust/README.md @@ -26,9 +26,10 @@ coverage and production evidence. |-------|------| | litellm-core | The SDK. Per-route entrypoints (`messages::messages()`), types, provider transforms (modules under `providers/`), provider resolution, auth, the provider HTTP call, and the router. | | litellm-ai-gateway | The axum server (behind the `server` feature) and WebSocket hosts. Translates HTTP/WS to core entrypoints; no provider handlers. | -| litellm-python-bridge | PyO3 cdylib exposing Rust to the litellm Python SDK — marshals Python objects and calls core entrypoints. | +| litellm-python-interop | Domain-neutral PyO3 foundation for GIL handling and typed Python/Serde conversion. | +| litellm-python-bridge | PyO3 cdylib exposing LiteLLM Rust APIs to the Python SDK. Owns API registration, domain wiring, and Python exception mapping. | -Dependency direction (acyclic): litellm-core ← litellm-ai-gateway ← litellm-python-bridge. +Dependency direction is acyclic: `litellm-python-bridge` depends on the domain layers and `litellm-python-interop`; the interop foundation depends on no LiteLLM domain crate. ## Layout @@ -38,7 +39,8 @@ crates/ src/messages/ mod.rs (entrypoint), types, transformation, prepare, handler, client src/providers/anthropic/messages/transformation.rs ai-gateway/ Axum server + WebSocket hosts; calls core entrypoints. - python-bridge/ PyO3 bridge for Python LiteLLM. + python-interop/ Domain-neutral PyO3 conversion and GIL primitives. + python-bridge/ PyO3 API adapter for Python LiteLLM. ``` The folder shape follows the Python provider tree: diff --git a/litellm-rust/crates/CODING_STANDARDS/PROVIDER_CODING_STANDARDS.md b/litellm-rust/crates/CODING_STANDARDS/PROVIDER_CODING_STANDARDS.md index a1860d8a9c9..4a689cb9579 100644 --- a/litellm-rust/crates/CODING_STANDARDS/PROVIDER_CODING_STANDARDS.md +++ b/litellm-rust/crates/CODING_STANDARDS/PROVIDER_CODING_STANDARDS.md @@ -54,6 +54,6 @@ Rules for adding or changing an LLM provider/route in `litellm-rust`. `messages` cd litellm-rust cargo fmt --check cargo clippy -p litellm-ai-gateway --all-targets --features server -- -D warnings - cargo clippy -p litellm-core -p litellm-python-bridge --all-targets -- -D warnings + cargo clippy -p litellm-core -p litellm-python-interop -p litellm-python-bridge --all-targets -- -D warnings cargo test --workspace ``` diff --git a/litellm-rust/crates/ai-gateway/README.md b/litellm-rust/crates/ai-gateway/README.md index 7a6c620ee84..5cbb47220be 100644 --- a/litellm-rust/crates/ai-gateway/README.md +++ b/litellm-rust/crates/ai-gateway/README.md @@ -6,15 +6,16 @@ dials OpenAI upstream, and splices the two sockets frame-by-frame. ## Crates -`litellm-rust` is exactly three crates (a crate is a **layer**, not a route): +`litellm-rust` has four crates. A crate is a layer or shared foundation, not a route: | Crate | Role | |-------|------| | litellm-core | The LiteLLM SDK in Rust — per-route entrypoints (`messages::messages()`) that resolve the provider, transform, and make the call; plus types, provider transforms, and the router. | | litellm-ai-gateway | The Axum server (behind the `server` feature) and WebSocket hosts. Translates HTTP/WS to core entrypoints; no provider handlers. | -| litellm-python-bridge | PyO3 cdylib exposing Rust to the litellm Python SDK — marshals Python objects and calls core entrypoints. | +| litellm-python-interop | Domain-neutral PyO3 foundation for GIL handling and typed Python/Serde conversion. | +| litellm-python-bridge | PyO3 cdylib exposing LiteLLM Rust APIs to the Python SDK. | -Dependency direction (acyclic): litellm-core ← litellm-ai-gateway ← litellm-python-bridge. +Dependency direction is acyclic: `litellm-python-bridge` depends on the domain layers and `litellm-python-interop`; the interop foundation depends on no LiteLLM domain crate. - **Client endpoint:** `wss:///v1/realtime?model=` (WebSocket) - **Auth:** `Authorization: Bearer $LITELLM_MASTER_KEY` (fails closed if unset) diff --git a/litellm-rust/crates/core/tests/workspace_crate_allowlist.rs b/litellm-rust/crates/core/tests/workspace_crate_allowlist.rs index 656ba033b62..8a8a5ea263a 100644 --- a/litellm-rust/crates/core/tests/workspace_crate_allowlist.rs +++ b/litellm-rust/crates/core/tests/workspace_crate_allowlist.rs @@ -1,7 +1,8 @@ -//! Enforcement: the litellm-rust workspace has exactly three crates. +//! Enforcement: the litellm-rust workspace has exactly four crates. //! -//! `core` (pure translation), `ai-gateway` (routes + all network I/O), and -//! `python-bridge` (the PyO3 cdylib). Adding or removing a crate must be a +//! `core` (the Rust SDK), `ai-gateway` (the HTTP/WebSocket host), +//! `python-interop` (domain-neutral PyO3 primitives), and `python-bridge` (the +//! PyO3 cdylib). Adding or removing a crate must be a //! deliberate act: this test fails until the allowlist here is updated, forcing //! whoever changes the crate set to justify the new crate per the rule that a //! crate is a layer needing independent compilation / its own deps / a separate @@ -16,10 +17,15 @@ use std::path::{Path, PathBuf}; /// The one true crate set. Update BOTH this and `litellm-rust/AGENTS.md` when the /// workspace legitimately gains or loses a crate. -const EXPECTED_MEMBERS: &[&str] = &["crates/core", "crates/ai-gateway", "crates/python-bridge"]; +const EXPECTED_MEMBERS: &[&str] = &[ + "crates/core", + "crates/ai-gateway", + "crates/python-interop", + "crates/python-bridge", +]; /// The crate subdirectory names that must exist under `crates/`. -const EXPECTED_CRATE_DIRS: &[&str] = &["core", "ai-gateway", "python-bridge"]; +const EXPECTED_CRATE_DIRS: &[&str] = &["core", "ai-gateway", "python-interop", "python-bridge"]; const MISMATCH: &str = "litellm-rust crate set changed — update this allowlist AND litellm-rust/AGENTS.md, and justify the crate per the rule (crate = layer needing independent compilation / its own deps / a separate artifact)."; diff --git a/litellm-rust/crates/python-bridge/AGENTS.md b/litellm-rust/crates/python-bridge/AGENTS.md index ad3cddfa5fd..42282ca4da4 100644 --- a/litellm-rust/crates/python-bridge/AGENTS.md +++ b/litellm-rust/crates/python-bridge/AGENTS.md @@ -1,3 +1,3 @@ -litellm-python-bridge is the PyO3 cdylib that exposes Rust to the litellm Python SDK — a thin adapter (Python objects → Rust calls → Python results) over the litellm-core route entrypoints (e.g. `litellm_core::messages::messages`). +litellm-python-bridge is the PyO3 cdylib that exposes LiteLLM Rust APIs to the Python SDK. Keep API registration, domain dependency wiring, request assembly, and Python exception mapping here. Put domain-neutral Python/Serde conversion and GIL primitives in litellm-python-interop. Keep it thin: no business logic, no transforms, no I/O orchestration — just marshal in/out and call the core entrypoint. diff --git a/litellm-rust/crates/python-bridge/CLAUDE.md b/litellm-rust/crates/python-bridge/CLAUDE.md index 3ce8b8c639a..d25ae5a8130 100644 --- a/litellm-rust/crates/python-bridge/CLAUDE.md +++ b/litellm-rust/crates/python-bridge/CLAUDE.md @@ -5,8 +5,9 @@ Rules for `litellm-rust/crates/python-bridge`. ## Responsibility `python-bridge` is the PyO3 boundary between Python LiteLLM and Rust transforms. -Keep this crate thin. It adapts Python objects to Rust payloads and returns -Python-compatible dictionaries. +Keep this crate thin. It exposes LiteLLM Rust APIs, assembles domain requests, +maps domain errors to Python exceptions, and delegates generic conversion and +GIL handling to `litellm-python-interop`. ## Bridge Shape diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index b1fdfd7677a..498003de149 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -18,10 +18,9 @@ panic-test = [] [dependencies] litellm-core = { workspace = true, features = ["bedrock-auth"] } litellm-ai-gateway = { workspace = true, default-features = false } +litellm-python-interop.workspace = true pyo3.workspace = true pyo3-async-runtimes.workspace = true -pythonize.workspace = true -serde.workspace = true serde_json.workspace = true tokio.workspace = true diff --git a/litellm-rust/crates/python-bridge/benches/serialization.rs b/litellm-rust/crates/python-bridge/benches/serialization.rs index 8a90cf667d0..0b9436d0cb7 100644 --- a/litellm-rust/crates/python-bridge/benches/serialization.rs +++ b/litellm-rust/crates/python-bridge/benches/serialization.rs @@ -2,6 +2,7 @@ use std::hint::black_box; use std::time::Duration; use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main}; +use litellm_python_interop::{from_py, to_py}; use pyo3::prelude::*; use pyo3::types::PyDict; use serde_json::{Value, json}; @@ -25,7 +26,7 @@ fn former_json_roundtrip_from_py(py: Python<'_>, value: &Bound<'_, PyAny>) -> Va } fn pythonize_from_py(value: &Bound<'_, PyAny>) -> Value { - pythonize::depythonize(value).expect("payload should depythonize") + from_py(value).expect("payload should depythonize") } fn former_json_roundtrip_to_py(py: Python<'_>, value: &Value) -> Py { @@ -37,12 +38,10 @@ fn former_json_roundtrip_to_py(py: Python<'_>, value: &Value) -> Py { } fn pythonize_to_py(py: Python<'_>, value: &Value) -> Py { - pythonize::pythonize(py, value) - .expect("response should pythonize") - .unbind() + to_py(py, value).expect("response should pythonize") } -fn serialization(c: &mut Criterion) { +fn bridge_serialization(c: &mut Criterion) { Python::initialize(); Python::attach(|py| { for &(label, payload_bytes) in PAYLOAD_SIZES { @@ -98,6 +97,6 @@ criterion_group! { .sample_size(20) .warm_up_time(Duration::from_secs(1)) .measurement_time(Duration::from_secs(4)); - targets = serialization + targets = bridge_serialization } criterion_main!(benches); diff --git a/litellm-rust/crates/python-bridge/src/gil.rs b/litellm-rust/crates/python-bridge/src/gil.rs deleted file mode 100644 index e887c8ec1e3..00000000000 --- a/litellm-rust/crates/python-bridge/src/gil.rs +++ /dev/null @@ -1,32 +0,0 @@ -//! GIL accounting. -//! -//! A single chokepoint for releasing the GIL around blocking work. Every -//! blocking call in the bridge goes through [`release_gil`] instead of calling -//! `Python::detach` directly, so the release count stays accurate and we -//! have one place to extend later (timing histograms, per-call labels, etc.). - -use std::sync::atomic::{AtomicU64, Ordering}; - -use pyo3::prelude::*; - -/// Number of times the bridge has released the GIL since process start. -static GIL_RELEASES: AtomicU64 = AtomicU64::new(0); - -/// Release the GIL around `f`, recording the release. -/// -/// `f` must not touch any Python state — that is what makes releasing the GIL -/// safe. Returning the value back to Python re-acquires the GIL at the call -/// site, after `f` has finished. -pub fn release_gil(py: Python<'_>, f: F) -> T -where - F: FnOnce() -> T + Send, - T: Send, -{ - GIL_RELEASES.fetch_add(1, Ordering::Relaxed); - py.detach(f) -} - -/// Total GIL releases performed by the bridge so far. -pub fn release_count() -> u64 { - GIL_RELEASES.load(Ordering::Relaxed) -} diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index 18e0b05cbb5..746e0770f9b 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -13,16 +13,12 @@ use litellm_core::chat_completions::{ use litellm_core::error::CoreError; use litellm_core::messages::messages as run_messages; use litellm_core::messages::types::{AnthropicMessagesResponse, MessagesRequest}; +use litellm_python_interop::{from_py, release_count, release_gil, to_py}; use pyo3::exceptions::{PyRuntimeError, PyValueError}; use pyo3::prelude::*; use pyo3::types::{PyAny, PyDict}; use serde_json::{Map, Value}; -mod gil; -mod marshal; - -use marshal::{from_py, to_py}; - pyo3::create_exception!( _native, RustBridgeDeclined, @@ -230,7 +226,7 @@ fn ocr( timeout_seconds, )?; - let result = gil::release_gil(py, || { + let result = release_gil(py, || { pyo3_async_runtimes::tokio::get_runtime().block_on(run_ocr(OcrRequest { model: &model, document, @@ -318,7 +314,7 @@ fn transcription( }; let optional_params = optional_object_to_map(py, "optional_params", optional_params)?; let timeout = optional_timeout(timeout_seconds); - let result = gil::release_gil(py, || { + let result = release_gil(py, || { pyo3_async_runtimes::tokio::get_runtime().block_on(run_audio_transcription( AudioTranscriptionRequest { model: &model, @@ -419,7 +415,7 @@ fn messages( let (body, extra_headers, timeout) = marshal_messages_inputs(py, body, extra_headers, timeout_seconds)?; - let result = gil::release_gil(py, || { + let result = release_gil(py, || { pyo3_async_runtimes::tokio::get_runtime().block_on(run_messages(MessagesRequest { model: &model, body, @@ -546,7 +542,7 @@ fn chat_completions( timeout_seconds, )?; - let result = gil::release_gil(py, || { + let result = release_gil(py, || { pyo3_async_runtimes::tokio::get_runtime().block_on(run_chat_completions( ChatCompletionsRequest { model: &model, @@ -610,7 +606,7 @@ fn achat_completions( #[pyfunction] fn gil_stats(py: Python<'_>) -> PyResult> { let stats = PyDict::new(py); - stats.set_item("releases", gil::release_count())?; + stats.set_item("releases", release_count())?; Ok(stats.into_any().unbind()) } diff --git a/litellm-rust/crates/python-bridge/tests/marshal_boundary.rs b/litellm-rust/crates/python-bridge/tests/marshal_boundary.rs index 6a6ede22e85..d397d20b9fd 100644 --- a/litellm-rust/crates/python-bridge/tests/marshal_boundary.rs +++ b/litellm-rust/crates/python-bridge/tests/marshal_boundary.rs @@ -1,7 +1,7 @@ use std::fs; use std::path::{Path, PathBuf}; -const DISALLOWED_OUTSIDE_MARSHAL: &[&str] = &[ +const DISALLOWED_OUTSIDE_INTEROP: &[&str] = &[ "py.import(\"json\")", "pythonize::", "serde_json::to_string", @@ -33,18 +33,15 @@ fn rust_sources(directory: &Path) -> Vec { } #[test] -fn serialization_is_centralized_in_marshal_module() { +fn serialization_uses_the_interop_boundary() { let root = source_root(); for path in rust_sources(&root) { - if path == root.join("marshal.rs") { - continue; - } let source = fs::read_to_string(&path).expect("bridge source should be readable"); - for disallowed in DISALLOWED_OUTSIDE_MARSHAL { + for disallowed in DISALLOWED_OUTSIDE_INTEROP { assert!( !source.contains(disallowed), - "{} bypasses the typed marshal module with `{disallowed}`", + "{} bypasses litellm-python-interop with `{disallowed}`", path.display() ); } diff --git a/litellm-rust/crates/python-interop/AGENTS.md b/litellm-rust/crates/python-interop/AGENTS.md new file mode 100644 index 00000000000..d1d61e5dfa0 --- /dev/null +++ b/litellm-rust/crates/python-interop/AGENTS.md @@ -0,0 +1 @@ +litellm-python-interop is the domain-neutral PyO3 foundation. Keep generic Python/Serde conversion and interpreter primitives here. Do not add LiteLLM domain crates, route types, API registration, or cdylib build features. diff --git a/litellm-rust/crates/python-interop/Cargo.toml b/litellm-rust/crates/python-interop/Cargo.toml new file mode 100644 index 00000000000..9da6af6e2e2 --- /dev/null +++ b/litellm-rust/crates/python-interop/Cargo.toml @@ -0,0 +1,15 @@ +[package] +name = "litellm-python-interop" +version = "0.1.0" +edition.workspace = true +license.workspace = true +repository.workspace = true + +[dependencies] +pyo3.workspace = true +pythonize.workspace = true +serde.workspace = true + +[dev-dependencies] +rstest.workspace = true +serde_json.workspace = true diff --git a/litellm-rust/crates/python-interop/src/gil.rs b/litellm-rust/crates/python-interop/src/gil.rs new file mode 100644 index 00000000000..04b966a6002 --- /dev/null +++ b/litellm-rust/crates/python-interop/src/gil.rs @@ -0,0 +1,21 @@ +use std::sync::atomic::{AtomicU64, Ordering}; + +use pyo3::prelude::*; + +static GIL_RELEASES: AtomicU64 = AtomicU64::new(0); + +/// Runs work detached from the interpreter and records the release. +/// +/// `f` must not access Python state while the interpreter is detached. +pub fn release_gil(py: Python<'_>, f: F) -> T +where + F: FnOnce() -> T + Send, + T: Send, +{ + GIL_RELEASES.fetch_add(1, Ordering::Relaxed); + py.detach(f) +} + +pub fn release_count() -> u64 { + GIL_RELEASES.load(Ordering::Relaxed) +} diff --git a/litellm-rust/crates/python-interop/src/lib.rs b/litellm-rust/crates/python-interop/src/lib.rs new file mode 100644 index 00000000000..df2bd260fdb --- /dev/null +++ b/litellm-rust/crates/python-interop/src/lib.rs @@ -0,0 +1,5 @@ +mod gil; +mod marshal; + +pub use gil::{release_count, release_gil}; +pub use marshal::{from_py, to_py}; diff --git a/litellm-rust/crates/python-bridge/src/marshal.rs b/litellm-rust/crates/python-interop/src/marshal.rs similarity index 100% rename from litellm-rust/crates/python-bridge/src/marshal.rs rename to litellm-rust/crates/python-interop/src/marshal.rs diff --git a/litellm-rust/crates/python-interop/tests/interop.rs b/litellm-rust/crates/python-interop/tests/interop.rs new file mode 100644 index 00000000000..9c456dcb938 --- /dev/null +++ b/litellm-rust/crates/python-interop/tests/interop.rs @@ -0,0 +1,44 @@ +use pyo3::Python; +use rstest::{fixture, rstest}; +use serde_json::{Value, json}; + +use litellm_python_interop::{from_py, release_count, release_gil, to_py}; + +struct InitializedPython; + +impl InitializedPython { + fn attach(&self, f: F) -> R + where + F: for<'py> FnOnce(Python<'py>) -> R, + { + Python::attach(f) + } +} + +#[fixture] +#[once] +fn initialized_python() -> InitializedPython { + Python::initialize(); + InitializedPython +} + +#[rstest] +fn serde_values_round_trip_through_python(#[from(initialized_python)] python: &InitializedPython) { + python.attach(|py| { + let expected = json!({"model": "test", "items": [1, true, null]}); + let python_value = to_py(py, &expected).expect("value should convert to Python"); + let actual: Value = + from_py(python_value.bind(py)).expect("Python value should convert to serde"); + + assert_eq!(actual, expected); + }); +} + +#[rstest] +fn release_gil_runs_work_and_records_it(#[from(initialized_python)] python: &InitializedPython) { + let before = release_count(); + let result = python.attach(|py| release_gil(py, || 42)); + + assert_eq!(result, 42); + assert_eq!(release_count(), before + 1); +} From 518a2a70f159c1e118d1ea0d387c7da8ba87d642 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Wed, 2 Sep 2026 05:47:01 -0700 Subject: [PATCH 92/97] refactor(rust): standardize the core Error type --- .../src/audio_transcription/common_utils.rs | 10 +-- .../src/audio_transcription/handler.rs | 31 ++++---- .../src/audio_transcription/hooks.rs | 58 +++++++------- .../ai-gateway/src/audio_transcription/mod.rs | 4 +- .../crates/ai-gateway/src/io/realtime.rs | 43 +++++------ .../crates/ai-gateway/src/io/realtime_pool.rs | 4 +- .../crates/ai-gateway/src/io/responses_ws.rs | 76 +++++++++---------- .../crates/ai-gateway/src/ocr/common_utils.rs | 73 +++++++++--------- .../crates/ai-gateway/src/ocr/handler.rs | 15 ++-- .../crates/ai-gateway/src/ocr/hooks.rs | 65 ++++++++-------- litellm-rust/crates/ai-gateway/src/ocr/mod.rs | 4 +- .../crates/ai-gateway/src/ocr/tests.rs | 8 +- .../crates/ai-gateway/src/python/config.rs | 12 ++- .../ai-gateway/src/routes/messages/mod.rs | 38 +++++----- .../ai-gateway/src/routes/messages/service.rs | 16 ++-- .../ai-gateway/src/routes/realtime/service.rs | 11 ++- .../src/routes/responses/service.rs | 12 +-- .../src/audio_transcription/transformation.rs | 11 ++- .../crates/core/src/call_lifecycle/mod.rs | 34 ++++----- .../core/src/chat_completions/common_utils.rs | 7 +- .../core/src/chat_completions/handler.rs | 30 ++++---- .../crates/core/src/chat_completions/mod.rs | 5 +- .../core/src/chat_completions/prepare.rs | 21 +++-- .../crates/core/src/chat_completions/tests.rs | 40 +++++----- .../src/chat_completions/transformation.rs | 11 ++- litellm-rust/crates/core/src/error.rs | 8 +- litellm-rust/crates/core/src/http_utils.rs | 8 +- litellm-rust/crates/core/src/lib.rs | 2 +- .../crates/core/src/messages/common_utils.rs | 7 +- .../crates/core/src/messages/handler.rs | 25 +++--- litellm-rust/crates/core/src/messages/mod.rs | 7 +- .../crates/core/src/messages/prepare.rs | 12 +-- .../crates/core/src/messages/tests.rs | 10 +-- .../core/src/messages/transformation.rs | 11 ++- .../crates/core/src/ocr/transformation.rs | 11 ++- .../anthropic/chat_completions/tests.rs | 16 ++-- .../chat_completions/transformation.rs | 25 +++--- .../anthropic/messages/transformation.rs | 12 +-- .../azure_ai/messages/transformation.rs | 22 +++--- .../providers/azure_ai/ocr/transformation.rs | 62 +++++++-------- .../providers/bedrock/audio_transcription.rs | 20 ++--- .../core/src/providers/bedrock/aws_base.rs | 38 +++++----- .../bedrock/chat_completions/tests.rs | 14 ++-- .../chat_completions/transformation.rs | 23 +++--- .../providers/mistral/ocr/transformation.rs | 28 +++---- .../openai/realtime/transformation.rs | 10 +-- .../openai/responses/transformation.rs | 6 +- .../providers/vertex_ai/ocr/transformation.rs | 50 ++++++------ .../core/src/realtime/transformation.rs | 6 +- .../core/src/responses/instrumentation.rs | 8 +- .../crates/core/src/responses/websocket.rs | 6 +- litellm-rust/crates/python-bridge/src/lib.rs | 36 ++++----- 52 files changed, 544 insertions(+), 578 deletions(-) diff --git a/litellm-rust/crates/ai-gateway/src/audio_transcription/common_utils.rs b/litellm-rust/crates/ai-gateway/src/audio_transcription/common_utils.rs index 270d5c2d97a..140bc8aeea8 100644 --- a/litellm-rust/crates/ai-gateway/src/audio_transcription/common_utils.rs +++ b/litellm-rust/crates/ai-gateway/src/audio_transcription/common_utils.rs @@ -1,10 +1,8 @@ -use std::collections::BTreeMap; - -use litellm_core::CoreResult; use litellm_core::audio_transcription::transformation::AudioTranscriptionProviderConfig; -use litellm_core::error::CoreError; +use litellm_core::error::Error; use litellm_core::providers::bedrock::audio_transcription::BEDROCK_AUDIO_TRANSCRIPTION_CONFIG; use serde_json::{Map, Value}; +use std::collections::BTreeMap; pub(super) fn audio_transcription_provider_config( provider: &str, @@ -17,7 +15,7 @@ pub(super) fn audio_transcription_provider_config( pub(super) fn string_headers( headers: Option>, -) -> CoreResult> { +) -> Result, Error> { headers .unwrap_or_default() .into_iter() @@ -26,7 +24,7 @@ pub(super) fn string_headers( .as_str() .map(|value| (key.clone(), value.to_string())) .ok_or_else(|| { - CoreError::InvalidRequest(format!( + Error::InvalidRequest(format!( "audio transcription extra_headers.{key} must be a string" )) }) diff --git a/litellm-rust/crates/ai-gateway/src/audio_transcription/handler.rs b/litellm-rust/crates/ai-gateway/src/audio_transcription/handler.rs index 33c13550f58..1bdd4ae72a2 100644 --- a/litellm-rust/crates/ai-gateway/src/audio_transcription/handler.rs +++ b/litellm-rust/crates/ai-gateway/src/audio_transcription/handler.rs @@ -1,11 +1,9 @@ -use std::time::SystemTime; - -use litellm_core::CoreResult; use litellm_core::audio_transcription::transformation::AudioTranscriptionAuth; -use litellm_core::error::CoreError; +use litellm_core::error::Error; use litellm_core::providers::bedrock::audio_transcription::aws_auth_config; use litellm_core::providers::bedrock::aws_base::{resolve_credentials, sign_bedrock_post}; use serde_json::Value; +use std::time::SystemTime; use super::common_utils::truncate_error_body; use super::types::ProviderAudioTranscriptionRequest; @@ -13,10 +11,9 @@ use crate::client::http_client; pub(crate) async fn execute_audio_transcription_provider_call( request: ProviderAudioTranscriptionRequest, -) -> CoreResult { - let body = serde_json::to_vec(&request.body).map_err(|error| { - CoreError::InvalidRequest(format!("invalid audio request body: {error}")) - })?; +) -> Result { + let body = serde_json::to_vec(&request.body) + .map_err(|error| Error::InvalidRequest(format!("invalid audio request body: {error}")))?; let mut request_builder = http_client().post(&request.url).body(body.clone()); for (key, value) in &request.upstream_headers { request_builder = request_builder.header(key, value); @@ -27,21 +24,20 @@ pub(crate) async fn execute_audio_transcription_provider_call( let response = request_builder .send() .await - .map_err(|error| CoreError::Network(error.to_string()))?; + .map_err(|error| Error::Network(error.to_string()))?; let status = response.status(); let text = response .text() .await - .map_err(|error| CoreError::Network(error.to_string()))?; + .map_err(|error| Error::Network(error.to_string()))?; if !status.is_success() { - return Err(CoreError::Http { + return Err(Error::Http { status: status.as_u16(), body: truncate_error_body(&text), }); } - let response_json: Value = serde_json::from_str(&text).map_err(|error| { - CoreError::InvalidResponse(format!("invalid audio response JSON: {error}")) - })?; + let response_json: Value = serde_json::from_str(&text) + .map_err(|error| Error::InvalidResponse(format!("invalid audio response JSON: {error}")))?; Ok(request .config .transform_transcription_response(&request.model, response_json)? @@ -51,14 +47,13 @@ pub(crate) async fn execute_audio_transcription_provider_call( pub(crate) async fn sign_request( request: &ProviderAudioTranscriptionRequest, optional_params: &serde_json::Map, -) -> CoreResult { +) -> Result { let env_lookup = environment_lookup; let auth = request .config .auth_strategy(&request.model, optional_params, &env_lookup)?; - let body = serde_json::to_vec(&request.body).map_err(|error| { - CoreError::InvalidRequest(format!("invalid audio request body: {error}")) - })?; + let body = serde_json::to_vec(&request.body) + .map_err(|error| Error::InvalidRequest(format!("invalid audio request body: {error}")))?; let mut headers = super::common_utils::string_headers(None)?; headers.insert("Content-Type".to_string(), "application/json".to_string()); headers.extend(request.upstream_headers.iter().cloned()); diff --git a/litellm-rust/crates/ai-gateway/src/audio_transcription/hooks.rs b/litellm-rust/crates/ai-gateway/src/audio_transcription/hooks.rs index 0c9faeda6e7..5e1240de759 100644 --- a/litellm-rust/crates/ai-gateway/src/audio_transcription/hooks.rs +++ b/litellm-rust/crates/ai-gateway/src/audio_transcription/hooks.rs @@ -1,11 +1,9 @@ -use std::future::Future; -use std::pin::Pin; - -use litellm_core::CoreResult; use litellm_core::audio_transcription::transformation::AudioTranscriptionAuth; use litellm_core::call_lifecycle::{CallLifecycleContext, CallLifecycleHooks, CallLifecycleTiming}; -use litellm_core::error::CoreError; +use litellm_core::error::Error; use serde_json::{Map, Value, json}; +use std::future::Future; +use std::pin::Pin; use super::common_utils::{audio_transcription_provider_config, has_header, string_headers}; use super::handler::sign_request; @@ -26,7 +24,7 @@ pub(crate) struct AudioTranscriptionLifecycleHooks { request_metadata: RequestMetadata, } -type AudioFuture<'a, T> = Pin> + Send + 'a>>; +type AudioFuture<'a, T> = Pin> + Send + 'a>>; type AudioLogFuture<'a> = Pin + Send + 'a>>; impl AudioTranscriptionLifecycleHooks { @@ -45,7 +43,7 @@ impl AudioTranscriptionLifecycleHooks { async fn run_pre_call_guardrails( &self, request: PreparedAudioTranscriptionRequest, - ) -> CoreResult { + ) -> Result { if self.guardrail_runner.is_empty() { return Ok(request); } @@ -63,17 +61,17 @@ impl AudioTranscriptionLifecycleHooks { .await .map_err(guardrail_error_to_core_error)?; let Value::Object(mut data) = guardrail_request.data else { - return Err(CoreError::InvalidRequest( + return Err(Error::InvalidRequest( "audio transcription pre_call guardrail must return an object".to_string(), )); }; let audio = data.remove("audio").ok_or_else(|| { - CoreError::InvalidRequest("audio transcription guardrail removed audio".to_string()) + Error::InvalidRequest("audio transcription guardrail removed audio".to_string()) })?; let optional_params = match data.remove("optional_params") { Some(Value::Object(value)) => value, Some(_) => { - return Err(CoreError::InvalidRequest( + return Err(Error::InvalidRequest( "audio transcription optional_params must be an object".to_string(), )); } @@ -89,9 +87,9 @@ impl AudioTranscriptionLifecycleHooks { async fn prepare_provider_request( &self, request: PreparedAudioTranscriptionRequest, - ) -> CoreResult { + ) -> Result { let config = audio_transcription_provider_config(&request.custom_llm_provider) - .ok_or_else(|| CoreError::InvalidProvider(request.custom_llm_provider.clone()))?; + .ok_or_else(|| Error::InvalidProvider(request.custom_llm_provider.clone()))?; let env_lookup = super::handler::environment_lookup; let headers = string_headers(request.extra_headers)?; let url = config.complete_url( @@ -135,7 +133,7 @@ impl AudioTranscriptionLifecycleHooks { async fn run_during_call_guardrails( &self, request: ProviderAudioTranscriptionRequest, - ) -> CoreResult { + ) -> Result { if self.guardrail_runner.is_empty() { return Ok(request); } @@ -153,12 +151,12 @@ impl AudioTranscriptionLifecycleHooks { .await .map_err(guardrail_error_to_core_error)?; let Value::Object(mut data) = guardrail_request.data else { - return Err(CoreError::InvalidRequest( + return Err(Error::InvalidRequest( "audio transcription during_call guardrail must return an object".to_string(), )); }; let body = data.remove("body").ok_or_else(|| { - CoreError::InvalidRequest("audio transcription guardrail removed body".to_string()) + Error::InvalidRequest("audio transcription guardrail removed body".to_string()) })?; Ok(ProviderAudioTranscriptionRequest { body, ..request }) } @@ -241,7 +239,7 @@ impl CallLifecycleHooks( &'a self, context: &'a CallLifecycleContext, - error: &'a CoreError, + error: &'a Error, timing: &'a CallLifecycleTiming, ) -> Self::FailureFuture<'a> { Box::pin(async move { @@ -281,22 +279,22 @@ fn guardrail_context(metadata: &RequestMetadata) -> GuardrailContext { } } -fn guardrail_error_to_core_error(error: GuardrailError) -> CoreError { - CoreError::InvalidRequest(format!("{}: {}", error.kind, error.message)) +fn guardrail_error_to_core_error(error: GuardrailError) -> Error { + Error::InvalidRequest(format!("{}: {}", error.kind, error.message)) } -fn core_error_kind(error: &CoreError) -> &'static str { +fn core_error_kind(error: &Error) -> &'static str { match error { - CoreError::Auth(_) => "AuthError", - CoreError::InvalidProvider(_) => "InvalidProvider", - CoreError::InvalidRequest(_) => "InvalidRequest", - CoreError::InvalidType { .. } => "InvalidType", - CoreError::MissingField(_) => "MissingField", - CoreError::Http { .. } => "HttpError", - CoreError::InvalidResponse(_) => "InvalidResponse", - CoreError::Network(_) => "NetworkError", - CoreError::Connect(_) => "ConnectError", - CoreError::Routing(_) => "RoutingError", - CoreError::Unsupported(_) => "UnsupportedRequest", + Error::Auth(_) => "AuthError", + Error::InvalidProvider(_) => "InvalidProvider", + Error::InvalidRequest(_) => "InvalidRequest", + Error::InvalidType { .. } => "InvalidType", + Error::MissingField(_) => "MissingField", + Error::Http { .. } => "HttpError", + Error::InvalidResponse(_) => "InvalidResponse", + Error::Network(_) => "NetworkError", + Error::Connect(_) => "ConnectError", + Error::Routing(_) => "RoutingError", + Error::Unsupported(_) => "UnsupportedRequest", } } diff --git a/litellm-rust/crates/ai-gateway/src/audio_transcription/mod.rs b/litellm-rust/crates/ai-gateway/src/audio_transcription/mod.rs index 5d33d912c40..3983846d7b6 100644 --- a/litellm-rust/crates/ai-gateway/src/audio_transcription/mod.rs +++ b/litellm-rust/crates/ai-gateway/src/audio_transcription/mod.rs @@ -1,4 +1,4 @@ -use litellm_core::CoreResult; +use litellm_core::Error; use litellm_core::call_lifecycle::CallLifecycle; use serde_json::Value; @@ -13,7 +13,7 @@ pub use types::AudioTranscriptionRequest; use handler::execute_audio_transcription_provider_call; use prepare::{PreparedAudioTranscriptionCall, prepare_audio_transcription_call}; -pub async fn audio_transcription(request: AudioTranscriptionRequest<'_>) -> CoreResult { +pub async fn audio_transcription(request: AudioTranscriptionRequest<'_>) -> Result { let PreparedAudioTranscriptionCall { request, hooks } = prepare_audio_transcription_call(request); CallLifecycle::default() diff --git a/litellm-rust/crates/ai-gateway/src/io/realtime.rs b/litellm-rust/crates/ai-gateway/src/io/realtime.rs index 845e7bf9527..662f7328982 100644 --- a/litellm-rust/crates/ai-gateway/src/io/realtime.rs +++ b/litellm-rust/crates/ai-gateway/src/io/realtime.rs @@ -15,8 +15,7 @@ use std::time::Duration; use futures_util::stream::{SplitSink, SplitStream}; use futures_util::{Sink, SinkExt, Stream, StreamExt}; -use litellm_core::CoreResult; -use litellm_core::error::CoreError; +use litellm_core::error::Error; use litellm_core::realtime::transformation::RealtimeProviderConfig; use litellm_core::realtime::types::RealtimeEvent; use tokio::net::TcpStream; @@ -48,7 +47,7 @@ pub(crate) type UpstreamRx = SplitStream; /// Resolve the OpenAI API key from the explicit param or the environment. /// /// Blank/whitespace values are treated as absent (guard at resolution time). -pub(crate) fn resolve_api_key(api_key: Option<&str>) -> CoreResult { +pub(crate) fn resolve_api_key(api_key: Option<&str>) -> Result { api_key .map(str::trim) .filter(|key| !key.is_empty()) @@ -58,7 +57,7 @@ pub(crate) fn resolve_api_key(api_key: Option<&str>) -> CoreResult { .ok() .filter(|key| !key.trim().is_empty()) }) - .ok_or_else(|| CoreError::Auth(MISSING_KEY_MESSAGE.to_string())) + .ok_or_else(|| Error::Auth(MISSING_KEY_MESSAGE.to_string())) } /// Open the upstream WebSocket to OpenAI for `(model, api_key, api_base)`. @@ -70,24 +69,24 @@ pub(crate) async fn dial_upstream( model: &str, api_key: &str, api_base: Option<&str>, -) -> CoreResult { +) -> Result { let url = OPENAI_REALTIME_CONFIG.complete_url(api_base, model); let mut request = url .as_str() .into_client_request() - .map_err(|err| CoreError::Network(err.to_string()))?; + .map_err(|err| Error::Network(err.to_string()))?; // GA realtime: only Authorization. The legacy OpenAI-Beta header triggers // beta_api_shape_disabled, so we do not send it. request.headers_mut().insert( AUTHORIZATION, HeaderValue::from_str(&format!("Bearer {api_key}")) - .map_err(|err| CoreError::Auth(err.to_string()))?, + .map_err(|err| Error::Auth(err.to_string()))?, ); let (upstream, _response) = connect_async(request) .await - .map_err(|err| CoreError::Network(err.to_string()))?; + .map_err(|err| Error::Network(err.to_string()))?; Ok(upstream) } @@ -96,22 +95,22 @@ pub(crate) async fn dial_upstream( /// Used by the pool to pre-read OpenAI's unprompted `session.created`. Returns an /// error on a non-text frame, a closed socket, or undecodable JSON so the pool can /// discard a misbehaving socket rather than warm it. -pub(crate) async fn read_event(upstream_rx: &mut UpstreamRx) -> CoreResult { +pub(crate) async fn read_event(upstream_rx: &mut UpstreamRx) -> Result { loop { let message = upstream_rx .next() .await - .ok_or_else(|| CoreError::Network("upstream closed before first event".to_string()))? - .map_err(|err| CoreError::Network(err.to_string()))?; + .ok_or_else(|| Error::Network("upstream closed before first event".to_string()))? + .map_err(|err| Error::Network(err.to_string()))?; match message { Message::Text(text) => { return serde_json::from_str(&text) - .map_err(|err| CoreError::InvalidResponse(err.to_string())); + .map_err(|err| Error::InvalidResponse(err.to_string())); } // Ignore protocol frames (ping/pong) while waiting for the first event. Message::Ping(_) | Message::Pong(_) => continue, Message::Close(_) => { - return Err(CoreError::Network( + return Err(Error::Network( "upstream closed before first event".to_string(), )); } @@ -139,7 +138,7 @@ pub(crate) async fn splice( mut observe: impl FnMut(&RealtimeEvent) + Send, mut client_in: In, mut client_out: Out, -) -> CoreResult<()> +) -> Result<(), Error> where In: Stream + Unpin + Send, Out: Sink + Unpin + Send, @@ -154,7 +153,7 @@ where client_out .send(outbound) .await - .map_err(|err| CoreError::Network(err.to_string()))?; + .map_err(|err| Error::Network(err.to_string()))?; } } @@ -175,26 +174,26 @@ where // inflate its own spend log. Logging observes upstream events only. for outbound in config.transform_realtime_request(&event, model)?.events { let payload = serde_json::to_string(&outbound) - .map_err(|err| CoreError::InvalidResponse(err.to_string()))?; + .map_err(|err| Error::InvalidResponse(err.to_string()))?; upstream_tx .send(Message::Text(payload)) .await - .map_err(|err| CoreError::Network(err.to_string()))?; + .map_err(|err| Error::Network(err.to_string()))?; } } // upstream -> client upstream_message = upstream_rx.next() => { let Some(message) = upstream_message else { break }; // upstream closed - match message.map_err(|err| CoreError::Network(err.to_string()))? { + match message.map_err(|err| Error::Network(err.to_string()))? { Message::Text(text) => { let event: RealtimeEvent = serde_json::from_str(&text) - .map_err(|err| CoreError::InvalidResponse(err.to_string()))?; + .map_err(|err| Error::InvalidResponse(err.to_string()))?; observe(&event); for outbound in config.transform_realtime_response(&event, model)?.events { client_out .send(outbound) .await - .map_err(|err| CoreError::Network(err.to_string()))?; + .map_err(|err| Error::Network(err.to_string()))?; } } Message::Close(_) => break, @@ -225,7 +224,7 @@ pub async fn realtime( observe: impl FnMut(&RealtimeEvent) + Send, client_in: In, client_out: Out, -) -> CoreResult<()> +) -> Result<(), Error> where In: Stream + Unpin + Send, Out: Sink + Unpin + Send, @@ -258,7 +257,7 @@ pub async fn realtime_warm( observe: impl FnMut(&RealtimeEvent) + Send, client_in: In, client_out: Out, -) -> CoreResult<()> +) -> Result<(), Error> where In: Stream + Unpin + Send, Out: Sink + Unpin + Send, diff --git a/litellm-rust/crates/ai-gateway/src/io/realtime_pool.rs b/litellm-rust/crates/ai-gateway/src/io/realtime_pool.rs index 4a1a3cd1166..49e9c459a88 100644 --- a/litellm-rust/crates/ai-gateway/src/io/realtime_pool.rs +++ b/litellm-rust/crates/ai-gateway/src/io/realtime_pool.rs @@ -28,7 +28,7 @@ use std::sync::{Arc, Mutex}; use std::time::{Duration, Instant}; use futures_util::StreamExt; -use litellm_core::CoreResult; +use litellm_core::Error; use litellm_core::realtime::types::RealtimeEvent; use crate::io::realtime::{ @@ -438,7 +438,7 @@ impl RealtimePool { /// /// `key.api_key` is already resolved (non-blank). The first frame OpenAI sends /// unprompted is `session.created`; we buffer exactly that and read nothing more. -async fn warm_one(key: &UpstreamKey) -> CoreResult { +async fn warm_one(key: &UpstreamKey) -> Result { let upstream: UpstreamWs = dial_upstream(&key.model, &key.api_key, key.api_base.as_deref()).await?; let (tx, mut rx) = upstream.split(); diff --git a/litellm-rust/crates/ai-gateway/src/io/responses_ws.rs b/litellm-rust/crates/ai-gateway/src/io/responses_ws.rs index 9b51019f4bc..0b01747b1a5 100644 --- a/litellm-rust/crates/ai-gateway/src/io/responses_ws.rs +++ b/litellm-rust/crates/ai-gateway/src/io/responses_ws.rs @@ -4,10 +4,10 @@ use std::time::Duration; use futures_util::stream::{SplitSink, SplitStream}; use futures_util::{Sink, SinkExt, Stream, StreamExt}; +use litellm_core::Error; use litellm_core::providers::openai::responses::transformation::OPENAI_RESPONSES_WS_CONFIG; use litellm_core::responses::types::ResponsesWsEvent; use litellm_core::responses::websocket::ResponsesWebSocketProviderConfig; -use litellm_core::{CoreError, CoreResult}; use tokio::net::TcpStream; use tokio::sync::Mutex; use tokio_tungstenite::tungstenite::Message; @@ -37,51 +37,49 @@ impl ResponsesWebSocketConnection { url: &str, headers: &HashMap, timeout: Option, - ) -> CoreResult { + ) -> Result { let mut request = url .into_client_request() - .map_err(|error| CoreError::Network(error.to_string()))?; + .map_err(|error| Error::Network(error.to_string()))?; for (name, value) in headers { let header_name = name .parse::() - .map_err(|error| CoreError::InvalidRequest(error.to_string()))?; + .map_err(|error| Error::InvalidRequest(error.to_string()))?; let header_value = HeaderValue::from_str(value) - .map_err(|error| CoreError::InvalidRequest(error.to_string()))?; + .map_err(|error| Error::InvalidRequest(error.to_string()))?; request.headers_mut().insert(header_name, header_value); } let connect = connect_async(request); let result = match timeout { Some(timeout) => tokio::time::timeout(timeout, connect).await.map_err(|_| { - CoreError::Network("Responses WebSocket connection timed out".to_string()) + Error::Network("Responses WebSocket connection timed out".to_string()) })?, None => connect.await, }; let (socket, _) = result.map_err(|error| match error { - tokio_tungstenite::tungstenite::Error::Http(response) => CoreError::Http { + tokio_tungstenite::tungstenite::Error::Http(response) => Error::Http { status: response.status().as_u16(), body: String::new(), }, - other => CoreError::Network(other.to_string()), + other => Error::Network(other.to_string()), })?; Ok(Self { socket: Arc::new(Mutex::new(Some(socket))), }) } - pub async fn send_text(&self, text: String) -> CoreResult<()> { + pub async fn send_text(&self, text: String) -> Result<(), Error> { let mut socket = self.socket.lock().await; let Some(socket) = socket.as_mut() else { - return Err(CoreError::Network( - "Responses WebSocket is closed".to_string(), - )); + return Err(Error::Network("Responses WebSocket is closed".to_string())); }; socket .send(Message::Text(text)) .await - .map_err(|error| CoreError::Network(error.to_string())) + .map_err(|error| Error::Network(error.to_string())) } - pub async fn recv_text(&self) -> CoreResult> { + pub async fn recv_text(&self) -> Result, Error> { let mut socket_guard = self.socket.lock().await; let Some(socket) = socket_guard.as_mut() else { return Ok(None); @@ -90,27 +88,27 @@ impl ResponsesWebSocketConnection { Some(Ok(Message::Text(text))) => Ok(Some(text)), Some(Ok(Message::Binary(bytes))) => String::from_utf8(bytes.to_vec()) .map(Some) - .map_err(|error| CoreError::InvalidResponse(error.to_string())), + .map_err(|error| Error::InvalidResponse(error.to_string())), Some(Ok(Message::Close(_))) | None => Ok(None), Some(Ok(_)) => Ok(None), - Some(Err(error)) => Err(CoreError::Network(error.to_string())), + Some(Err(error)) => Err(Error::Network(error.to_string())), } } - pub async fn close(&self) -> CoreResult<()> { + pub async fn close(&self) -> Result<(), Error> { let mut socket = self.socket.lock().await; if let Some(socket) = socket.as_mut() { socket .close(None) .await - .map_err(|error| CoreError::Network(error.to_string()))?; + .map_err(|error| Error::Network(error.to_string()))?; } *socket = None; Ok(()) } } -pub(crate) fn resolve_api_key(api_key: Option<&str>) -> CoreResult { +pub(crate) fn resolve_api_key(api_key: Option<&str>) -> Result { api_key .map(str::trim) .filter(|value| !value.is_empty()) @@ -120,38 +118,38 @@ pub(crate) fn resolve_api_key(api_key: Option<&str>) -> CoreResult { .ok() .filter(|value| !value.trim().is_empty()) }) - .ok_or_else(|| CoreError::Auth(MISSING_KEY_MESSAGE.to_string())) + .ok_or_else(|| Error::Auth(MISSING_KEY_MESSAGE.to_string())) } async fn dial_upstream( model: &str, api_key: &str, api_base: Option<&str>, -) -> CoreResult { +) -> Result { let url = OPENAI_RESPONSES_WS_CONFIG.complete_websocket_url(api_base, model); let mut request = url .as_str() .into_client_request() - .map_err(|error| CoreError::Network(error.to_string()))?; + .map_err(|error| Error::Network(error.to_string()))?; request.headers_mut().insert( AUTHORIZATION, HeaderValue::from_str(&format!("Bearer {api_key}")) - .map_err(|error| CoreError::Auth(error.to_string()))?, + .map_err(|error| Error::Auth(error.to_string()))?, ); let result = tokio::time::timeout( Duration::from_secs(DEFAULT_RESPONSES_WS_CONNECT_TIMEOUT_SECS), connect_async(request), ) .await - .map_err(|_| CoreError::Network("Responses WebSocket connection timed out".to_string()))?; + .map_err(|_| Error::Network("Responses WebSocket connection timed out".to_string()))?; result .map(|(socket, _)| socket) .map_err(|error| match error { - tokio_tungstenite::tungstenite::Error::Http(response) => CoreError::Http { + tokio_tungstenite::tungstenite::Error::Http(response) => Error::Http { status: response.status().as_u16(), body: String::new(), }, - other => CoreError::Network(other.to_string()), + other => Error::Network(other.to_string()), }) } @@ -166,7 +164,7 @@ impl ResponsesWebSocketStreaming { observe: impl FnMut(&ResponsesWsEvent) + Send, client_in: In, client_out: Out, - ) -> CoreResult<()> + ) -> Result<(), Error> where In: Stream + Unpin + Send, Out: Sink + Unpin + Send, @@ -193,7 +191,7 @@ pub(crate) async fn splice( mut observe: impl FnMut(&ResponsesWsEvent) + Send, mut client_in: In, mut client_out: Out, -) -> CoreResult<()> +) -> Result<(), Error> where In: Stream + Unpin + Send, Out: Sink + Unpin + Send, @@ -210,18 +208,18 @@ where .events { let payload = serde_json::to_string(&outbound) - .map_err(|error| CoreError::InvalidResponse(error.to_string()))?; + .map_err(|error| Error::InvalidResponse(error.to_string()))?; upstream_tx.send(Message::Text(payload)) .await - .map_err(|error| CoreError::Network(error.to_string()))?; + .map_err(|error| Error::Network(error.to_string()))?; } } message = upstream_rx.next() => { let Some(message) = message else { break }; - match message.map_err(|error| CoreError::Network(error.to_string()))? { + match message.map_err(|error| Error::Network(error.to_string()))? { Message::Text(text) => { let event = serde_json::from_str::(&text) - .map_err(|error| CoreError::InvalidResponse(error.to_string()))?; + .map_err(|error| Error::InvalidResponse(error.to_string()))?; observe(&event); for outbound in OPENAI_RESPONSES_WS_CONFIG .transform_ws_response(&event, model)? @@ -229,7 +227,7 @@ where { client_out.send(outbound) .await - .map_err(|error| CoreError::Network(error.to_string()))?; + .map_err(|error| Error::Network(error.to_string()))?; } } Message::Close(_) => break, @@ -252,7 +250,7 @@ pub async fn async_responses_websocket( mut observe: impl FnMut(&ResponsesWsEvent) + Send, client_in: In, client_out: Out, -) -> CoreResult<()> +) -> Result<(), Error> where In: Stream + Unpin + Send, Out: Sink + Unpin + Send, @@ -267,11 +265,11 @@ where .events { let payload = serde_json::to_string(&outbound) - .map_err(|error| CoreError::InvalidResponse(error.to_string()))?; + .map_err(|error| Error::InvalidResponse(error.to_string()))?; upstream_tx .send(Message::Text(payload)) .await - .map_err(|error| CoreError::Network(error.to_string()))?; + .map_err(|error| Error::Network(error.to_string()))?; } } ResponsesWebSocketStreaming::bidirectional_forward( @@ -296,7 +294,7 @@ pub async fn responses_ws( observe: impl FnMut(&ResponsesWsEvent) + Send, client_in: In, client_out: Out, -) -> CoreResult<()> +) -> Result<(), Error> where In: Stream + Unpin + Send, Out: Sink + Unpin + Send, @@ -514,7 +512,7 @@ mod tests { ) .await .expect_err("status error"); - assert!(matches!(error, CoreError::Http { status: 401, .. })); + assert!(matches!(error, Error::Http { status: 401, .. })); server.await.expect("server task"); } @@ -543,7 +541,7 @@ mod tests { ) .await .expect_err("status error"); - assert!(matches!(error, CoreError::Http { status: 500, .. })); + assert!(matches!(error, Error::Http { status: 500, .. })); server.await.expect("server task"); } } diff --git a/litellm-rust/crates/ai-gateway/src/ocr/common_utils.rs b/litellm-rust/crates/ai-gateway/src/ocr/common_utils.rs index 9bc2818b6e7..e0ce165dc93 100644 --- a/litellm-rust/crates/ai-gateway/src/ocr/common_utils.rs +++ b/litellm-rust/crates/ai-gateway/src/ocr/common_utils.rs @@ -3,8 +3,7 @@ use std::time::{Duration, Instant}; use base64::Engine; use base64::engine::general_purpose::STANDARD as BASE64_STANDARD; -use litellm_core::CoreResult; -use litellm_core::error::CoreError; +use litellm_core::error::Error; use litellm_core::ocr::transformation::OcrProviderConfig; use reqwest::Url; use serde_json::{Map, Value}; @@ -56,7 +55,7 @@ fn is_azure_document_intelligence_model(model: &str) -> bool { pub(super) fn string_headers( extra_headers: Option>, -) -> CoreResult> { +) -> Result, Error> { extra_headers .unwrap_or_default() .into_iter() @@ -65,7 +64,7 @@ pub(super) fn string_headers( .as_str() .map(|value| (key.clone(), value.to_string())) .ok_or_else(|| { - CoreError::InvalidRequest(format!( + Error::InvalidRequest(format!( "OCR extra_headers.{key} must be a string, got {}", litellm_core::error::json_type_name(&value) )) @@ -80,7 +79,7 @@ pub(super) fn has_header(headers: &[(String, String)], name: &str) -> bool { .any(|(key, _)| key.eq_ignore_ascii_case(name)) } -fn document_url_field(document: &Value) -> CoreResult> { +fn document_url_field(document: &Value) -> Result, Error> { let Some(object) = document.as_object() else { return Ok(None); }; @@ -138,13 +137,13 @@ fn is_blocked_ip(ip: IpAddr) -> bool { } } -fn blocked_url_error(url: &Url) -> CoreError { - CoreError::InvalidRequest(format!( +fn blocked_url_error(url: &Url) -> Error { + Error::InvalidRequest(format!( "OCR document URL rejected by SSRF protection: {url}" )) } -async fn validate_safe_fetch_url(url: &Url) -> CoreResult<()> { +async fn validate_safe_fetch_url(url: &Url) -> Result<(), Error> { if !matches!(url.scheme(), "http" | "https") { return Err(blocked_url_error(url)); } @@ -162,7 +161,7 @@ async fn validate_safe_fetch_url(url: &Url) -> CoreResult<()> { .ok_or_else(|| blocked_url_error(url))?; let addresses = tokio::net::lookup_host((host, port)) .await - .map_err(|err| CoreError::Network(err.to_string()))?; + .map_err(|err| Error::Network(err.to_string()))?; let mut saw_address = false; for address in addresses { saw_address = true; @@ -176,25 +175,25 @@ async fn validate_safe_fetch_url(url: &Url) -> CoreResult<()> { Ok(()) } -fn redirect_location(response: &reqwest::Response, url: &Url) -> CoreResult { +fn redirect_location(response: &reqwest::Response, url: &Url) -> Result { let location = response .headers() .get(reqwest::header::LOCATION) .and_then(|value| value.to_str().ok()) .ok_or_else(|| { - CoreError::InvalidResponse("OCR document redirect missing Location header".to_string()) + Error::InvalidResponse("OCR document redirect missing Location header".to_string()) })?; url.join(location) - .map_err(|err| CoreError::InvalidResponse(format!("invalid OCR document redirect: {err}"))) + .map_err(|err| Error::InvalidResponse(format!("invalid OCR document redirect: {err}"))) } -async fn safe_get_document_url(url: &str) -> CoreResult<(Url, reqwest::Response)> { +async fn safe_get_document_url(url: &str) -> Result<(Url, reqwest::Response), Error> { let client = reqwest::Client::builder() .redirect(reqwest::redirect::Policy::none()) .build() - .map_err(|err| CoreError::Network(err.to_string()))?; + .map_err(|err| Error::Network(err.to_string()))?; let mut current_url = Url::parse(url) - .map_err(|err| CoreError::InvalidRequest(format!("invalid OCR document URL: {err}")))?; + .map_err(|err| Error::InvalidRequest(format!("invalid OCR document URL: {err}")))?; for _ in 0..MAX_SAFE_FETCH_REDIRECTS { validate_safe_fetch_url(¤t_url).await?; @@ -202,28 +201,28 @@ async fn safe_get_document_url(url: &str) -> CoreResult<(Url, reqwest::Response) .get(current_url.clone()) .send() .await - .map_err(|err| CoreError::Network(err.to_string()))?; + .map_err(|err| Error::Network(err.to_string()))?; if !response.status().is_redirection() { return Ok((current_url, response)); } current_url = redirect_location(&response, ¤t_url)?; } - Err(CoreError::InvalidRequest( + Err(Error::InvalidRequest( "Too many redirects while fetching OCR document URL".to_string(), )) } -fn enforce_download_size(content_length: u64, max_bytes: u64, url: &Url) -> CoreResult<()> { +fn enforce_download_size(content_length: u64, max_bytes: u64, url: &Url) -> Result<(), Error> { if max_bytes == 0 { - return Err(CoreError::InvalidRequest(format!( + return Err(Error::InvalidRequest(format!( "OCR document URL download is disabled (MAX_IMAGE_URL_DOWNLOAD_SIZE_MB=0). url={url}" ))); } if content_length > max_bytes { let size_mb = content_length as f64 / (1024.0 * 1024.0); let max_size_mb = max_bytes as f64 / (1024.0 * 1024.0); - return Err(CoreError::InvalidRequest(format!( + return Err(Error::InvalidRequest(format!( "OCR document size ({size_mb:.2}MB) exceeds maximum allowed size ({max_size_mb:.2}MB). url={url}" ))); } @@ -233,7 +232,7 @@ fn enforce_download_size(content_length: u64, max_bytes: u64, url: &Url) -> Core async fn read_response_with_limit( mut response: reqwest::Response, url: &Url, -) -> CoreResult> { +) -> Result, Error> { let max_bytes = max_document_download_bytes(); if let Some(content_length) = response.content_length() { enforce_download_size(content_length, max_bytes, url)?; @@ -246,7 +245,7 @@ async fn read_response_with_limit( while let Some(chunk) = response .chunk() .await - .map_err(|err| CoreError::Network(err.to_string()))? + .map_err(|err| Error::Network(err.to_string()))? { bytes_downloaded += chunk.len() as u64; enforce_download_size(bytes_downloaded, max_bytes, url)?; @@ -255,7 +254,7 @@ async fn read_response_with_limit( Ok(bytes) } -pub(super) async fn convert_document_url_to_data_uri(document: Value) -> CoreResult { +pub(super) async fn convert_document_url_to_data_uri(document: Value) -> Result { let Some((field, url)) = document_url_field(&document)? else { return Ok(document); }; @@ -267,7 +266,7 @@ pub(super) async fn convert_document_url_to_data_uri(document: Value) -> CoreRes let status = response.status(); if !status.is_success() { let body = response.text().await.unwrap_or_default(); - return Err(CoreError::Http { + return Err(Error::Http { status: status.as_u16(), body: truncate_error_body(&body), }); @@ -290,7 +289,7 @@ pub(super) async fn convert_document_url_to_data_uri(document: Value) -> CoreRes let mut transformed = document .as_object() .cloned() - .ok_or_else(|| CoreError::InvalidRequest("OCR document must be an object".to_string()))?; + .ok_or_else(|| Error::InvalidRequest("OCR document must be an object".to_string()))?; transformed.insert(field.to_string(), Value::String(data_uri)); Ok(Value::Object(transformed)) } @@ -316,11 +315,11 @@ fn retry_after_secs(response: &reqwest::Response) -> u64 { .unwrap_or(2) } -fn operation_status(response_json: &Value) -> CoreResult<&str> { +fn operation_status(response_json: &Value) -> Result<&str, Error> { let status = response_json .get("status") .and_then(Value::as_str) - .ok_or(CoreError::MissingField("status"))?; + .ok_or(Error::MissingField("status"))?; match status { "succeeded" => Ok("succeeded"), "running" | "notStarted" => Ok("running"), @@ -330,11 +329,11 @@ fn operation_status(response_json: &Value) -> CoreResult<&str> { .and_then(|error| error.get("message")) .and_then(Value::as_str) .unwrap_or("Unknown error"); - Err(CoreError::InvalidResponse(format!( + Err(Error::InvalidResponse(format!( "Azure Document Intelligence analysis failed: {message}" ))) } - other => Err(CoreError::InvalidResponse(format!( + other => Err(Error::InvalidResponse(format!( "Unknown operation status: {other}" ))), } @@ -345,9 +344,9 @@ pub(super) async fn poll_document_intelligence( original_url: &str, headers: &[(String, String)], timeout: Option, -) -> CoreResult { +) -> Result { if !same_origin(operation_url, original_url) { - return Err(CoreError::InvalidResponse( + return Err(Error::InvalidResponse( "Azure Document Intelligence: rejected cross-origin polling URL".to_string(), )); } @@ -358,7 +357,7 @@ pub(super) async fn poll_document_intelligence( )); loop { if start.elapsed() > timeout { - return Err(CoreError::Network(format!( + return Err(Error::Network(format!( "Azure Document Intelligence operation polling timed out after {} seconds", timeout.as_secs() ))); @@ -373,21 +372,21 @@ pub(super) async fn poll_document_intelligence( let response = request_builder .send() .await - .map_err(|err| CoreError::Network(err.to_string()))?; + .map_err(|err| Error::Network(err.to_string()))?; let retry_after = retry_after_secs(&response); let status = response.status(); let text = response .text() .await - .map_err(|err| CoreError::Network(err.to_string()))?; + .map_err(|err| Error::Network(err.to_string()))?; if !status.is_success() { - return Err(CoreError::Http { + return Err(Error::Http { status: status.as_u16(), body: truncate_error_body(&text), }); } let response_json: Value = serde_json::from_str(&text).map_err(|err| { - CoreError::InvalidResponse(format!("invalid Azure DI poll response JSON: {err}")) + Error::InvalidResponse(format!("invalid Azure DI poll response JSON: {err}")) })?; if operation_status(&response_json)? == "succeeded" { return Ok(response_json); @@ -426,7 +425,7 @@ mod tests { assert!(matches!( error, - CoreError::InvalidRequest(message) + Error::InvalidRequest(message) if message.contains("SSRF protection") )); } diff --git a/litellm-rust/crates/ai-gateway/src/ocr/handler.rs b/litellm-rust/crates/ai-gateway/src/ocr/handler.rs index 1de34eb400e..815bc84363a 100644 --- a/litellm-rust/crates/ai-gateway/src/ocr/handler.rs +++ b/litellm-rust/crates/ai-gateway/src/ocr/handler.rs @@ -1,5 +1,4 @@ -use litellm_core::CoreResult; -use litellm_core::error::CoreError; +use litellm_core::error::Error; use litellm_core::ocr::transformation::OcrResponseHandling; use serde_json::Value; @@ -7,7 +6,7 @@ use super::common_utils::{poll_document_intelligence, truncate_error_body}; use super::types::ProviderOcrRequest; use crate::client::http_client; -pub(crate) async fn execute_ocr_provider_call(request: ProviderOcrRequest) -> CoreResult { +pub(crate) async fn execute_ocr_provider_call(request: ProviderOcrRequest) -> Result { let mut request_builder = http_client().post(&request.url).json(&request.body); for (key, value) in &request.upstream_headers { request_builder = request_builder.header(key, value); @@ -19,7 +18,7 @@ pub(crate) async fn execute_ocr_provider_call(request: ProviderOcrRequest) -> Co let response = request_builder .send() .await - .map_err(|err| CoreError::Network(err.to_string()))?; + .map_err(|err| Error::Network(err.to_string()))?; let status = response.status(); if request.config.response_handling() == OcrResponseHandling::AzureDocumentIntelligencePoll @@ -31,7 +30,7 @@ pub(crate) async fn execute_ocr_provider_call(request: ProviderOcrRequest) -> Co .and_then(|value| value.to_str().ok()) .map(str::to_string) .ok_or_else(|| { - CoreError::InvalidResponse( + Error::InvalidResponse( "Azure Document Intelligence returned 202 but no Operation-Location header found" .to_string(), ) @@ -52,17 +51,17 @@ pub(crate) async fn execute_ocr_provider_call(request: ProviderOcrRequest) -> Co let text = response .text() .await - .map_err(|err| CoreError::Network(err.to_string()))?; + .map_err(|err| Error::Network(err.to_string()))?; if !status.is_success() { - return Err(CoreError::Http { + return Err(Error::Http { status: status.as_u16(), body: truncate_error_body(&text), }); } let response_json: Value = serde_json::from_str(&text) - .map_err(|err| CoreError::InvalidResponse(format!("invalid OCR response JSON: {err}")))?; + .map_err(|err| Error::InvalidResponse(format!("invalid OCR response JSON: {err}")))?; Ok(request .config diff --git a/litellm-rust/crates/ai-gateway/src/ocr/hooks.rs b/litellm-rust/crates/ai-gateway/src/ocr/hooks.rs index 95df566dc53..401e26d3b29 100644 --- a/litellm-rust/crates/ai-gateway/src/ocr/hooks.rs +++ b/litellm-rust/crates/ai-gateway/src/ocr/hooks.rs @@ -1,11 +1,9 @@ -use std::future::Future; -use std::pin::Pin; - -use litellm_core::CoreResult; use litellm_core::call_lifecycle::{CallLifecycleContext, CallLifecycleHooks, CallLifecycleTiming}; -use litellm_core::error::CoreError; +use litellm_core::error::Error; use litellm_core::ocr::transformation::OcrAuthStrategy; use serde_json::{Map, Value, json}; +use std::future::Future; +use std::pin::Pin; use super::common_utils::{ convert_document_url_to_data_uri, has_header, ocr_provider_config, string_headers, @@ -27,7 +25,7 @@ pub(crate) struct OcrLifecycleHooks { request_metadata: RequestMetadata, } -type OcrFuture<'a, T> = Pin> + Send + 'a>>; +type OcrFuture<'a, T> = Pin> + Send + 'a>>; type OcrLogFuture<'a> = Pin + Send + 'a>>; impl OcrLifecycleHooks { @@ -46,7 +44,7 @@ impl OcrLifecycleHooks { async fn run_pre_call_guardrails( &self, request: PreparedOcrRequest, - ) -> CoreResult { + ) -> Result { if self.guardrail_runner.is_empty() { return Ok(request); } @@ -74,9 +72,9 @@ impl OcrLifecycleHooks { async fn prepare_provider_request( &self, request: PreparedOcrRequest, - ) -> CoreResult { + ) -> Result { let config = ocr_provider_config(&request.custom_llm_provider, &request.model) - .ok_or_else(|| CoreError::InvalidProvider(request.custom_llm_provider.clone()))?; + .ok_or_else(|| Error::InvalidProvider(request.custom_llm_provider.clone()))?; let env_lookup = |key: &str| std::env::var(key).ok(); let headers = string_headers(request.extra_headers)?; let auth_strategy = config.auth_strategy(); @@ -120,7 +118,7 @@ impl OcrLifecycleHooks { custom_llm_provider: &str, url: &str, body: Value, - ) -> CoreResult { + ) -> Result { if self.guardrail_runner.is_empty() { return Ok(body); } @@ -217,7 +215,7 @@ impl CallLifecycleHooks for OcrLi fn async_log_failure_event<'a>( &'a self, context: &'a CallLifecycleContext, - error: &'a CoreError, + error: &'a Error, timing: &'a CallLifecycleTiming, ) -> Self::FailureFuture<'a> { Box::pin(async move { @@ -278,19 +276,19 @@ fn guardrail_context(metadata: &RequestMetadata) -> GuardrailContext { fn parse_ocr_pre_call_guardrail_request( request: GuardrailRequest, -) -> CoreResult<(Value, Map)> { +) -> Result<(Value, Map), Error> { let Value::Object(mut data) = request.data else { - return Err(CoreError::InvalidRequest( + return Err(Error::InvalidRequest( "OCR pre_call guardrail must return an object".to_string(), )); }; let document = data.remove("document").ok_or_else(|| { - CoreError::InvalidRequest("OCR pre_call guardrail removed document".to_string()) + Error::InvalidRequest("OCR pre_call guardrail removed document".to_string()) })?; let optional_params = match data.remove("optional_params") { Some(Value::Object(params)) => params, Some(_) => { - return Err(CoreError::InvalidRequest( + return Err(Error::InvalidRequest( "OCR pre_call guardrail optional_params must be an object".to_string(), )); } @@ -299,33 +297,32 @@ fn parse_ocr_pre_call_guardrail_request( Ok((document, optional_params)) } -fn parse_ocr_during_call_guardrail_request(request: GuardrailRequest) -> CoreResult { +fn parse_ocr_during_call_guardrail_request(request: GuardrailRequest) -> Result { let Value::Object(mut data) = request.data else { - return Err(CoreError::InvalidRequest( + return Err(Error::InvalidRequest( "OCR during_call guardrail must return an object".to_string(), )); }; - data.remove("body").ok_or_else(|| { - CoreError::InvalidRequest("OCR during_call guardrail removed body".to_string()) - }) + data.remove("body") + .ok_or_else(|| Error::InvalidRequest("OCR during_call guardrail removed body".to_string())) } -fn guardrail_error_to_core_error(error: GuardrailError) -> CoreError { - CoreError::InvalidRequest(format!("{}: {}", error.kind, error.message)) +fn guardrail_error_to_core_error(error: GuardrailError) -> Error { + Error::InvalidRequest(format!("{}: {}", error.kind, error.message)) } -fn core_error_kind(error: &CoreError) -> &'static str { +fn core_error_kind(error: &Error) -> &'static str { match error { - CoreError::Auth(_) => "AuthError", - CoreError::InvalidProvider(_) => "InvalidProvider", - CoreError::InvalidRequest(_) => "InvalidRequest", - CoreError::InvalidType { .. } => "InvalidType", - CoreError::MissingField(_) => "MissingField", - CoreError::Http { .. } => "HttpError", - CoreError::InvalidResponse(_) => "InvalidResponse", - CoreError::Network(_) => "NetworkError", - CoreError::Connect(_) => "ConnectError", - CoreError::Routing(_) => "RoutingError", - CoreError::Unsupported(_) => "UnsupportedRequest", + Error::Auth(_) => "AuthError", + Error::InvalidProvider(_) => "InvalidProvider", + Error::InvalidRequest(_) => "InvalidRequest", + Error::InvalidType { .. } => "InvalidType", + Error::MissingField(_) => "MissingField", + Error::Http { .. } => "HttpError", + Error::InvalidResponse(_) => "InvalidResponse", + Error::Network(_) => "NetworkError", + Error::Connect(_) => "ConnectError", + Error::Routing(_) => "RoutingError", + Error::Unsupported(_) => "UnsupportedRequest", } } diff --git a/litellm-rust/crates/ai-gateway/src/ocr/mod.rs b/litellm-rust/crates/ai-gateway/src/ocr/mod.rs index c4c13e2300c..b59ab626fd3 100644 --- a/litellm-rust/crates/ai-gateway/src/ocr/mod.rs +++ b/litellm-rust/crates/ai-gateway/src/ocr/mod.rs @@ -1,4 +1,4 @@ -use litellm_core::CoreResult; +use litellm_core::Error; use litellm_core::call_lifecycle::CallLifecycle; use serde_json::Value; @@ -13,7 +13,7 @@ pub use types::OcrRequest; use handler::execute_ocr_provider_call; use prepare::{PreparedOcrCall, prepare_ocr_call}; -pub async fn ocr(request: OcrRequest<'_>) -> CoreResult { +pub async fn ocr(request: OcrRequest<'_>) -> Result { let PreparedOcrCall { request, hooks } = prepare_ocr_call(request); CallLifecycle::default() .run_request(request, &hooks, execute_ocr_provider_call) diff --git a/litellm-rust/crates/ai-gateway/src/ocr/tests.rs b/litellm-rust/crates/ai-gateway/src/ocr/tests.rs index bb2a6b06501..8c3f0425149 100644 --- a/litellm-rust/crates/ai-gateway/src/ocr/tests.rs +++ b/litellm-rust/crates/ai-gateway/src/ocr/tests.rs @@ -1,7 +1,7 @@ use std::sync::{Arc, Mutex}; use std::time::Duration; -use litellm_core::error::CoreError; +use litellm_core::error::Error; use litellm_core::ocr::transformation::OcrResponseHandling; use serde_json::{Map, Value, json}; use tokio::io::{AsyncReadExt, AsyncWriteExt}; @@ -395,7 +395,7 @@ async fn ocr_lifecycle_runs_failure_hook_on_provider_error() { .await .expect_err("provider error propagates"); - assert!(matches!(err, CoreError::Http { status: 500, .. })); + assert!(matches!(err, Error::Http { status: 500, .. })); server.await.expect("server task completes"); assert_eq!( logger.events(), @@ -439,7 +439,7 @@ async fn ocr_lifecycle_pre_call_block_skips_provider_socket() { .await .expect_err("guardrail blocks request"); - assert!(matches!(err, CoreError::InvalidRequest(_))); + assert!(matches!(err, Error::InvalidRequest(_))); assert_eq!(guardrail.events(), vec!["async_pre_call_hook"]); assert_eq!( logger.events(), @@ -607,7 +607,7 @@ fn string_headers_rejects_non_string_values() { let err = string_headers(Some(headers)).expect_err("non-string header rejected"); assert_eq!( err, - CoreError::InvalidRequest( + Error::InvalidRequest( "OCR extra_headers.x-retry-count must be a string, got number".to_string() ) ); diff --git a/litellm-rust/crates/ai-gateway/src/python/config.rs b/litellm-rust/crates/ai-gateway/src/python/config.rs index c028d3d6b51..d5a4dd69c8d 100644 --- a/litellm-rust/crates/ai-gateway/src/python/config.rs +++ b/litellm-rust/crates/ai-gateway/src/python/config.rs @@ -6,33 +6,31 @@ //! (and recorded in [`crate::gil`]); the realtime hot path never touches Python. //! //! Compiled only under the `python-config` feature. - -use litellm_core::CoreResult; -use litellm_core::error::CoreError; +use litellm_core::error::Error; use litellm_core::router::{Deployment, Router}; use pyo3::prelude::*; use crate::gil; /// Load the router's `model_list` from `config_path` via the Python reader. -pub fn load_router_from_config(config_path: &str) -> CoreResult { +pub fn load_router_from_config(config_path: &str) -> Result { gil::record_acquisition(); Python::attach(|py| { let model_list = py .import("litellm.proxy.read_model_list") .and_then(|module| module.getattr("read_model_list")) .and_then(|reader| reader.call1((config_path,))) - .map_err(|err| CoreError::Routing(format!("read_model_list failed: {err}")))?; + .map_err(|err| Error::Routing(format!("read_model_list failed: {err}")))?; let model_list_json: String = py .import("json") .and_then(|json| json.getattr("dumps")) .and_then(|dumps| dumps.call1((model_list,))) .and_then(|encoded| encoded.extract()) - .map_err(|err| CoreError::Routing(format!("serializing model_list failed: {err}")))?; + .map_err(|err| Error::Routing(format!("serializing model_list failed: {err}")))?; let deployments: Vec = serde_json::from_str(&model_list_json) - .map_err(|err| CoreError::Routing(format!("parsing model_list failed: {err}")))?; + .map_err(|err| Error::Routing(format!("parsing model_list failed: {err}")))?; Ok(Router::new(deployments)) }) diff --git a/litellm-rust/crates/ai-gateway/src/routes/messages/mod.rs b/litellm-rust/crates/ai-gateway/src/routes/messages/mod.rs index 7e38d10c6ff..e9f8c477f36 100644 --- a/litellm-rust/crates/ai-gateway/src/routes/messages/mod.rs +++ b/litellm-rust/crates/ai-gateway/src/routes/messages/mod.rs @@ -9,7 +9,7 @@ use axum::http::StatusCode; use axum::http::header::{CACHE_CONTROL, CONTENT_TYPE, HeaderMap, HeaderValue}; use axum::response::{IntoResponse, Response}; use axum::routing::post; -use litellm_core::CoreError; +use litellm_core::Error; use serde_json::{Map, Value}; use crate::auth::RequireMasterKey; @@ -46,7 +46,7 @@ fn stream_response(upstream: reqwest::Response) -> Result Result Result>, CoreError> { +fn forwarded_headers(headers: &HeaderMap) -> Result>, Error> { let forwarded = headers .iter() .filter(|(name, _)| { @@ -74,19 +74,19 @@ fn forwarded_headers(headers: &HeaderMap) -> Result>, }) .map(|(name, value)| { let value = value.to_str().map_err(|_| { - CoreError::InvalidRequest(format!("invalid value for header {}", name.as_str())) + Error::InvalidRequest(format!("invalid value for header {}", name.as_str())) })?; Ok((name.to_string(), Value::String(value.to_string()))) }) - .collect::, CoreError>>()?; + .collect::, Error>>()?; Ok((!forwarded.is_empty()).then_some(forwarded)) } #[derive(Debug)] -struct MessagesRouteError(CoreError); +struct MessagesRouteError(Error); -impl From for MessagesRouteError { - fn from(error: CoreError) -> Self { +impl From for MessagesRouteError { + fn from(error: Error) -> Self { Self(error) } } @@ -94,28 +94,28 @@ impl From for MessagesRouteError { impl IntoResponse for MessagesRouteError { fn into_response(self) -> Response { let (status, message) = match self.0 { - CoreError::InvalidRequest(message) => (StatusCode::BAD_REQUEST, message), - CoreError::InvalidProvider(_) | CoreError::Routing(_) => ( + Error::InvalidRequest(message) => (StatusCode::BAD_REQUEST, message), + Error::InvalidProvider(_) | Error::Routing(_) => ( StatusCode::NOT_FOUND, "no messages deployment is configured for this model".to_string(), ), - CoreError::Auth(_) => ( + Error::Auth(_) => ( StatusCode::BAD_GATEWAY, "messages provider authentication failed".to_string(), ), - CoreError::Http { .. } - | CoreError::Network(_) - | CoreError::Connect(_) - | CoreError::InvalidResponse(_) - | CoreError::InvalidType { .. } - | CoreError::MissingField(_) => ( + Error::Http { .. } + | Error::Network(_) + | Error::Connect(_) + | Error::InvalidResponse(_) + | Error::InvalidType { .. } + | Error::MissingField(_) => ( StatusCode::BAD_GATEWAY, "messages provider request failed".to_string(), ), // The gateway has no Python implementation to decline to, so a // request the core cannot serve is reported to the caller. The // reason is a fixed internal string, never provider content. - CoreError::Unsupported(reason) => ( + Error::Unsupported(reason) => ( StatusCode::BAD_REQUEST, format!("messages request is not supported: {reason}"), ), diff --git a/litellm-rust/crates/ai-gateway/src/routes/messages/service.rs b/litellm-rust/crates/ai-gateway/src/routes/messages/service.rs index 5f4c5fe8de4..4fd29db05d6 100644 --- a/litellm-rust/crates/ai-gateway/src/routes/messages/service.rs +++ b/litellm-rust/crates/ai-gateway/src/routes/messages/service.rs @@ -1,10 +1,10 @@ use std::sync::Arc; +use litellm_core::Error; use litellm_core::constants::ANTHROPIC_MESSAGES_PROVIDER; use litellm_core::messages::types::MessagesRequest; use litellm_core::messages::{messages, messages_stream}; use litellm_core::router::Router; -use litellm_core::{CoreError, CoreResult}; use serde_json::{Map, Value}; pub(crate) enum MessagesResponse { @@ -16,16 +16,16 @@ pub async fn run( router: &Arc, body: Value, extra_headers: Option>, -) -> CoreResult { +) -> Result { let model = body .get("model") .and_then(Value::as_str) .map(str::trim) .filter(|model| !model.is_empty()) - .ok_or_else(|| CoreError::InvalidRequest("messages body requires a model".to_string()))?; - let deployment = router.get_available_deployment(model).ok_or_else(|| { - CoreError::Routing(format!("no deployment available for model '{model}'")) - })?; + .ok_or_else(|| Error::InvalidRequest("messages body requires a model".to_string()))?; + let deployment = router + .get_available_deployment(model) + .ok_or_else(|| Error::Routing(format!("no deployment available for model '{model}'")))?; let provider_model = deployment.litellm_params.model.as_str(); let upstream_model = provider_model .split_once('/') @@ -37,7 +37,7 @@ pub async fn run( }; let mut body = body; body.as_object_mut() - .ok_or_else(|| CoreError::InvalidRequest("messages body must be an object".to_string()))? + .ok_or_else(|| Error::InvalidRequest("messages body must be an object".to_string()))? .insert( "model".to_string(), Value::String(upstream_model.to_string()), @@ -60,6 +60,6 @@ pub async fn run( serde_json::to_value(response) .map(MessagesResponse::Json) .map_err(|err| { - CoreError::InvalidResponse(format!("failed to serialize messages response: {err}")) + Error::InvalidResponse(format!("failed to serialize messages response: {err}")) }) } diff --git a/litellm-rust/crates/ai-gateway/src/routes/realtime/service.rs b/litellm-rust/crates/ai-gateway/src/routes/realtime/service.rs index 4ae8cfe7379..b8ee77c4269 100644 --- a/litellm-rust/crates/ai-gateway/src/routes/realtime/service.rs +++ b/litellm-rust/crates/ai-gateway/src/routes/realtime/service.rs @@ -11,8 +11,7 @@ use std::time::Duration; use crate::io::realtime_pool::{RealtimePool, upstream_key}; use futures_util::{Sink, Stream}; -use litellm_core::CoreResult; -use litellm_core::error::CoreError; +use litellm_core::error::Error; use litellm_core::realtime::types::RealtimeEvent; use litellm_core::router::Router; @@ -29,15 +28,15 @@ pub async fn run( observe: impl FnMut(&RealtimeEvent) + Send, client_in: In, client_out: Out, -) -> CoreResult<()> +) -> Result<(), Error> where In: Stream + Unpin + Send, Out: Sink + Unpin + Send, >::Error: std::fmt::Display, { - let deployment = router.get_available_deployment(model).ok_or_else(|| { - CoreError::Routing(format!("no deployment available for model '{model}'")) - })?; + let deployment = router + .get_available_deployment(model) + .ok_or_else(|| Error::Routing(format!("no deployment available for model '{model}'")))?; let params = &deployment.litellm_params; // Strip a leading `openai/` so the OpenAI-only realtime fn gets the bare model. let provider_model = params diff --git a/litellm-rust/crates/ai-gateway/src/routes/responses/service.rs b/litellm-rust/crates/ai-gateway/src/routes/responses/service.rs index 165c95695d3..e8f840c0c8e 100644 --- a/litellm-rust/crates/ai-gateway/src/routes/responses/service.rs +++ b/litellm-rust/crates/ai-gateway/src/routes/responses/service.rs @@ -2,13 +2,13 @@ use std::sync::Arc; use std::time::Duration; use futures_util::{Sink, Stream}; +use litellm_core::Error; use litellm_core::call_lifecycle::{CallLifecycle, CallLifecycleContext}; use litellm_core::responses::instrumentation::{ ResponsesWsCallbackPayload, ResponsesWsInstrumentation, ResponsesWsLogOutcome, ResponsesWsMetadata, }; use litellm_core::responses::types::ResponsesWsEvent; -use litellm_core::{CoreError, CoreResult}; use crate::integrations::custom_logger::{ CallbackTiming, CallbackValue, CustomLogger, CustomLoggerRunner, LoggingError, ModelCallDetails, @@ -26,22 +26,22 @@ pub async fn run( metadata: RequestMetadata, client_in: In, client_out: Out, -) -> CoreResult<()> +) -> Result<(), Error> where In: Stream + Unpin + Send, Out: Sink + Unpin + Send, Out::Error: std::fmt::Display, { - let deployment = router.get_available_deployment(model).ok_or_else(|| { - CoreError::Routing(format!("no deployment available for model '{model}'")) - })?; + let deployment = router + .get_available_deployment(model) + .ok_or_else(|| Error::Routing(format!("no deployment available for model '{model}'")))?; let params = &deployment.litellm_params; let provider_model = params .model .strip_prefix("openai/") .unwrap_or(¶ms.model); if params.model.contains('/') && !params.model.starts_with("openai/") { - return Err(CoreError::InvalidProvider( + return Err(Error::InvalidProvider( "Responses WebSocket route supports OpenAI deployments only".to_string(), )); } diff --git a/litellm-rust/crates/core/src/audio_transcription/transformation.rs b/litellm-rust/crates/core/src/audio_transcription/transformation.rs index eab34c13843..16a28fbcac0 100644 --- a/litellm-rust/crates/core/src/audio_transcription/transformation.rs +++ b/litellm-rust/crates/core/src/audio_transcription/transformation.rs @@ -1,7 +1,6 @@ +use crate::Error; use serde_json::{Map, Value}; -use crate::CoreResult; - use super::types::{AudioTranscriptionRequestData, AudioTranscriptionResponseData}; #[derive(Clone, Debug, PartialEq, Eq)] @@ -32,13 +31,13 @@ pub trait AudioTranscriptionProviderConfig: Sync { model: &str, audio: Value, optional_params: Map, - ) -> CoreResult; + ) -> Result; fn transform_transcription_response( &self, model: &str, response_json: Value, - ) -> CoreResult; + ) -> Result; fn complete_url( &self, @@ -46,12 +45,12 @@ pub trait AudioTranscriptionProviderConfig: Sync { model: &str, optional_params: &Map, env_lookup: &dyn Fn(&str) -> Option, - ) -> CoreResult; + ) -> Result; fn auth_strategy( &self, model: &str, optional_params: &Map, env_lookup: &dyn Fn(&str) -> Option, - ) -> CoreResult; + ) -> Result; } diff --git a/litellm-rust/crates/core/src/call_lifecycle/mod.rs b/litellm-rust/crates/core/src/call_lifecycle/mod.rs index d9b68a1b726..637c156e192 100644 --- a/litellm-rust/crates/core/src/call_lifecycle/mod.rs +++ b/litellm-rust/crates/core/src/call_lifecycle/mod.rs @@ -1,7 +1,7 @@ use std::future::Future; use std::time::{Instant, SystemTime, UNIX_EPOCH}; -use crate::{CoreError, CoreResult}; +use crate::Error; pub mod types; @@ -11,14 +11,14 @@ pub use types::{ }; pub trait CallLifecycleHooks: Send + Sync { - type PreCallFuture<'a>: Future> + Send + 'a + type PreCallFuture<'a>: Future> + Send + 'a where Self: 'a, InitialReq: 'a, ProviderReq: 'a, Resp: 'a; - type DuringCallFuture<'a>: Future> + Send + 'a + type DuringCallFuture<'a>: Future> + Send + 'a where Self: 'a, InitialReq: 'a, @@ -56,7 +56,7 @@ pub trait CallLifecycleHooks: Send + Sync { fn async_log_failure_event<'a>( &'a self, context: &'a CallLifecycleContext, - error: &'a CoreError, + error: &'a Error, timing: &'a CallLifecycleTiming, ) -> Self::FailureFuture<'a>; } @@ -86,12 +86,12 @@ impl<'a> CallLifecycle<'a> { request: InitialReq, hooks: &Hooks, provider_call: ProviderCall, - ) -> CoreResult + ) -> Result where InitialReq: CallLifecycleRequest, Hooks: CallLifecycleHooks, ProviderCall: FnOnce(ProviderReq) -> ProviderFuture, - ProviderFuture: Future>, + ProviderFuture: Future>, { let context = request.lifecycle_context(); self.run(context, request, hooks, provider_call).await @@ -103,11 +103,11 @@ impl<'a> CallLifecycle<'a> { request: InitialReq, hooks: &Hooks, provider_call: ProviderCall, - ) -> CoreResult + ) -> Result where Hooks: CallLifecycleHooks, ProviderCall: FnOnce(ProviderReq) -> ProviderFuture, - ProviderFuture: Future>, + ProviderFuture: Future>, { let call_start = epoch_seconds(); let mut phases = Vec::new(); @@ -166,7 +166,7 @@ impl<'a> CallLifecycle<'a> { &self, context: &CallLifecycleContext, hooks: &Hooks, - error: &CoreError, + error: &Error, call_start: f64, phases: &mut Vec, ) where @@ -251,8 +251,8 @@ mod tests { } impl CallLifecycleHooks for RecordingHooks { - type PreCallFuture<'a> = BoxFuture<'a, CoreResult>; - type DuringCallFuture<'a> = BoxFuture<'a, CoreResult>; + type PreCallFuture<'a> = BoxFuture<'a, Result>; + type DuringCallFuture<'a> = BoxFuture<'a, Result>; type SuccessFuture<'a> = BoxFuture<'a, ()>; type FailureFuture<'a> = BoxFuture<'a, ()>; @@ -294,7 +294,7 @@ mod tests { fn async_log_failure_event<'a>( &'a self, _context: &'a CallLifecycleContext, - _error: &'a CoreError, + _error: &'a Error, _timing: &'a CallLifecycleTiming, ) -> Self::FailureFuture<'a> { Box::pin(async move { @@ -304,8 +304,8 @@ mod tests { } impl CallLifecycleHooks for RecordingHooks { - type PreCallFuture<'a> = BoxFuture<'a, CoreResult>; - type DuringCallFuture<'a> = BoxFuture<'a, CoreResult>; + type PreCallFuture<'a> = BoxFuture<'a, Result>; + type DuringCallFuture<'a> = BoxFuture<'a, Result>; type SuccessFuture<'a> = BoxFuture<'a, ()>; type FailureFuture<'a> = BoxFuture<'a, ()>; @@ -345,7 +345,7 @@ mod tests { fn async_log_failure_event<'a>( &'a self, _context: &'a CallLifecycleContext, - _error: &'a CoreError, + _error: &'a Error, _timing: &'a CallLifecycleTiming, ) -> Self::FailureFuture<'a> { Box::pin(async move { @@ -383,13 +383,13 @@ mod tests { "request".to_string(), &hooks, |_request| async move { - Err::(CoreError::Network("provider down".to_string())) + Err::(Error::Network("provider down".to_string())) }, ) .await .expect_err("call fails"); - assert_eq!(error, CoreError::Network("provider down".to_string())); + assert_eq!(error, Error::Network("provider down".to_string())); assert_eq!(hooks.events(), vec!["pre_call", "during_call", "failure"]); } diff --git a/litellm-rust/crates/core/src/chat_completions/common_utils.rs b/litellm-rust/crates/core/src/chat_completions/common_utils.rs index 36eaf242a5a..ca51471eb7c 100644 --- a/litellm-rust/crates/core/src/chat_completions/common_utils.rs +++ b/litellm-rust/crates/core/src/chat_completions/common_utils.rs @@ -1,8 +1,7 @@ -use serde_json::{Map, Value}; - -use crate::error::CoreResult; +use crate::Error; use crate::http_utils::string_headers as shared_string_headers; use crate::providers::anthropic::chat_completions::transformation::ANTHROPIC_CHAT_COMPLETIONS_CONFIG; +use serde_json::{Map, Value}; use super::transformation::ChatCompletionsProviderConfig; @@ -23,6 +22,6 @@ pub(super) fn chat_completions_provider_config( pub(super) fn string_headers( extra_headers: Option>, -) -> CoreResult> { +) -> Result, Error> { shared_string_headers(HEADER_CONTEXT, extra_headers) } diff --git a/litellm-rust/crates/core/src/chat_completions/handler.rs b/litellm-rust/crates/core/src/chat_completions/handler.rs index afc4529fd26..7e2731442cc 100644 --- a/litellm-rust/crates/core/src/chat_completions/handler.rs +++ b/litellm-rust/crates/core/src/chat_completions/handler.rs @@ -1,6 +1,6 @@ use serde_json::Value; -use crate::error::{CoreError, CoreResult}; +use crate::error::Error; use crate::http_utils::truncate_error_body; use super::client::http_client; @@ -11,9 +11,9 @@ use super::types::{ pub(super) async fn execute_chat_completions_provider_call( request: ProviderChatCompletionsRequest, -) -> CoreResult { +) -> Result { let body = serde_json::to_vec(&request.body).map_err(|err| { - CoreError::InvalidRequest(format!( + Error::InvalidRequest(format!( "failed to serialize chat completions request: {err}" )) })?; @@ -32,9 +32,9 @@ pub(super) async fn execute_chat_completions_provider_call( // so the host can still serve it. Everything else here, a timeout // above all, may have reached the provider and been answered. if err.is_connect() || err.is_builder() { - CoreError::Connect(err.to_string()) + Error::Connect(err.to_string()) } else { - CoreError::Network(err.to_string()) + Error::Network(err.to_string()) } })?; @@ -42,17 +42,17 @@ pub(super) async fn execute_chat_completions_provider_call( let text = response .text() .await - .map_err(|err| CoreError::Network(err.to_string()))?; + .map_err(|err| Error::Network(err.to_string()))?; if !status.is_success() { - return Err(CoreError::Http { + return Err(Error::Http { status: status.as_u16(), body: truncate_error_body(&text), }); } let body: Value = serde_json::from_str(&text).map_err(|err| { - CoreError::InvalidResponse(format!("invalid chat completions response JSON: {err}")) + Error::InvalidResponse(format!("invalid chat completions response JSON: {err}")) })?; request .config @@ -69,10 +69,10 @@ pub(super) async fn execute_chat_completions_provider_call( /// second kind has already been billed, and a host that keeps a reference /// implementation must not retry those, so collapse them to one variant that /// can only mean the provider was already called. -pub(super) fn as_response_error(err: CoreError) -> CoreError { +pub(super) fn as_response_error(err: Error) -> Error { match err { - already @ (CoreError::InvalidResponse(_) | CoreError::Http { .. }) => already, - other => CoreError::InvalidResponse(other.to_string()), + already @ (Error::InvalidResponse(_) | Error::Http { .. }) => already, + other => Error::InvalidResponse(other.to_string()), } } @@ -80,7 +80,7 @@ pub(super) fn as_response_error(err: CoreError) -> CoreError { pub(super) async fn signed_headers( request: &ProviderChatCompletionsRequest, body: &[u8], -) -> CoreResult> { +) -> Result, Error> { use std::collections::BTreeMap; use std::time::SystemTime; @@ -101,7 +101,7 @@ pub(super) async fn signed_headers( .iter() .any(|(name, _)| is_sigv4_computed_header(name)) { - return Err(CoreError::Unsupported( + return Err(Error::Unsupported( "request forwards a header AWS SigV4 computes", )); } @@ -137,9 +137,9 @@ pub(super) async fn signed_headers( pub(super) async fn signed_headers( request: &ProviderChatCompletionsRequest, _body: &[u8], -) -> CoreResult> { +) -> Result, Error> { match &request.auth { - ChatCompletionsAuth::AwsSigV4 { .. } => Err(CoreError::Unsupported( + ChatCompletionsAuth::AwsSigV4 { .. } => Err(Error::Unsupported( "AWS SigV4 requires the bedrock-auth feature", )), _ => Ok(request.upstream_headers.clone()), diff --git a/litellm-rust/crates/core/src/chat_completions/mod.rs b/litellm-rust/crates/core/src/chat_completions/mod.rs index f30ac1a24bf..0d009d36d16 100644 --- a/litellm-rust/crates/core/src/chat_completions/mod.rs +++ b/litellm-rust/crates/core/src/chat_completions/mod.rs @@ -6,6 +6,7 @@ //! credentials, and it resolves the provider, translates the conversation, //! calls the provider, and returns a typed OpenAI-shaped response. +use crate::Error; mod client; mod common_utils; pub mod conversation; @@ -17,15 +18,13 @@ pub mod types; use serde_json::{Map, Value}; -use crate::error::CoreResult; - use handler::execute_chat_completions_provider_call; use prepare::{parse_messages, prepare_chat_completions_call, resolve_provider_config}; use types::{ChatCompletionsRequest, ChatCompletionsResponse}; pub async fn chat_completions( request: ChatCompletionsRequest<'_>, -) -> CoreResult { +) -> Result { execute_chat_completions_provider_call(prepare_chat_completions_call(request)?).await } diff --git a/litellm-rust/crates/core/src/chat_completions/prepare.rs b/litellm-rust/crates/core/src/chat_completions/prepare.rs index 1e1c8d1bafd..142b2f2aaed 100644 --- a/litellm-rust/crates/core/src/chat_completions/prepare.rs +++ b/litellm-rust/crates/core/src/chat_completions/prepare.rs @@ -1,6 +1,6 @@ use serde_json::Value; -use crate::error::{CoreError, CoreResult}; +use crate::error::Error; use crate::http_utils::has_header; use crate::routing_utils::provider::{CustomLlmProvider, get_custom_llm_provider}; @@ -11,7 +11,7 @@ use super::types::{ChatCompletionsRequest, ChatMessage, ProviderChatCompletionsR pub(super) fn resolve_provider_config<'a>( model: &'a str, custom_llm_provider: Option<&'a str>, -) -> CoreResult<(String, &'static dyn ChatCompletionsProviderConfig)> { +) -> Result<(String, &'static dyn ChatCompletionsProviderConfig), Error> { let provider_info = get_custom_llm_provider(model, custom_llm_provider) .or_else(|| { custom_llm_provider.map(|provider| CustomLlmProvider { @@ -20,35 +20,34 @@ pub(super) fn resolve_provider_config<'a>( }) }) .ok_or_else(|| { - CoreError::InvalidProvider( + Error::InvalidProvider( "unable to resolve custom_llm_provider for chat completions request".to_string(), ) })?; let config = chat_completions_provider_config(provider_info.custom_llm_provider) - .ok_or_else(|| CoreError::InvalidProvider(provider_info.custom_llm_provider.to_string()))?; + .ok_or_else(|| Error::InvalidProvider(provider_info.custom_llm_provider.to_string()))?; Ok((provider_info.model.to_string(), config)) } -pub(super) fn parse_messages(messages: Value) -> CoreResult> { - serde_json::from_value(messages).map_err(|err| { - CoreError::InvalidRequest(format!("invalid chat completions messages: {err}")) - }) +pub(super) fn parse_messages(messages: Value) -> Result, Error> { + serde_json::from_value(messages) + .map_err(|err| Error::InvalidRequest(format!("invalid chat completions messages: {err}"))) } pub(super) fn prepare_chat_completions_call( request: ChatCompletionsRequest<'_>, -) -> CoreResult { +) -> Result { let (model, config) = resolve_provider_config(request.model, request.custom_llm_provider)?; let env_lookup = |key: &str| std::env::var(key).ok(); let messages = parse_messages(request.messages)?; if messages.is_empty() { - return Err(CoreError::InvalidRequest( + return Err(Error::InvalidRequest( "chat completions requires at least one message".to_string(), )); } if let Some(reason) = config.unsupported_reason(&messages, &request.optional_params) { - return Err(CoreError::Unsupported(reason.0)); + return Err(Error::Unsupported(reason.0)); } let mut headers = string_headers(request.extra_headers)?; diff --git a/litellm-rust/crates/core/src/chat_completions/tests.rs b/litellm-rust/crates/core/src/chat_completions/tests.rs index e2383723cb0..2858d180e27 100644 --- a/litellm-rust/crates/core/src/chat_completions/tests.rs +++ b/litellm-rust/crates/core/src/chat_completions/tests.rs @@ -1,6 +1,6 @@ use serde_json::{Map, Value, json}; -use crate::error::CoreError; +use crate::error::Error; use super::prepare::prepare_chat_completions_call; use super::transformation::ChatCompletionsAuth; @@ -29,7 +29,7 @@ fn request<'a>( /// `ProviderChatCompletionsRequest` deliberately has no `Debug` (its headers /// carry resolved credentials), so unwrap the failure case by hand. -fn decline(request: ChatCompletionsRequest<'_>) -> CoreError { +fn decline(request: ChatCompletionsRequest<'_>) -> Error { match prepare_chat_completions_call(request) { Err(error) => error, Ok(prepared) => panic!("expected a decline, prepared a call to {}", prepared.url), @@ -196,7 +196,7 @@ fn declines_an_unsupported_request_before_resolving_credentials() { call.api_key = None; // No api_key is set and no env is consulted: the gate must run first, so the // error is the decline rather than a missing-credential error. - assert_eq!(decline(call), CoreError::Unsupported("streaming")); + assert_eq!(decline(call), Error::Unsupported("streaming")); } #[test] @@ -208,7 +208,7 @@ fn rejects_an_unknown_provider() { json!([{"role": "user", "content": "hi"}]), json!({}), )), - CoreError::InvalidProvider("openai".to_string()) + Error::InvalidProvider("openai".to_string()) ); } @@ -221,7 +221,7 @@ fn rejects_a_model_with_no_resolvable_provider() { json!([{"role": "user", "content": "hi"}]), json!({}), )), - CoreError::InvalidProvider(_) + Error::InvalidProvider(_) )); } @@ -234,7 +234,7 @@ fn rejects_an_empty_or_malformed_message_list() { json!([]), json!({}), )), - CoreError::InvalidRequest("chat completions requires at least one message".to_string()) + Error::InvalidRequest("chat completions requires at least one message".to_string()) ); assert!(matches!( decline(request( @@ -243,7 +243,7 @@ fn rejects_an_empty_or_malformed_message_list() { json!("not a list"), json!({}), )), - CoreError::InvalidRequest(_) + Error::InvalidRequest(_) )); } @@ -258,7 +258,7 @@ fn rejects_non_string_extra_headers() { call.extra_headers = Some(Map::from_iter([("x-trace".to_string(), json!(7))])); assert_eq!( decline(call), - CoreError::InvalidRequest( + Error::InvalidRequest( "chat completions extra_headers.x-trace must be a string, got number".to_string() ) ); @@ -374,7 +374,7 @@ async fn a_forwarded_header_the_signer_computes_declines_to_python() { .await .expect_err("{forwarded} should decline instead of being signed"); assert!( - matches!(error, CoreError::Unsupported(_)), + matches!(error, Error::Unsupported(_)), "{forwarded} declined as {error:?}, which the host would not fall back on" ); } @@ -727,7 +727,7 @@ mod round_trip { .expect_err("response cannot be normalized"); handle.await.expect("server task"); assert!( - matches!(err, CoreError::InvalidResponse(_)), + matches!(err, Error::InvalidResponse(_)), "expected a post-send error, got {err:?}" ); } @@ -745,7 +745,7 @@ mod round_trip { .expect_err("response cannot be normalized"); handle.await.expect("server task"); assert!( - matches!(err, CoreError::InvalidResponse(_)), + matches!(err, Error::InvalidResponse(_)), "expected a post-send error, got {err:?}" ); } @@ -763,7 +763,7 @@ mod round_trip { .expect_err("upstream rejects"); handle.await.expect("server task"); assert!( - matches!(err, CoreError::Http { status: 429, .. }), + matches!(err, Error::Http { status: 429, .. }), "expected a 429, got {err:?}" ); } @@ -787,7 +787,7 @@ mod round_trip { .await .expect_err("nothing is listening"); assert!( - matches!(err, CoreError::Connect(_)), + matches!(err, Error::Connect(_)), "expected a pre-send connect failure, got {err:?}" ); } @@ -797,24 +797,24 @@ mod round_trip { use crate::chat_completions::handler::as_response_error; for original in [ - CoreError::MissingField("usage"), - CoreError::Unsupported("non-text response content block"), - CoreError::InvalidRequest("whatever".to_string()), - CoreError::Auth("whatever".to_string()), + Error::MissingField("usage"), + Error::Unsupported("non-text response content block"), + Error::InvalidRequest("whatever".to_string()), + Error::Auth("whatever".to_string()), ] { let label = format!("{original:?}"); assert!( - matches!(as_response_error(original), CoreError::InvalidResponse(_)), + matches!(as_response_error(original), Error::InvalidResponse(_)), "{label} must not stay retryable once the provider has answered" ); } // An upstream status is already unambiguous, so it survives intact. assert!(matches!( - as_response_error(CoreError::Http { + as_response_error(Error::Http { status: 500, body: "boom".to_string() }), - CoreError::Http { status: 500, .. } + Error::Http { status: 500, .. } )); } } diff --git a/litellm-rust/crates/core/src/chat_completions/transformation.rs b/litellm-rust/crates/core/src/chat_completions/transformation.rs index a30ce9dc77c..a0868209305 100644 --- a/litellm-rust/crates/core/src/chat_completions/transformation.rs +++ b/litellm-rust/crates/core/src/chat_completions/transformation.rs @@ -1,7 +1,6 @@ +use crate::Error; use serde_json::{Map, Value}; -use crate::error::CoreResult; - use super::types::{ ChatCompletionsResponse, ChatMessage, ChatMessageContent, ProviderChatRequestData, ProviderChatResponseData, @@ -39,7 +38,7 @@ pub trait ChatCompletionsProviderConfig: Sync { model: &str, optional_params: &Map, env_lookup: &dyn Fn(&str) -> Option, - ) -> CoreResult; + ) -> Result; fn auth( &self, @@ -47,7 +46,7 @@ pub trait ChatCompletionsProviderConfig: Sync { model: &str, optional_params: &Map, env_lookup: &dyn Fn(&str) -> Option, - ) -> CoreResult; + ) -> Result; fn default_headers(&self) -> &'static [(&'static str, &'static str)] { &[("content-type", "application/json")] @@ -91,13 +90,13 @@ pub trait ChatCompletionsProviderConfig: Sync { model: &str, messages: Vec, optional_params: Map, - ) -> CoreResult; + ) -> Result; fn transform_response( &self, model: &str, response: ProviderChatResponseData, - ) -> CoreResult; + ) -> Result; } pub fn unsupported_param( diff --git a/litellm-rust/crates/core/src/error.rs b/litellm-rust/crates/core/src/error.rs index 739532f8cb5..db3fa2ec704 100644 --- a/litellm-rust/crates/core/src/error.rs +++ b/litellm-rust/crates/core/src/error.rs @@ -1,9 +1,7 @@ -use thiserror::Error; +use thiserror::Error as ThisError; -pub type CoreResult = Result; - -#[derive(Debug, Error, PartialEq, Eq)] -pub enum CoreError { +#[derive(Debug, ThisError, PartialEq, Eq)] +pub enum Error { #[error("expected {expected}, got {actual}")] InvalidType { expected: &'static str, diff --git a/litellm-rust/crates/core/src/http_utils.rs b/litellm-rust/crates/core/src/http_utils.rs index c541f50275b..10661fadf96 100644 --- a/litellm-rust/crates/core/src/http_utils.rs +++ b/litellm-rust/crates/core/src/http_utils.rs @@ -3,7 +3,7 @@ use serde_json::{Map, Value}; use crate::constants::UPSTREAM_ERROR_BODY_MAX_CHARS; -use crate::error::{CoreError, CoreResult, json_type_name}; +use crate::error::{Error, json_type_name}; /// Bound an upstream error body before it crosses a host boundary, so provider /// bodies stay data-minimized. @@ -18,7 +18,7 @@ pub fn truncate_error_body(body: &str) -> String { pub fn string_headers( context: &'static str, extra_headers: Option>, -) -> CoreResult> { +) -> Result, Error> { extra_headers .unwrap_or_default() .into_iter() @@ -27,7 +27,7 @@ pub fn string_headers( .as_str() .map(|value| (key.clone(), value.to_string())) .ok_or_else(|| { - CoreError::InvalidRequest(format!( + Error::InvalidRequest(format!( "{context} extra_headers.{key} must be a string, got {}", json_type_name(&value) )) @@ -81,7 +81,7 @@ mod tests { let err = string_headers("chat completions", Some(headers)).expect_err("non-string value"); assert_eq!( err, - CoreError::InvalidRequest( + Error::InvalidRequest( "chat completions extra_headers.x-trace must be a string, got number".to_string() ) ); diff --git a/litellm-rust/crates/core/src/lib.rs b/litellm-rust/crates/core/src/lib.rs index dce4a425ea0..0e18d24e5d8 100644 --- a/litellm-rust/crates/core/src/lib.rs +++ b/litellm-rust/crates/core/src/lib.rs @@ -13,4 +13,4 @@ pub mod responses; pub mod router; pub mod routing_utils; -pub use error::{CoreError, CoreResult}; +pub use error::Error; diff --git a/litellm-rust/crates/core/src/messages/common_utils.rs b/litellm-rust/crates/core/src/messages/common_utils.rs index a14dffbc1fe..8dfdb2e361a 100644 --- a/litellm-rust/crates/core/src/messages/common_utils.rs +++ b/litellm-rust/crates/core/src/messages/common_utils.rs @@ -1,9 +1,8 @@ -use serde_json::{Map, Value}; - -use crate::error::CoreResult; +use crate::Error; use crate::http_utils::string_headers as shared_string_headers; use crate::providers::anthropic::messages::transformation::ANTHROPIC_MESSAGES_CONFIG; use crate::providers::azure_ai::messages::transformation::AZURE_ANTHROPIC_MESSAGES_CONFIG; +use serde_json::{Map, Value}; use super::transformation::AnthropicMessagesProviderConfig; @@ -23,6 +22,6 @@ pub(super) fn messages_provider_config( pub(super) fn string_headers( extra_headers: Option>, -) -> CoreResult> { +) -> Result, Error> { shared_string_headers(HEADER_CONTEXT, extra_headers) } diff --git a/litellm-rust/crates/core/src/messages/handler.rs b/litellm-rust/crates/core/src/messages/handler.rs index 1c895f66eba..13a65d86131 100644 --- a/litellm-rust/crates/core/src/messages/handler.rs +++ b/litellm-rust/crates/core/src/messages/handler.rs @@ -1,5 +1,5 @@ use crate::constants::ANTHROPIC_MESSAGES_PROVIDER; -use crate::error::{CoreError, CoreResult}; +use crate::error::Error; use super::client::http_client; use super::common_utils::truncate_error_body; @@ -7,7 +7,7 @@ use super::types::{AnthropicMessagesResponse, ProviderMessagesRequest}; pub(super) async fn execute_messages_provider_call( request: ProviderMessagesRequest, -) -> CoreResult { +) -> Result { let mut request_builder = http_client().post(&request.url).json(&request.body); for (key, value) in &request.upstream_headers { request_builder = request_builder.header(key, value); @@ -19,32 +19,31 @@ pub(super) async fn execute_messages_provider_call( let response = request_builder .send() .await - .map_err(|err| CoreError::Network(err.to_string()))?; + .map_err(|err| Error::Network(err.to_string()))?; let status = response.status(); let text = response .text() .await - .map_err(|err| CoreError::Network(err.to_string()))?; + .map_err(|err| Error::Network(err.to_string()))?; if !status.is_success() { - return Err(CoreError::Http { + return Err(Error::Http { status: status.as_u16(), body: truncate_error_body(&text), }); } - let response = serde_json::from_str(&text).map_err(|err| { - CoreError::InvalidResponse(format!("invalid messages response JSON: {err}")) - })?; + let response = serde_json::from_str(&text) + .map_err(|err| Error::InvalidResponse(format!("invalid messages response JSON: {err}")))?; request.config.transform_response(&request.model, response) } pub(super) async fn execute_messages_provider_stream( request: ProviderMessagesRequest, -) -> CoreResult { +) -> Result { if request.provider != ANTHROPIC_MESSAGES_PROVIDER { - return Err(CoreError::InvalidRequest( + return Err(Error::InvalidRequest( "streaming messages is not supported for this provider".to_string(), )); } @@ -60,14 +59,14 @@ pub(super) async fn execute_messages_provider_stream( let response = request_builder .send() .await - .map_err(|err| CoreError::Network(err.to_string()))?; + .map_err(|err| Error::Network(err.to_string()))?; let status = response.status(); if !status.is_success() { let text = response .text() .await - .map_err(|err| CoreError::Network(err.to_string()))?; - return Err(CoreError::Http { + .map_err(|err| Error::Network(err.to_string()))?; + return Err(Error::Http { status: status.as_u16(), body: truncate_error_body(&text), }); diff --git a/litellm-rust/crates/core/src/messages/mod.rs b/litellm-rust/crates/core/src/messages/mod.rs index acb36d89daf..ee2877e61fc 100644 --- a/litellm-rust/crates/core/src/messages/mod.rs +++ b/litellm-rust/crates/core/src/messages/mod.rs @@ -7,6 +7,7 @@ //! is the streaming variant; it hands the raw upstream response back so a host //! can splice the event stream to its own caller. +use crate::Error; mod client; mod common_utils; mod handler; @@ -14,17 +15,15 @@ mod prepare; pub mod transformation; pub mod types; -use crate::error::CoreResult; - use handler::{execute_messages_provider_call, execute_messages_provider_stream}; use prepare::prepare_messages_call; use types::{AnthropicMessagesResponse, MessagesRequest}; -pub async fn messages(request: MessagesRequest<'_>) -> CoreResult { +pub async fn messages(request: MessagesRequest<'_>) -> Result { execute_messages_provider_call(prepare_messages_call(request)?).await } -pub async fn messages_stream(request: MessagesRequest<'_>) -> CoreResult { +pub async fn messages_stream(request: MessagesRequest<'_>) -> Result { execute_messages_provider_stream(prepare_messages_call(request)?).await } diff --git a/litellm-rust/crates/core/src/messages/prepare.rs b/litellm-rust/crates/core/src/messages/prepare.rs index 94b5b1eaed7..3b253ac3766 100644 --- a/litellm-rust/crates/core/src/messages/prepare.rs +++ b/litellm-rust/crates/core/src/messages/prepare.rs @@ -1,4 +1,4 @@ -use crate::error::{CoreError, CoreResult}; +use crate::error::Error; use crate::routing_utils::provider::{CustomLlmProvider, get_custom_llm_provider}; use super::common_utils::{has_bearer_auth, has_header, messages_provider_config, string_headers}; @@ -7,7 +7,7 @@ use super::types::{MessagesRequest, ProviderMessagesRequest}; pub(super) fn prepare_messages_call( request: MessagesRequest<'_>, -) -> CoreResult { +) -> Result { let provider_info = get_custom_llm_provider(request.model, request.custom_llm_provider) .or_else(|| { request @@ -18,7 +18,7 @@ pub(super) fn prepare_messages_call( }) }) .ok_or_else(|| { - CoreError::InvalidProvider( + Error::InvalidProvider( "unable to resolve custom_llm_provider for messages request".to_string(), ) })?; @@ -26,7 +26,7 @@ pub(super) fn prepare_messages_call( let provider = provider_info.custom_llm_provider; let config = messages_provider_config(provider) - .ok_or_else(|| CoreError::InvalidProvider(provider.to_string()))?; + .ok_or_else(|| Error::InvalidProvider(provider.to_string()))?; let env_lookup = |key: &str| std::env::var(key).ok(); let mut headers = string_headers(request.extra_headers)?; @@ -53,11 +53,11 @@ pub(super) fn prepare_messages_call( let url = config.complete_url(request.api_base, &model, &env_lookup)?; let typed_request = serde_json::from_value(request.body).map_err(|err| { - CoreError::InvalidRequest(format!("invalid Anthropic messages request: {err}")) + Error::InvalidRequest(format!("invalid Anthropic messages request: {err}")) })?; let transformed = config.transform_request(typed_request)?; let body = serde_json::to_value(transformed).map_err(|err| { - CoreError::InvalidRequest(format!( + Error::InvalidRequest(format!( "failed to serialize Anthropic messages request: {err}" )) })?; diff --git a/litellm-rust/crates/core/src/messages/tests.rs b/litellm-rust/crates/core/src/messages/tests.rs index 9fc1763683b..df9f7051011 100644 --- a/litellm-rust/crates/core/src/messages/tests.rs +++ b/litellm-rust/crates/core/src/messages/tests.rs @@ -4,7 +4,7 @@ use serde_json::{Map, Value, json}; use tokio::io::{AsyncReadExt, AsyncWriteExt}; use tokio::net::{TcpListener, TcpStream}; -use crate::error::CoreError; +use crate::error::Error; use super::common_utils::{ has_bearer_auth, has_header, messages_provider_config, string_headers, truncate_error_body, @@ -77,7 +77,7 @@ fn truncate_error_body_caps_long_payloads() { fn string_headers_rejects_non_string_values() { let headers = json!({"x-count": 3}).as_object().unwrap().clone(); let err = string_headers(Some(headers)).expect_err("non-string header rejected"); - assert!(matches!(err, CoreError::InvalidRequest(_))); + assert!(matches!(err, Error::InvalidRequest(_))); } #[test] @@ -341,7 +341,7 @@ async fn messages_requires_auth_when_no_key_and_no_header() { .await .expect_err("missing auth errors"); - assert!(matches!(err, CoreError::Auth(_))); + assert!(matches!(err, Error::Auth(_))); } #[tokio::test] @@ -420,7 +420,7 @@ async fn messages_maps_provider_error_status_to_http_error() { .await .expect_err("provider error propagates"); - assert!(matches!(err, CoreError::Http { status: 401, .. })); + assert!(matches!(err, Error::Http { status: 401, .. })); } #[tokio::test] @@ -437,5 +437,5 @@ async fn messages_rejects_unsupported_provider() { .await .expect_err("unsupported provider errors"); - assert!(matches!(err, CoreError::InvalidProvider(provider) if provider == "openai")); + assert!(matches!(err, Error::InvalidProvider(provider) if provider == "openai")); } diff --git a/litellm-rust/crates/core/src/messages/transformation.rs b/litellm-rust/crates/core/src/messages/transformation.rs index b478e20d24b..673a5728aca 100644 --- a/litellm-rust/crates/core/src/messages/transformation.rs +++ b/litellm-rust/crates/core/src/messages/transformation.rs @@ -1,6 +1,5 @@ -use crate::error::CoreResult; - use super::types::{AnthropicMessagesRequest, AnthropicMessagesResponse}; +use crate::Error; #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub enum MessagesAuthStrategy { @@ -23,13 +22,13 @@ pub trait AnthropicMessagesProviderConfig: Sync { api_base: Option<&str>, model: &str, env_lookup: &dyn Fn(&str) -> Option, - ) -> CoreResult; + ) -> Result; fn resolve_api_key( &self, api_key: Option<&str>, env_lookup: &dyn Fn(&str) -> Option, - ) -> CoreResult; + ) -> Result; fn auth_strategy(&self) -> MessagesAuthStrategy { MessagesAuthStrategy::Header("x-api-key") @@ -49,7 +48,7 @@ pub trait AnthropicMessagesProviderConfig: Sync { fn transform_request( &self, request: AnthropicMessagesRequest, - ) -> CoreResult { + ) -> Result { Ok(request) } @@ -57,7 +56,7 @@ pub trait AnthropicMessagesProviderConfig: Sync { &self, _model: &str, response: AnthropicMessagesResponse, - ) -> CoreResult { + ) -> Result { Ok(response) } } diff --git a/litellm-rust/crates/core/src/ocr/transformation.rs b/litellm-rust/crates/core/src/ocr/transformation.rs index cb3e735e533..3d3c16c8cb6 100644 --- a/litellm-rust/crates/core/src/ocr/transformation.rs +++ b/litellm-rust/crates/core/src/ocr/transformation.rs @@ -1,7 +1,6 @@ +use crate::Error; use serde_json::{Map, Value}; -use crate::CoreResult; - use super::types::{OcrRequestData, OcrResponseData}; #[derive(Clone, Copy, Debug, PartialEq, Eq)] @@ -43,13 +42,13 @@ pub trait OcrProviderConfig: Sync { model: &str, document: Value, optional_params: Map, - ) -> CoreResult; + ) -> Result; fn transform_ocr_response( &self, model: &str, response_json: Value, - ) -> CoreResult; + ) -> Result; fn complete_url( &self, @@ -57,13 +56,13 @@ pub trait OcrProviderConfig: Sync { model: &str, optional_params: &Map, env_lookup: &dyn Fn(&str) -> Option, - ) -> CoreResult; + ) -> Result; fn resolve_api_key( &self, api_key: Option<&str>, env_lookup: &dyn Fn(&str) -> Option, - ) -> CoreResult; + ) -> Result; fn auth_strategy(&self) -> OcrAuthStrategy { OcrAuthStrategy::Bearer diff --git a/litellm-rust/crates/core/src/providers/anthropic/chat_completions/tests.rs b/litellm-rust/crates/core/src/providers/anthropic/chat_completions/tests.rs index 4534ac0182c..b22de6c47de 100644 --- a/litellm-rust/crates/core/src/providers/anthropic/chat_completions/tests.rs +++ b/litellm-rust/crates/core/src/providers/anthropic/chat_completions/tests.rs @@ -1,4 +1,5 @@ use super::*; +use crate::Error; use serde_json::json; fn messages(value: Value) -> Vec { @@ -19,7 +20,7 @@ fn transform(model: &str, msgs: Value, opts: Value) -> Value { .body } -fn transform_response(body: Value) -> CoreResult { +fn transform_response(body: Value) -> Result { ANTHROPIC_CHAT_COMPLETIONS_CONFIG .transform_response("claude-sonnet-4-5", ProviderChatResponseData { body }) } @@ -390,29 +391,26 @@ fn declines_a_response_carrying_a_non_text_block() { "usage": {"input_tokens": 1, "output_tokens": 1} })) .expect_err("non-text block"); - assert_eq!( - err, - CoreError::Unsupported("non-text response content block") - ); + assert_eq!(err, Error::Unsupported("non-text response content block")); } #[test] fn errors_on_a_response_missing_required_fields() { assert_eq!( transform_response(json!("nope")).expect_err("not an object"), - CoreError::InvalidResponse("messages response is not an object".to_string()) + Error::InvalidResponse("messages response is not an object".to_string()) ); assert_eq!( transform_response(json!({"model": "m", "usage": {}})).expect_err("no content"), - CoreError::MissingField("content") + Error::MissingField("content") ); assert_eq!( transform_response(json!({"model": "m", "content": []})).expect_err("no usage"), - CoreError::MissingField("usage") + Error::MissingField("usage") ); assert_eq!( transform_response(json!({"content": [], "usage": {}})).expect_err("no model"), - CoreError::MissingField("model") + Error::MissingField("model") ); } diff --git a/litellm-rust/crates/core/src/providers/anthropic/chat_completions/transformation.rs b/litellm-rust/crates/core/src/providers/anthropic/chat_completions/transformation.rs index 3658642b539..97cc48aa6f2 100644 --- a/litellm-rust/crates/core/src/providers/anthropic/chat_completions/transformation.rs +++ b/litellm-rust/crates/core/src/providers/anthropic/chat_completions/transformation.rs @@ -10,7 +10,7 @@ use crate::chat_completions::types::{ ProviderChatRequestData, ProviderChatResponseData, }; use crate::constants::ANTHROPIC_OAUTH_TOKEN_PREFIX; -use crate::error::{CoreError, CoreResult}; +use crate::error::Error; use crate::providers::anthropic::messages::transformation::{ complete_anthropic_url, resolve_anthropic_api_key, }; @@ -74,7 +74,7 @@ impl ChatCompletionsProviderConfig for AnthropicChatCompletionsConfig { _model: &str, _optional_params: &Map, env_lookup: &dyn Fn(&str) -> Option, - ) -> CoreResult { + ) -> Result { Ok(complete_anthropic_url(api_base, env_lookup)) } @@ -84,7 +84,7 @@ impl ChatCompletionsProviderConfig for AnthropicChatCompletionsConfig { _model: &str, _optional_params: &Map, env_lookup: &dyn Fn(&str) -> Option, - ) -> CoreResult { + ) -> Result { Ok(ChatCompletionsAuth::Header { name: "x-api-key", value: resolve_anthropic_api_key(api_key, env_lookup)?, @@ -137,7 +137,7 @@ impl ChatCompletionsProviderConfig for AnthropicChatCompletionsConfig { model: &str, messages: Vec, optional_params: Map, - ) -> CoreResult { + ) -> Result { Ok(ProviderChatRequestData { body: anthropic_body(model, &build_conversation(&messages), optional_params), }) @@ -147,15 +147,16 @@ impl ChatCompletionsProviderConfig for AnthropicChatCompletionsConfig { &self, _model: &str, response: ProviderChatResponseData, - ) -> CoreResult { - let body = response.body.as_object().ok_or_else(|| { - CoreError::InvalidResponse("messages response is not an object".into()) - })?; + ) -> Result { + let body = response + .body + .as_object() + .ok_or_else(|| Error::InvalidResponse("messages response is not an object".into()))?; let content = body .get("content") .and_then(Value::as_array) - .ok_or(CoreError::MissingField("content"))?; + .ok_or(Error::MissingField("content"))?; // The route declines tool and thinking requests, so a non-text block // means the response carries something this path never asked for. // Decline rather than silently dropping it; the host falls back. @@ -163,7 +164,7 @@ impl ChatCompletionsProviderConfig for AnthropicChatCompletionsConfig { .iter() .any(|block| block.get("type").and_then(Value::as_str) != Some("text")) { - return Err(CoreError::Unsupported("non-text response content block")); + return Err(Error::Unsupported("non-text response content block")); } let text: String = content .iter() @@ -173,7 +174,7 @@ impl ChatCompletionsProviderConfig for AnthropicChatCompletionsConfig { let usage = body .get("usage") .and_then(Value::as_object) - .ok_or(CoreError::MissingField("usage"))?; + .ok_or(Error::MissingField("usage"))?; let field = |name: &str| usage.get(name).and_then(Value::as_u64).unwrap_or(0); Ok(ChatCompletionsResponse { @@ -181,7 +182,7 @@ impl ChatCompletionsProviderConfig for AnthropicChatCompletionsConfig { model: body .get("model") .and_then(Value::as_str) - .ok_or(CoreError::MissingField("model"))? + .ok_or(Error::MissingField("model"))? .to_string(), choices: vec![ChatCompletionsChoice { index: 0, diff --git a/litellm-rust/crates/core/src/providers/anthropic/messages/transformation.rs b/litellm-rust/crates/core/src/providers/anthropic/messages/transformation.rs index 829f2260d3c..8fcc0f36c7d 100644 --- a/litellm-rust/crates/core/src/providers/anthropic/messages/transformation.rs +++ b/litellm-rust/crates/core/src/providers/anthropic/messages/transformation.rs @@ -1,4 +1,4 @@ -use crate::error::{CoreError, CoreResult}; +use crate::error::Error; use crate::messages::transformation::{AnthropicMessagesProviderConfig, MessagesAuthStrategy}; const ANTHROPIC_API_KEY_ENV: &str = "ANTHROPIC_API_KEY"; @@ -17,12 +17,12 @@ pub fn non_empty(value: Option<&str>) -> Option<&str> { pub fn resolve_anthropic_api_key( api_key: Option<&str>, env_lookup: &dyn Fn(&str) -> Option, -) -> CoreResult { +) -> Result { non_empty(api_key) .map(str::to_string) .or_else(|| env_lookup(ANTHROPIC_API_KEY_ENV).filter(|value| !value.trim().is_empty())) .ok_or_else(|| { - CoreError::Auth( + Error::Auth( "Missing Anthropic API Key - Set `api_key` or the ANTHROPIC_API_KEY \ environment variable" .to_string(), @@ -52,7 +52,7 @@ impl AnthropicMessagesProviderConfig for AnthropicMessagesConfig { api_base: Option<&str>, _model: &str, env_lookup: &dyn Fn(&str) -> Option, - ) -> CoreResult { + ) -> Result { Ok(complete_anthropic_url(api_base, env_lookup)) } @@ -60,7 +60,7 @@ impl AnthropicMessagesProviderConfig for AnthropicMessagesConfig { &self, api_key: Option<&str>, env_lookup: &dyn Fn(&str) -> Option, - ) -> CoreResult { + ) -> Result { resolve_anthropic_api_key(api_key, env_lookup) } @@ -121,7 +121,7 @@ mod tests { ); assert!(matches!( resolve_anthropic_api_key(None, &|_| None).expect_err("missing key"), - CoreError::Auth(_) + Error::Auth(_) )); } diff --git a/litellm-rust/crates/core/src/providers/azure_ai/messages/transformation.rs b/litellm-rust/crates/core/src/providers/azure_ai/messages/transformation.rs index 7b958c77ba3..70dad0300f1 100644 --- a/litellm-rust/crates/core/src/providers/azure_ai/messages/transformation.rs +++ b/litellm-rust/crates/core/src/providers/azure_ai/messages/transformation.rs @@ -1,4 +1,4 @@ -use crate::error::{CoreError, CoreResult}; +use crate::error::Error; use crate::messages::transformation::{AnthropicMessagesProviderConfig, MessagesAuthStrategy}; use crate::messages::types::{ AnthropicMessage, AnthropicMessagesRequest, AnthropicMessagesResponse, ContentBlock, @@ -28,12 +28,12 @@ pub const AZURE_ANTHROPIC_MESSAGES_CONFIG: AzureAnthropicMessagesConfig = pub fn resolve_azure_api_key( api_key: Option<&str>, env_lookup: &dyn Fn(&str) -> Option, -) -> CoreResult { +) -> Result { non_empty(api_key) .map(str::to_string) .or_else(|| env_lookup(AZURE_API_KEY_ENV).filter(|value| !value.trim().is_empty())) .ok_or_else(|| { - CoreError::Auth( + Error::Auth( "Missing Azure API Key - Set `api_key` or the AZURE_API_KEY environment variable" .to_string(), ) @@ -43,12 +43,12 @@ pub fn resolve_azure_api_key( pub fn complete_azure_anthropic_url( api_base: Option<&str>, env_lookup: &dyn Fn(&str) -> Option, -) -> CoreResult { +) -> Result { let api_base = non_empty(api_base) .map(str::to_string) .or_else(|| env_lookup(AZURE_API_BASE_ENV).filter(|value| !value.trim().is_empty())) .ok_or_else(|| { - CoreError::Auth( + Error::Auth( "Missing Azure API Base - Set `api_base` or the AZURE_API_BASE environment variable. \ Expected format: https://.services.ai.azure.com/anthropic" .to_string(), @@ -147,7 +147,7 @@ impl AnthropicMessagesProviderConfig for AzureAnthropicMessagesConfig { api_base: Option<&str>, _model: &str, env_lookup: &dyn Fn(&str) -> Option, - ) -> CoreResult { + ) -> Result { complete_azure_anthropic_url(api_base, env_lookup) } @@ -155,7 +155,7 @@ impl AnthropicMessagesProviderConfig for AzureAnthropicMessagesConfig { &self, api_key: Option<&str>, env_lookup: &dyn Fn(&str) -> Option, - ) -> CoreResult { + ) -> Result { resolve_azure_api_key(api_key, env_lookup) } @@ -174,7 +174,7 @@ impl AnthropicMessagesProviderConfig for AzureAnthropicMessagesConfig { fn transform_request( &self, request: AnthropicMessagesRequest, - ) -> CoreResult { + ) -> Result { let mut request = fold_system_role_messages(request); if let Some(system) = request.system.as_mut() { strip_scope_from_system(system); @@ -190,7 +190,7 @@ impl AnthropicMessagesProviderConfig for AzureAnthropicMessagesConfig { &self, model: &str, response: AnthropicMessagesResponse, - ) -> CoreResult { + ) -> Result { self.anthropic.transform_response(model, response) } } @@ -268,7 +268,7 @@ mod tests { "https://env.services.ai.azure.com/anthropic/v1/messages" ); let err = complete_azure_anthropic_url(Some(" "), &|_| None).expect_err("missing base"); - assert!(matches!(err, CoreError::Auth(_))); + assert!(matches!(err, Error::Auth(_))); } #[test] @@ -284,7 +284,7 @@ mod tests { ); assert!(matches!( resolve_azure_api_key(None, &|_| None).expect_err("missing key"), - CoreError::Auth(_) + Error::Auth(_) )); } diff --git a/litellm-rust/crates/core/src/providers/azure_ai/ocr/transformation.rs b/litellm-rust/crates/core/src/providers/azure_ai/ocr/transformation.rs index eabd15677cc..b26a7925e8a 100644 --- a/litellm-rust/crates/core/src/providers/azure_ai/ocr/transformation.rs +++ b/litellm-rust/crates/core/src/providers/azure_ai/ocr/transformation.rs @@ -1,6 +1,6 @@ use std::collections::BTreeSet; -use crate::error::{CoreError, CoreResult, json_type_name}; +use crate::error::{Error, json_type_name}; use crate::ocr::transformation::{OcrAuthStrategy, OcrProviderConfig, OcrResponseHandling}; use crate::ocr::types::{OcrRequestData, OcrResponseData}; use serde_json::{Map, Value, json}; @@ -32,17 +32,17 @@ fn resolve_value( env_name: &str, env_lookup: &dyn Fn(&str) -> Option, missing_message: &str, -) -> CoreResult { +) -> Result { non_empty(explicit) .map(str::to_string) .or_else(|| env_lookup(env_name).filter(|value| !value.trim().is_empty())) - .ok_or_else(|| CoreError::Auth(missing_message.to_string())) + .ok_or_else(|| Error::Auth(missing_message.to_string())) } pub fn resolve_azure_ai_api_key( api_key: Option<&str>, env_lookup: &dyn Fn(&str) -> Option, -) -> CoreResult { +) -> Result { resolve_value( api_key, AZURE_AI_API_KEY_ENV, @@ -54,7 +54,7 @@ pub fn resolve_azure_ai_api_key( pub fn resolve_azure_ai_api_base( api_base: Option<&str>, env_lookup: &dyn Fn(&str) -> Option, -) -> CoreResult { +) -> Result { resolve_value( api_base, AZURE_AI_API_BASE_ENV, @@ -66,7 +66,7 @@ pub fn resolve_azure_ai_api_base( pub fn complete_azure_ai_url( api_base: Option<&str>, env_lookup: &dyn Fn(&str) -> Option, -) -> CoreResult { +) -> Result { let base = resolve_azure_ai_api_base(api_base, env_lookup)?; Ok(format!( "{}/providers/mistral/azure/ocr", @@ -77,7 +77,7 @@ pub fn complete_azure_ai_url( pub fn resolve_document_intelligence_api_key( api_key: Option<&str>, env_lookup: &dyn Fn(&str) -> Option, -) -> CoreResult { +) -> Result { resolve_value( api_key, AZURE_DOCUMENT_INTELLIGENCE_API_KEY_ENV, @@ -89,7 +89,7 @@ pub fn resolve_document_intelligence_api_key( pub fn resolve_document_intelligence_endpoint( api_base: Option<&str>, env_lookup: &dyn Fn(&str) -> Option, -) -> CoreResult { +) -> Result { resolve_value( api_base, AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT_ENV, @@ -127,7 +127,7 @@ fn pages_token_is_valid(token: &str) -> bool { } } -fn normalize_pages_param(pages: &Value) -> CoreResult> { +fn normalize_pages_param(pages: &Value) -> Result, Error> { match pages { Value::String(value) => { let normalized = value @@ -138,7 +138,7 @@ fn normalize_pages_param(pages: &Value) -> CoreResult> { if normalized.split(',').all(pages_token_is_valid) { Ok(Some(normalized)) } else { - Err(CoreError::InvalidRequest(format!( + Err(Error::InvalidRequest(format!( "Invalid `pages` string for Azure Document Intelligence: {value:?}. Expected format like '1-3,5,7-9'." ))) } @@ -152,7 +152,7 @@ fn normalize_pages_param(pages: &Value) -> CoreResult> { for value in values { let page = value.as_i64().expect("checked is_i64"); if page < 0 { - return Err(CoreError::InvalidRequest( + return Err(Error::InvalidRequest( "`pages` integers must be >= 0 (Mistral 0-based indices)".to_string(), )); } @@ -176,16 +176,16 @@ fn normalize_pages_param(pages: &Value) -> CoreResult> { if normalized.split(',').all(pages_token_is_valid) { return Ok(Some(normalized)); } - return Err(CoreError::InvalidRequest(format!( + return Err(Error::InvalidRequest(format!( "Invalid `pages` list for Azure Document Intelligence: {values:?}. Expected tokens like '1' or '3-5'." ))); } - Err(CoreError::InvalidRequest( + Err(Error::InvalidRequest( "`pages` must be a list[int] (0-based, Mistral-style) or a string like '1-3,5,7-9'." .to_string(), )) } - _ => Err(CoreError::InvalidRequest( + _ => Err(Error::InvalidRequest( "`pages` must be a list[int] (0-based, Mistral-style) or a string like '1-3,5,7-9'." .to_string(), )), @@ -197,7 +197,7 @@ pub fn complete_document_intelligence_url( model: &str, optional_params: &Map, env_lookup: &dyn Fn(&str) -> Option, -) -> CoreResult { +) -> Result { let endpoint = resolve_document_intelligence_endpoint(api_base, env_lookup)?; let mut url = format!( "{}/documentintelligence/documentModels/{}:analyze?api-version={}", @@ -216,20 +216,20 @@ pub fn complete_document_intelligence_url( Ok(url) } -fn document_url_from_mistral_document(document: &Value) -> CoreResult<&str> { - let object = document.as_object().ok_or_else(|| CoreError::InvalidType { +fn document_url_from_mistral_document(document: &Value) -> Result<&str, Error> { + let object = document.as_object().ok_or_else(|| Error::InvalidType { expected: "object", actual: json_type_name(document), })?; let doc_type = object .get("type") .and_then(Value::as_str) - .ok_or(CoreError::MissingField("document.type"))?; + .ok_or(Error::MissingField("document.type"))?; let field_name = match doc_type { "document_url" => "document_url", "image_url" => "image_url", other => { - return Err(CoreError::InvalidRequest(format!( + return Err(Error::InvalidRequest(format!( "Invalid document type: {other}. Must be 'document_url' or 'image_url'" ))); } @@ -238,7 +238,7 @@ fn document_url_from_mistral_document(document: &Value) -> CoreResult<&str> { .get(field_name) .and_then(Value::as_str) .filter(|value| !value.is_empty()) - .ok_or(CoreError::MissingField(field_name)) + .ok_or(Error::MissingField(field_name)) } fn extract_base64_from_data_uri(data_uri: &str) -> &str { @@ -290,7 +290,7 @@ impl OcrProviderConfig for AzureAiOcrConfig { model: &str, document: Value, optional_params: Map, - ) -> CoreResult { + ) -> Result { MISTRAL_OCR_CONFIG.transform_ocr_request(model, document, optional_params) } @@ -298,7 +298,7 @@ impl OcrProviderConfig for AzureAiOcrConfig { &self, model: &str, response_json: Value, - ) -> CoreResult { + ) -> Result { MISTRAL_OCR_CONFIG.transform_ocr_response(model, response_json) } @@ -308,7 +308,7 @@ impl OcrProviderConfig for AzureAiOcrConfig { _model: &str, _optional_params: &Map, env_lookup: &dyn Fn(&str) -> Option, - ) -> CoreResult { + ) -> Result { complete_azure_ai_url(api_base, env_lookup) } @@ -316,7 +316,7 @@ impl OcrProviderConfig for AzureAiOcrConfig { &self, api_key: Option<&str>, env_lookup: &dyn Fn(&str) -> Option, - ) -> CoreResult { + ) -> Result { resolve_azure_ai_api_key(api_key, env_lookup) } @@ -335,7 +335,7 @@ impl OcrProviderConfig for AzureDocumentIntelligenceOcrConfig { _model: &str, document: Value, _optional_params: Map, - ) -> CoreResult { + ) -> Result { let document_url = document_url_from_mistral_document(&document)?; let mut data = Map::new(); if document_url.starts_with("data:") { @@ -359,19 +359,19 @@ impl OcrProviderConfig for AzureDocumentIntelligenceOcrConfig { &self, model: &str, response_json: Value, - ) -> CoreResult { + ) -> Result { let response = response_json .as_object() - .ok_or_else(|| CoreError::InvalidType { + .ok_or_else(|| Error::InvalidType { expected: "object", actual: json_type_name(&response_json), })?; let status = response .get("status") .and_then(Value::as_str) - .ok_or(CoreError::MissingField("status"))?; + .ok_or(Error::MissingField("status"))?; if status != "succeeded" { - return Err(CoreError::InvalidResponse(format!( + return Err(Error::InvalidResponse(format!( "Azure Document Intelligence analysis failed with status: {status}" ))); } @@ -414,7 +414,7 @@ impl OcrProviderConfig for AzureDocumentIntelligenceOcrConfig { model: &str, optional_params: &Map, env_lookup: &dyn Fn(&str) -> Option, - ) -> CoreResult { + ) -> Result { complete_document_intelligence_url(api_base, model, optional_params, env_lookup) } @@ -422,7 +422,7 @@ impl OcrProviderConfig for AzureDocumentIntelligenceOcrConfig { &self, api_key: Option<&str>, env_lookup: &dyn Fn(&str) -> Option, - ) -> CoreResult { + ) -> Result { resolve_document_intelligence_api_key(api_key, env_lookup) } diff --git a/litellm-rust/crates/core/src/providers/bedrock/audio_transcription.rs b/litellm-rust/crates/core/src/providers/bedrock/audio_transcription.rs index 5e885734182..bb4f6afe5f9 100644 --- a/litellm-rust/crates/core/src/providers/bedrock/audio_transcription.rs +++ b/litellm-rust/crates/core/src/providers/bedrock/audio_transcription.rs @@ -6,7 +6,7 @@ use crate::audio_transcription::transformation::{ use crate::audio_transcription::types::{ AudioTranscriptionRequestData, AudioTranscriptionResponseData, }; -use crate::error::{CoreError, CoreResult, json_type_name}; +use crate::error::{Error, json_type_name}; pub use super::aws_base::{aws_auth_config, bedrock_model_id_and_region, resolve_bedrock_region}; use super::constants::{BEDROCK_RUNTIME_ENDPOINT_TEMPLATE, BEDROCK_SERVICE}; @@ -18,8 +18,8 @@ pub static BEDROCK_AUDIO_TRANSCRIPTION_CONFIG: BedrockAudioTranscriptionConfig = pub struct BedrockAudioTranscriptionConfig; -fn audio_fields(audio: Value) -> CoreResult<(String, String)> { - let object = audio.as_object().ok_or_else(|| CoreError::InvalidType { +fn audio_fields(audio: Value) -> Result<(String, String), Error> { + let object = audio.as_object().ok_or_else(|| Error::InvalidType { expected: "object", actual: json_type_name(&audio), })?; @@ -27,13 +27,13 @@ fn audio_fields(audio: Value) -> CoreResult<(String, String)> { .get("data") .and_then(Value::as_str) .filter(|value| !value.is_empty()) - .ok_or(CoreError::MissingField("audio.data"))?; + .ok_or(Error::MissingField("audio.data"))?; let format = object .get("format") .and_then(Value::as_str) .filter(|value| matches!(*value, "wav" | "mp3" | "flac" | "ogg")) .ok_or_else(|| { - CoreError::InvalidRequest("audio.format must be wav, mp3, flac, or ogg".to_string()) + Error::InvalidRequest("audio.format must be wav, mp3, flac, or ogg".to_string()) })?; Ok((data.to_string(), format.to_string())) } @@ -55,7 +55,7 @@ impl AudioTranscriptionProviderConfig for BedrockAudioTranscriptionConfig { _model: &str, audio: Value, optional_params: Map, - ) -> CoreResult { + ) -> Result { let (data, format) = audio_fields(audio)?; let mut instruction = "Transcribe the audio. Respond with only the transcript.".to_string(); if let Some(language) = optional_string(&optional_params, "language") { @@ -87,14 +87,14 @@ impl AudioTranscriptionProviderConfig for BedrockAudioTranscriptionConfig { &self, _model: &str, response_json: Value, - ) -> CoreResult { + ) -> Result { let content = response_json .get("output") .and_then(|value| value.get("message")) .and_then(|value| value.get("content")) .and_then(Value::as_array) .ok_or_else(|| { - CoreError::InvalidResponse("Bedrock response has no output content".to_string()) + Error::InvalidResponse("Bedrock response has no output content".to_string()) })?; let mut text = String::new(); for block in content { @@ -111,7 +111,7 @@ impl AudioTranscriptionProviderConfig for BedrockAudioTranscriptionConfig { model: &str, optional_params: &Map, env_lookup: &dyn Fn(&str) -> Option, - ) -> CoreResult { + ) -> Result { let (model_id, model_region) = bedrock_model_id_and_region(model); let region = resolve_bedrock_region(model_region.as_deref(), optional_params, env_lookup); let endpoint = optional_params @@ -133,7 +133,7 @@ impl AudioTranscriptionProviderConfig for BedrockAudioTranscriptionConfig { model: &str, optional_params: &Map, env_lookup: &dyn Fn(&str) -> Option, - ) -> CoreResult { + ) -> Result { let (_, model_region) = bedrock_model_id_and_region(model); Ok(AudioTranscriptionAuth::AwsSigV4 { region: resolve_bedrock_region(model_region.as_deref(), optional_params, env_lookup), diff --git a/litellm-rust/crates/core/src/providers/bedrock/aws_base.rs b/litellm-rust/crates/core/src/providers/bedrock/aws_base.rs index b11639aa09b..e5e52bfce95 100644 --- a/litellm-rust/crates/core/src/providers/bedrock/aws_base.rs +++ b/litellm-rust/crates/core/src/providers/bedrock/aws_base.rs @@ -4,7 +4,7 @@ use std::time::Duration; use std::time::{SystemTime, UNIX_EPOCH}; use crate::caching::in_memory_cache::InMemoryCache; -use crate::error::{CoreError, CoreResult}; +use crate::error::Error; use aws_credential_types::Credentials; use aws_credential_types::provider::ProvideCredentials; use aws_sigv4::http_request::{ @@ -197,7 +197,7 @@ pub fn classify_auth( pub async fn resolve_credentials( config: AwsAuthConfig, env_lookup: &(dyn Fn(&str) -> Option + Sync), -) -> CoreResult { +) -> Result { let resolved = config.clone().with_environment(env_lookup); let flow = classify_auth(config, env_lookup); match flow { @@ -244,9 +244,10 @@ pub async fn resolve_credentials( let provider = aws_config::profile::ProfileFileCredentialsProvider::builder() .profile_name(name) .build(); - provider.provide_credentials().await.map_err(|error| { - CoreError::Auth(format!("AWS profile credentials failed: {error}")) - }) + provider + .provide_credentials() + .await + .map_err(|error| Error::Auth(format!("AWS profile credentials failed: {error}"))) } AwsAuthFlow::AssumeRole { role, session_name } => { if is_already_running_as_role(&role, &resolved).await? { @@ -260,7 +261,7 @@ pub async fn resolve_credentials( .build() .await; let credentials = provider.provide_credentials().await.map_err(|error| { - CoreError::Auth(format!("AWS default credentials failed: {error}")) + Error::Auth(format!("AWS default credentials failed: {error}")) })?; set_cached_credentials( key, @@ -301,7 +302,7 @@ pub async fn resolve_credentials( provider .provide_credentials() .await - .map_err(|error| CoreError::Auth(format!("AWS role credentials failed: {error}"))) + .map_err(|error| Error::Auth(format!("AWS role credentials failed: {error}"))) } AwsAuthFlow::WebIdentity { token, @@ -325,13 +326,13 @@ pub async fn resolve_credentials( .send() .await .map_err(|error| { - CoreError::Auth(format!("AWS web identity credentials failed: {error}")) + Error::Auth(format!("AWS web identity credentials failed: {error}")) })?; let credentials = response.credentials().ok_or_else(|| { - CoreError::Auth("AWS web identity response had no credentials".to_string()) + Error::Auth("AWS web identity response had no credentials".to_string()) })?; let expiration = SystemTime::try_from(*credentials.expiration()).map_err(|error| { - CoreError::Auth(format!("AWS web identity expiration was invalid: {error}")) + Error::Auth(format!("AWS web identity expiration was invalid: {error}")) })?; Ok(Credentials::new( credentials.access_key_id(), @@ -350,9 +351,10 @@ pub async fn resolve_credentials( aws_config::default_provider::credentials::DefaultCredentialsChain::builder() .build() .await; - let credentials = provider.provide_credentials().await.map_err(|error| { - CoreError::Auth(format!("AWS default credentials failed: {error}")) - })?; + let credentials = provider + .provide_credentials() + .await + .map_err(|error| Error::Auth(format!("AWS default credentials failed: {error}")))?; set_cached_credentials( key, credentials.clone(), @@ -363,7 +365,7 @@ pub async fn resolve_credentials( } } -async fn is_already_running_as_role(role: &str, config: &AwsAuthConfig) -> CoreResult { +async fn is_already_running_as_role(role: &str, config: &AwsAuthConfig) -> Result { if role_identity(role).is_none() { return Ok(false); } @@ -437,7 +439,7 @@ pub fn sign_bedrock_post( region: &str, credentials: &Credentials, signing_time: SystemTime, -) -> CoreResult> { +) -> Result, Error> { let identity: Identity = credentials.clone().into(); let params = v4::SigningParams::builder() .identity(&identity) @@ -447,14 +449,14 @@ pub fn sign_bedrock_post( .settings(SigningSettings::default()) .build() .map(SigningParams::from) - .map_err(|error| CoreError::Auth(format!("AWS signing parameters failed: {error}")))?; + .map_err(|error| Error::Auth(format!("AWS signing parameters failed: {error}")))?; let header_refs = headers .iter() .map(|(name, value)| (name.as_str(), value.as_str())); let request = SignableRequest::new("POST", url, header_refs, SignableBody::Bytes(body)) - .map_err(|error| CoreError::Auth(format!("AWS signable request failed: {error}")))?; + .map_err(|error| Error::Auth(format!("AWS signable request failed: {error}")))?; let (instructions, _) = sign(request, ¶ms) - .map_err(|error| CoreError::Auth(format!("AWS request signing failed: {error}")))? + .map_err(|error| Error::Auth(format!("AWS request signing failed: {error}")))? .into_parts(); Ok(instructions .headers() diff --git a/litellm-rust/crates/core/src/providers/bedrock/chat_completions/tests.rs b/litellm-rust/crates/core/src/providers/bedrock/chat_completions/tests.rs index 4b75dcb8e9d..c86f061b9ca 100644 --- a/litellm-rust/crates/core/src/providers/bedrock/chat_completions/tests.rs +++ b/litellm-rust/crates/core/src/providers/bedrock/chat_completions/tests.rs @@ -1,4 +1,5 @@ use super::*; +use crate::Error; use serde_json::json; fn messages(value: Value) -> Vec { @@ -23,7 +24,7 @@ fn transform(msgs: Value, opts: Value) -> Value { .body } -fn transform_response(body: Value) -> CoreResult { +fn transform_response(body: Value) -> Result { BEDROCK_CHAT_COMPLETIONS_CONFIG.transform_response( "anthropic.claude-sonnet-4-5-v1:0", ProviderChatResponseData { body }, @@ -478,25 +479,22 @@ fn declines_a_response_carrying_a_tool_use_block() { "usage": {"inputTokens": 1, "outputTokens": 1} })) .expect_err("tool use block"); - assert_eq!( - err, - CoreError::Unsupported("non-text response content block") - ); + assert_eq!(err, Error::Unsupported("non-text response content block")); } #[test] fn errors_on_a_response_missing_required_fields() { assert_eq!( transform_response(json!("nope")).expect_err("not an object"), - CoreError::InvalidResponse("converse response is not an object".to_string()) + Error::InvalidResponse("converse response is not an object".to_string()) ); assert_eq!( transform_response(json!({"usage": {}})).expect_err("no output"), - CoreError::MissingField("output.message.content") + Error::MissingField("output.message.content") ); assert_eq!( transform_response(json!({"output": {"message": {"content": []}}})).expect_err("no usage"), - CoreError::MissingField("usage") + Error::MissingField("usage") ); } diff --git a/litellm-rust/crates/core/src/providers/bedrock/chat_completions/transformation.rs b/litellm-rust/crates/core/src/providers/bedrock/chat_completions/transformation.rs index b107950748e..ef5f44b4a14 100644 --- a/litellm-rust/crates/core/src/providers/bedrock/chat_completions/transformation.rs +++ b/litellm-rust/crates/core/src/providers/bedrock/chat_completions/transformation.rs @@ -11,7 +11,7 @@ use crate::chat_completions::types::{ ChatCompletionsUsage, ChatMessage, ChatMessageContent, ProviderChatRequestData, ProviderChatResponseData, }; -use crate::error::{CoreError, CoreResult}; +use crate::error::Error; use super::super::aws_base::{bedrock_model_id_and_region, resolve_bedrock_region}; use super::super::constants::{AWS_BEARER_TOKEN_BEDROCK, BEDROCK_RUNTIME_ENDPOINT_TEMPLATE}; @@ -110,7 +110,7 @@ impl ChatCompletionsProviderConfig for BedrockChatCompletionsConfig { model: &str, optional_params: &Map, env_lookup: &dyn Fn(&str) -> Option, - ) -> CoreResult { + ) -> Result { let (model_id, model_region) = bedrock_model_id_and_region(model); let region = resolve_bedrock_region(model_region.as_deref(), optional_params, env_lookup); let endpoint = optional_params @@ -137,7 +137,7 @@ impl ChatCompletionsProviderConfig for BedrockChatCompletionsConfig { model: &str, optional_params: &Map, env_lookup: &dyn Fn(&str) -> Option, - ) -> CoreResult { + ) -> Result { // Python reads `api_key` as the Bedrock bearer token and consults the // env only when the caller passed none, so a caller-supplied empty key // falls through to SigV4 without reaching for the environment. An @@ -208,7 +208,7 @@ impl ChatCompletionsProviderConfig for BedrockChatCompletionsConfig { _model: &str, messages: Vec, optional_params: Map, - ) -> CoreResult { + ) -> Result { Ok(ProviderChatRequestData { body: converse_body(&build_conversation(&messages), &optional_params), }) @@ -218,17 +218,18 @@ impl ChatCompletionsProviderConfig for BedrockChatCompletionsConfig { &self, model: &str, response: ProviderChatResponseData, - ) -> CoreResult { - let body = response.body.as_object().ok_or_else(|| { - CoreError::InvalidResponse("converse response is not an object".into()) - })?; + ) -> Result { + let body = response + .body + .as_object() + .ok_or_else(|| Error::InvalidResponse("converse response is not an object".into()))?; let content = body .get("output") .and_then(|output| output.get("message")) .and_then(|message| message.get("content")) .and_then(Value::as_array) - .ok_or(CoreError::MissingField("output.message.content"))?; + .ok_or(Error::MissingField("output.message.content"))?; // The route declines tool requests, so anything other than a text block // is something this path never asked for. Decline; the host falls back. if content.iter().any(|block| { @@ -236,7 +237,7 @@ impl ChatCompletionsProviderConfig for BedrockChatCompletionsConfig { .as_object() .is_none_or(|block| block.len() != 1 || !block.contains_key("text")) }) { - return Err(CoreError::Unsupported("non-text response content block")); + return Err(Error::Unsupported("non-text response content block")); } let text: String = content .iter() @@ -246,7 +247,7 @@ impl ChatCompletionsProviderConfig for BedrockChatCompletionsConfig { let usage = body .get("usage") .and_then(Value::as_object) - .ok_or(CoreError::MissingField("usage"))?; + .ok_or(Error::MissingField("usage"))?; let field = |name: &str| usage.get(name).and_then(Value::as_u64).unwrap_or(0); let computed = usage_from_parts( field("inputTokens"), diff --git a/litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs b/litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs index dc720cc4244..6a8a38204a9 100644 --- a/litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs +++ b/litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs @@ -1,4 +1,4 @@ -use crate::error::{CoreError, CoreResult, json_type_name}; +use crate::error::{Error, json_type_name}; use crate::ocr::transformation::OcrProviderConfig; use crate::ocr::types::{OcrRequestData, OcrResponseData}; use serde_json::{Map, Value}; @@ -47,7 +47,7 @@ pub fn complete_url(api_base: Option<&str>) -> String { /// Resolve the Mistral API key from the explicit param or the environment. /// -/// Blank/whitespace values are treated as absent. Returns `CoreError::Auth` +/// Blank/whitespace values are treated as absent. Returns `Error::Auth` /// when no usable key is available. /// /// Note: the env fallback only reads the process environment. Secret-manager @@ -56,13 +56,13 @@ pub fn complete_url(api_base: Option<&str>) -> String { pub fn resolve_api_key( api_key: Option<&str>, env_lookup: &dyn Fn(&str) -> Option, -) -> CoreResult { +) -> Result { api_key .map(str::trim) .filter(|key| !key.is_empty()) .map(str::to_string) .or_else(|| env_lookup(MISTRAL_API_KEY_ENV).filter(|key| !key.trim().is_empty())) - .ok_or_else(|| CoreError::Auth(MISSING_KEY_MESSAGE.to_string())) + .ok_or_else(|| Error::Auth(MISSING_KEY_MESSAGE.to_string())) } pub struct MistralOcrConfig; @@ -79,9 +79,9 @@ impl OcrProviderConfig for MistralOcrConfig { model: &str, document: Value, optional_params: Map, - ) -> CoreResult { + ) -> Result { if !document.is_object() { - return Err(CoreError::InvalidType { + return Err(Error::InvalidType { expected: "object", actual: json_type_name(&document), }); @@ -104,10 +104,10 @@ impl OcrProviderConfig for MistralOcrConfig { &self, model: &str, response_json: Value, - ) -> CoreResult { + ) -> Result { let response_object = response_json .as_object() - .ok_or_else(|| CoreError::InvalidType { + .ok_or_else(|| Error::InvalidType { expected: "object", actual: json_type_name(&response_json), })?; @@ -140,7 +140,7 @@ impl OcrProviderConfig for MistralOcrConfig { _model: &str, _optional_params: &Map, _env_lookup: &dyn Fn(&str) -> Option, - ) -> CoreResult { + ) -> Result { Ok(complete_url(api_base)) } @@ -148,7 +148,7 @@ impl OcrProviderConfig for MistralOcrConfig { &self, api_key: Option<&str>, env_lookup: &dyn Fn(&str) -> Option, - ) -> CoreResult { + ) -> Result { resolve_api_key(api_key, env_lookup) } } @@ -165,11 +165,11 @@ pub fn transform_ocr_request( model: &str, document: Value, optional_params: Map, -) -> CoreResult { +) -> Result { MISTRAL_OCR_CONFIG.transform_ocr_request(model, document, optional_params) } -pub fn transform_ocr_response(model: &str, response_json: Value) -> CoreResult { +pub fn transform_ocr_response(model: &str, response_json: Value) -> Result { MISTRAL_OCR_CONFIG.transform_ocr_response(model, response_json) } @@ -250,7 +250,7 @@ mod tests { assert_eq!( err, - CoreError::InvalidType { + Error::InvalidType { expected: "object", actual: "string", } @@ -307,6 +307,6 @@ mod tests { #[test] fn resolve_api_key_errors_when_absent() { let err = resolve_api_key(None, &|_| None).expect_err("missing key should error"); - assert_eq!(err, CoreError::Auth(MISSING_KEY_MESSAGE.to_string())); + assert_eq!(err, Error::Auth(MISSING_KEY_MESSAGE.to_string())); } } diff --git a/litellm-rust/crates/core/src/providers/openai/realtime/transformation.rs b/litellm-rust/crates/core/src/providers/openai/realtime/transformation.rs index b3f6b03b28a..f1985f81b7d 100644 --- a/litellm-rust/crates/core/src/providers/openai/realtime/transformation.rs +++ b/litellm-rust/crates/core/src/providers/openai/realtime/transformation.rs @@ -1,4 +1,4 @@ -use crate::CoreResult; +use crate::Error; use crate::realtime::transformation::RealtimeProviderConfig; use crate::realtime::types::{RealtimeEvent, RealtimeTransformResult}; @@ -72,7 +72,7 @@ impl RealtimeProviderConfig for OpenAiRealtimeConfig { &self, event: &RealtimeEvent, _model: &str, - ) -> CoreResult { + ) -> Result { Ok(RealtimeTransformResult::passthrough(event.clone())) } @@ -80,7 +80,7 @@ impl RealtimeProviderConfig for OpenAiRealtimeConfig { &self, event: &RealtimeEvent, _model: &str, - ) -> CoreResult { + ) -> Result { Ok(RealtimeTransformResult::passthrough(event.clone())) } } @@ -88,14 +88,14 @@ impl RealtimeProviderConfig for OpenAiRealtimeConfig { pub fn transform_realtime_request( event: &RealtimeEvent, model: &str, -) -> CoreResult { +) -> Result { OPENAI_REALTIME_CONFIG.transform_realtime_request(event, model) } pub fn transform_realtime_response( event: &RealtimeEvent, model: &str, -) -> CoreResult { +) -> Result { OPENAI_REALTIME_CONFIG.transform_realtime_response(event, model) } diff --git a/litellm-rust/crates/core/src/providers/openai/responses/transformation.rs b/litellm-rust/crates/core/src/providers/openai/responses/transformation.rs index e15197c468c..be86bb90311 100644 --- a/litellm-rust/crates/core/src/providers/openai/responses/transformation.rs +++ b/litellm-rust/crates/core/src/providers/openai/responses/transformation.rs @@ -1,4 +1,4 @@ -use crate::CoreResult; +use crate::Error; use crate::responses::types::{ResponsesWsEvent, ResponsesWsTransformResult}; use crate::responses::websocket::{ResponsesWebSocketProviderConfig, enforce_model}; @@ -15,7 +15,7 @@ impl ResponsesWebSocketProviderConfig for OpenAIResponsesWsConfig { &self, event: &ResponsesWsEvent, model: &str, - ) -> CoreResult { + ) -> Result { Ok(ResponsesWsTransformResult::passthrough(enforce_model( event, model, ))) @@ -25,7 +25,7 @@ impl ResponsesWebSocketProviderConfig for OpenAIResponsesWsConfig { &self, event: &ResponsesWsEvent, _model: &str, - ) -> CoreResult { + ) -> Result { Ok(ResponsesWsTransformResult::passthrough(event.clone())) } } diff --git a/litellm-rust/crates/core/src/providers/vertex_ai/ocr/transformation.rs b/litellm-rust/crates/core/src/providers/vertex_ai/ocr/transformation.rs index 6300149c237..ee095447028 100644 --- a/litellm-rust/crates/core/src/providers/vertex_ai/ocr/transformation.rs +++ b/litellm-rust/crates/core/src/providers/vertex_ai/ocr/transformation.rs @@ -1,4 +1,4 @@ -use crate::error::{CoreError, CoreResult, json_type_name}; +use crate::error::{Error, json_type_name}; use crate::ocr::transformation::OcrProviderConfig; use crate::ocr::types::{OcrRequestData, OcrResponseData}; use serde_json::{Map, Value, json}; @@ -43,7 +43,7 @@ pub fn is_deepseek_model(model: &str) -> bool { pub fn resolve_vertex_api_key( api_key: Option<&str>, env_lookup: &dyn Fn(&str) -> Option, -) -> CoreResult { +) -> Result { api_key .map(str::trim) .filter(|key| !key.is_empty()) @@ -51,7 +51,7 @@ pub fn resolve_vertex_api_key( .or_else(|| env_lookup(VERTEX_AI_API_KEY_ENV).filter(|key| !key.trim().is_empty())) .or_else(|| env_lookup(VERTEXAI_API_KEY_ENV).filter(|key| !key.trim().is_empty())) .ok_or_else(|| { - CoreError::Auth( + Error::Auth( "Missing Vertex AI access token - pass api_key or provide Authorization via extra_headers" .to_string(), ) @@ -61,12 +61,12 @@ pub fn resolve_vertex_api_key( fn vertex_project( params: &Map, env_lookup: &dyn Fn(&str) -> Option, -) -> CoreResult { +) -> Result { string_param(params, &["vertex_project", "vertex_ai_project"]) .map(str::to_string) .or_else(|| env_lookup(VERTEXAI_PROJECT_ENV).filter(|value| !value.trim().is_empty())) .ok_or_else(|| { - CoreError::InvalidRequest( + Error::InvalidRequest( "Missing vertex_project - Set VERTEXAI_PROJECT environment variable or pass vertex_project parameter" .to_string(), ) @@ -99,7 +99,7 @@ pub fn complete_vertex_mistral_url( model: &str, optional_params: &Map, env_lookup: &dyn Fn(&str) -> Option, -) -> CoreResult { +) -> Result { let project = vertex_project(optional_params, env_lookup)?; let location = vertex_location(optional_params, env_lookup); let base = vertex_mistral_api_base(api_base, &location); @@ -112,7 +112,7 @@ pub fn complete_vertex_deepseek_url( api_base: Option<&str>, optional_params: &Map, env_lookup: &dyn Fn(&str) -> Option, -) -> CoreResult { +) -> Result { let project = vertex_project(optional_params, env_lookup)?; let location = vertex_location(optional_params, env_lookup); let base = api_base @@ -125,20 +125,20 @@ pub fn complete_vertex_deepseek_url( )) } -fn document_content_item(document: &Value) -> CoreResult { - let object = document.as_object().ok_or_else(|| CoreError::InvalidType { +fn document_content_item(document: &Value) -> Result { + let object = document.as_object().ok_or_else(|| Error::InvalidType { expected: "object", actual: json_type_name(document), })?; let doc_type = object .get("type") .and_then(Value::as_str) - .ok_or(CoreError::MissingField("document.type"))?; + .ok_or(Error::MissingField("document.type"))?; let url_field = match doc_type { "image_url" => "image_url", "document_url" => "document_url", other => { - return Err(CoreError::InvalidRequest(format!( + return Err(Error::InvalidRequest(format!( "Unsupported document type: {other}. Expected 'image_url' or 'document_url'" ))); } @@ -147,7 +147,7 @@ fn document_content_item(document: &Value) -> CoreResult { .get(url_field) .and_then(Value::as_str) .filter(|value| !value.is_empty()) - .ok_or(CoreError::MissingField(url_field))?; + .ok_or(Error::MissingField(url_field))?; Ok(json!({ "type": "image_url", @@ -163,7 +163,7 @@ fn deepseek_model_name(model: &str) -> String { } } -fn first_choice_content(response: &Value) -> CoreResult { +fn first_choice_content(response: &Value) -> Result { response .get("choices") .and_then(Value::as_array) @@ -176,9 +176,7 @@ fn first_choice_content(response: &Value) -> CoreResult { Value::Object(_) => true, _ => false, }) - .ok_or_else(|| { - CoreError::InvalidResponse("No content in DeepSeek OCR response".to_string()) - }) + .ok_or_else(|| Error::InvalidResponse("No content in DeepSeek OCR response".to_string())) } fn ocr_data_from_content(content: Value, usage: Option, model: &str) -> Value { @@ -219,7 +217,7 @@ impl OcrProviderConfig for VertexAiOcrConfig { model: &str, document: Value, optional_params: Map, - ) -> CoreResult { + ) -> Result { MISTRAL_OCR_CONFIG.transform_ocr_request(model, document, optional_params) } @@ -227,7 +225,7 @@ impl OcrProviderConfig for VertexAiOcrConfig { &self, model: &str, response_json: Value, - ) -> CoreResult { + ) -> Result { MISTRAL_OCR_CONFIG.transform_ocr_response(model, response_json) } @@ -237,7 +235,7 @@ impl OcrProviderConfig for VertexAiOcrConfig { model: &str, optional_params: &Map, env_lookup: &dyn Fn(&str) -> Option, - ) -> CoreResult { + ) -> Result { complete_vertex_mistral_url(api_base, model, optional_params, env_lookup) } @@ -245,7 +243,7 @@ impl OcrProviderConfig for VertexAiOcrConfig { &self, api_key: Option<&str>, env_lookup: &dyn Fn(&str) -> Option, - ) -> CoreResult { + ) -> Result { resolve_vertex_api_key(api_key, env_lookup) } @@ -264,7 +262,7 @@ impl OcrProviderConfig for VertexAiDeepSeekOcrConfig { model: &str, document: Value, optional_params: Map, - ) -> CoreResult { + ) -> Result { let mut data = Map::new(); data.insert( "model".to_string(), @@ -289,10 +287,10 @@ impl OcrProviderConfig for VertexAiDeepSeekOcrConfig { &self, model: &str, response_json: Value, - ) -> CoreResult { + ) -> Result { let response = response_json .as_object() - .ok_or_else(|| CoreError::InvalidType { + .ok_or_else(|| Error::InvalidType { expected: "object", actual: json_type_name(&response_json), })?; @@ -314,7 +312,7 @@ impl OcrProviderConfig for VertexAiDeepSeekOcrConfig { }); } - let object = ocr_data.as_object().ok_or_else(|| CoreError::InvalidType { + let object = ocr_data.as_object().ok_or_else(|| Error::InvalidType { expected: "object", actual: json_type_name(&ocr_data), })?; @@ -346,7 +344,7 @@ impl OcrProviderConfig for VertexAiDeepSeekOcrConfig { _model: &str, optional_params: &Map, env_lookup: &dyn Fn(&str) -> Option, - ) -> CoreResult { + ) -> Result { complete_vertex_deepseek_url(api_base, optional_params, env_lookup) } @@ -354,7 +352,7 @@ impl OcrProviderConfig for VertexAiDeepSeekOcrConfig { &self, api_key: Option<&str>, env_lookup: &dyn Fn(&str) -> Option, - ) -> CoreResult { + ) -> Result { resolve_vertex_api_key(api_key, env_lookup) } } diff --git a/litellm-rust/crates/core/src/realtime/transformation.rs b/litellm-rust/crates/core/src/realtime/transformation.rs index 69b88687000..b08084514ef 100644 --- a/litellm-rust/crates/core/src/realtime/transformation.rs +++ b/litellm-rust/crates/core/src/realtime/transformation.rs @@ -1,4 +1,4 @@ -use crate::CoreResult; +use crate::Error; use crate::realtime::types::{RealtimeEvent, RealtimeTransformResult}; pub trait RealtimeProviderConfig { @@ -11,12 +11,12 @@ pub trait RealtimeProviderConfig { &self, event: &RealtimeEvent, model: &str, - ) -> CoreResult; + ) -> Result; /// Transform a backend → client event before it is forwarded downstream. fn transform_realtime_response( &self, event: &RealtimeEvent, model: &str, - ) -> CoreResult; + ) -> Result; } diff --git a/litellm-rust/crates/core/src/responses/instrumentation.rs b/litellm-rust/crates/core/src/responses/instrumentation.rs index ec04571da14..b1098f4d386 100644 --- a/litellm-rust/crates/core/src/responses/instrumentation.rs +++ b/litellm-rust/crates/core/src/responses/instrumentation.rs @@ -5,9 +5,9 @@ use std::time::{SystemTime, UNIX_EPOCH}; use serde_json::Value; +use crate::Error; use crate::call_lifecycle::{CallLifecycleContext, CallLifecycleHooks, CallLifecycleTiming}; use crate::responses::types::{ResponsesWsEvent, ResponsesWsEventType}; -use crate::{CoreError, CoreResult}; #[derive(Clone, Debug, Default, PartialEq, Eq)] pub struct ResponsesWsUsage { @@ -205,7 +205,7 @@ impl ResponsesWsInstrumentation { } } -type LifecycleFuture<'a, T> = Pin> + Send + 'a>>; +type LifecycleFuture<'a, T> = Pin> + Send + 'a>>; impl CallLifecycleHooks<(), (), ()> for ResponsesWsInstrumentation { type PreCallFuture<'a> = LifecycleFuture<'a, ()>; @@ -246,7 +246,7 @@ impl CallLifecycleHooks<(), (), ()> for ResponsesWsInstrumentation { fn async_log_failure_event<'a>( &'a self, _context: &'a CallLifecycleContext, - _error: &'a CoreError, + _error: &'a Error, _timing: &'a CallLifecycleTiming, ) -> Self::FailureFuture<'a> { Box::pin(async move { @@ -342,7 +342,7 @@ mod tests { ), (), &instrumentation, - |_| async { Ok::<(), CoreError>(()) }, + |_| async { Ok::<(), Error>(()) }, ) .await; diff --git a/litellm-rust/crates/core/src/responses/websocket.rs b/litellm-rust/crates/core/src/responses/websocket.rs index 92dc19627a0..5d037e9cf1b 100644 --- a/litellm-rust/crates/core/src/responses/websocket.rs +++ b/litellm-rust/crates/core/src/responses/websocket.rs @@ -1,4 +1,4 @@ -use crate::CoreResult; +use crate::Error; use crate::constants::{OPENAI_RESPONSES_DEFAULT_API_BASE, OPENAI_RESPONSES_PATH}; use crate::responses::types::{ResponsesWsEvent, ResponsesWsEventType, ResponsesWsTransformResult}; @@ -19,13 +19,13 @@ pub trait ResponsesWebSocketProviderConfig: Sync { &self, event: &ResponsesWsEvent, model: &str, - ) -> CoreResult; + ) -> Result; fn transform_ws_response( &self, event: &ResponsesWsEvent, model: &str, - ) -> CoreResult; + ) -> Result; } pub fn complete_websocket_url( diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index 746e0770f9b..68aa9436b15 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -10,7 +10,7 @@ use litellm_core::chat_completions::types::{ChatCompletionsRequest, ChatCompleti use litellm_core::chat_completions::{ chat_completions as run_chat_completions, chat_completions_decline_reason, }; -use litellm_core::error::CoreError; +use litellm_core::error::Error; use litellm_core::messages::messages as run_messages; use litellm_core::messages::types::{AnthropicMessagesResponse, MessagesRequest}; use litellm_python_interop::{from_py, release_count, release_gil, to_py}; @@ -54,13 +54,13 @@ fn chat_completions_response_to_py( to_py(py, &response) } -fn core_error_to_pyerr(err: CoreError) -> PyErr { +fn core_error_to_pyerr(err: Error) -> PyErr { match err { - CoreError::Auth(message) => PyValueError::new_err(message), - CoreError::InvalidProvider(_) - | CoreError::InvalidRequest(_) - | CoreError::InvalidType { .. } - | CoreError::MissingField(_) => PyValueError::new_err(err.to_string()), + Error::Auth(message) => PyValueError::new_err(message), + Error::InvalidProvider(_) + | Error::InvalidRequest(_) + | Error::InvalidType { .. } + | Error::MissingField(_) => PyValueError::new_err(err.to_string()), other => PyRuntimeError::new_err(other.to_string()), } } @@ -71,22 +71,22 @@ fn core_error_to_pyerr(err: CoreError) -> PyErr { /// Everything raised before the request goes out is safe for the host to retry /// on its own path; anything after it is not, because the provider has already /// done the work and billed for it. -fn chat_completions_error_to_pyerr(err: CoreError) -> PyErr { +fn chat_completions_error_to_pyerr(err: Error) -> PyErr { match err { - CoreError::Unsupported(_) - | CoreError::Auth(_) - | CoreError::InvalidProvider(_) - | CoreError::InvalidRequest(_) - | CoreError::InvalidType { .. } - | CoreError::MissingField(_) - | CoreError::Routing(_) + Error::Unsupported(_) + | Error::Auth(_) + | Error::InvalidProvider(_) + | Error::InvalidRequest(_) + | Error::InvalidType { .. } + | Error::MissingField(_) + | Error::Routing(_) // Nothing reached the provider, so serving it on Python cannot double // bill and is the only way the caller gets an answer at all. - | CoreError::Connect(_) => RustBridgeDeclined::new_err(err.to_string()), - CoreError::Http { status, body } => { + | Error::Connect(_) => RustBridgeDeclined::new_err(err.to_string()), + Error::Http { status, body } => { RustUpstreamError::new_err((status, format!("{status}: {body}"))) } - CoreError::Network(message) | CoreError::InvalidResponse(message) => { + Error::Network(message) | Error::InvalidResponse(message) => { RustUpstreamError::new_err((0u16, message)) } } 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 93/97] 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 94/97] 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 95/97] 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 From 711430216ea73eb0ad45773a81c61155062e0567 Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Wed, 2 Sep 2026 13:33:14 -0700 Subject: [PATCH 96/97] fix(ui): preserve full AgentCore runtime ARN in agent edit form (#39382) parseDynamicAgentForForm recovered a credential field's value from a stored model string by splitting both the model_template and the model on "/" and matching by array index. That breaks for any placeholder value that itself contains "/", such as a Bedrock AgentCore runtime ARN resource path (runtime/), silently dropping everything after the first slash when populating the edit form. Saving without touching the field then persisted the truncated ARN. Replace the index-matching split with a non-mutating template parse (split on the placeholder pattern, escape and rejoin the literal segments into a regex) so a placeholder captures everything it needs regardless of embedded slashes. Also add a lightweight ARN-shape validator for the AgentCore runtime ARN field, guarded against a malformed pattern string, so a truncated value is rejected client-side before it reaches the backend. Resolves LIT-6737 --- .../public_endpoints/agent_create_fields.json | 4 +- .../public_endpoints/public_endpoints.py | 2 + .../public_endpoints/test_public_endpoints.py | 40 ++++++ .../agent_info.integration.test.tsx | 97 ++++++++++++++ .../_components/agent_type_utils.test.ts | 73 ++++++++++ .../agents/_components/agent_type_utils.ts | 45 +++++-- .../_components/dynamic_agent_form_fields.tsx | 126 ++++++++++-------- .../src/components/networking.tsx | 2 + ui/litellm-dashboard/src/lib/http/schema.d.ts | 4 + 9 files changed, 326 insertions(+), 67 deletions(-) create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_type_utils.test.ts diff --git a/litellm/proxy/public_endpoints/agent_create_fields.json b/litellm/proxy/public_endpoints/agent_create_fields.json index 36484cc1065..cc2fc17d759 100644 --- a/litellm/proxy/public_endpoints/agent_create_fields.json +++ b/litellm/proxy/public_endpoints/agent_create_fields.json @@ -107,7 +107,9 @@ "required": true, "field_type": "text", "default_value": null, - "include_in_litellm_params": false + "include_in_litellm_params": false, + "validation_pattern": "^arn:aws[a-zA-Z0-9-]*:bedrock-agentcore:[a-z0-9-]+:[0-9]{12}:runtime/.+$", + "validation_message": "Enter the complete Bedrock AgentCore runtime ARN, including the runtime ID after \"runtime/\" (e.g. arn:aws:bedrock-agentcore:us-west-2:123456789012:runtime/my-agent-runtime)." } ], "litellm_params_template": { diff --git a/litellm/types/proxy/public_endpoints/public_endpoints.py b/litellm/types/proxy/public_endpoints/public_endpoints.py index f6ee054ceaa..c7f80a61e0f 100644 --- a/litellm/types/proxy/public_endpoints/public_endpoints.py +++ b/litellm/types/proxy/public_endpoints/public_endpoints.py @@ -43,6 +43,8 @@ class AgentCredentialField(BaseModel): options: list[str] | None = None default_value: str | None = None include_in_litellm_params: bool | None = None + validation_pattern: str | None = None + validation_message: str | None = None class AgentCreateInfo(BaseModel): diff --git a/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py b/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py index 31430da71e8..8006f64ba41 100644 --- a/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py +++ b/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py @@ -1,3 +1,4 @@ +import re from datetime import datetime, timezone from unittest.mock import AsyncMock, MagicMock, patch @@ -756,6 +757,45 @@ def test_public_agent_hub_returns_empty_when_no_public_groups(): assert response.json() == [] +# --------------------------------------------------------------------------- +# /public/agents/fields +# --------------------------------------------------------------------------- + + +def test_bedrock_agentcore_runtime_arn_validation_pattern_accepts_full_resource_path(): + """Regression for LIT-6737: the AgentCore agent_runtime_arn field's + validation_pattern must accept a complete runtime ARN whose resource part + is itself multi-segment (``runtime/``), and reject the exact + truncated shape a naive split("/")-by-position parse used to produce (the + ARN cut off right after the ``runtime`` resource type). + """ + app_instance = FastAPI() + app_instance.include_router(router) + test_client = TestClient(app_instance) + + response = test_client.get("/public/agents/fields") + assert response.status_code == 200 + agents = response.json() + + bedrock_agentcore = next((a for a in agents if a["agent_type"] == "bedrock_agentcore"), None) + assert bedrock_agentcore is not None, "bedrock_agentcore agent type not found" + assert bedrock_agentcore["model_template"] == "bedrock/agentcore/{agent_runtime_arn}" + + fields_by_key = {f["key"]: f for f in bedrock_agentcore["credential_fields"]} + arn_field = fields_by_key["agent_runtime_arn"] + assert arn_field["required"] is True + assert arn_field["include_in_litellm_params"] is False + + pattern = arn_field.get("validation_pattern") + assert pattern, "agent_runtime_arn must ship a validation_pattern so the UI can reject a truncated ARN" + + full_arn = "arn:aws:bedrock-agentcore:eu-central-1:123456789012:runtime/hosted_agent_4vm3i-BaTdfOELAs" + truncated_arn = "arn:aws:bedrock-agentcore:eu-central-1:123456789012:runtime" + + assert re.match(pattern, full_arn), "the validator must accept a complete runtime ARN" + assert not re.match(pattern, truncated_arn), "the validator must reject the truncated ARN" + + # --------------------------------------------------------------------------- # /public/endpoints # --------------------------------------------------------------------------- diff --git a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.integration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.integration.test.tsx index 79bd2f6a21b..de1eb153f6f 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.integration.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.integration.test.tsx @@ -75,6 +75,40 @@ const langgraphInfo: AgentCreateInfo = { ], }; +const FULL_RUNTIME_ARN = "arn:aws:bedrock-agentcore:eu-central-1:123456789012:runtime/hosted_agent_4vm3i-BaTdfOELAs"; + +const BEDROCK_AGENTCORE_AGENT = { + agent_id: "agent-3", + agent_name: "bedrock-agent", + agent_card_params: { name: "bedrock-agent", description: "agentcore agent", url: "", version: "1.0.0", skills: [] }, + litellm_params: { + custom_llm_provider: "bedrock", + model: `bedrock/agentcore/${FULL_RUNTIME_ARN}`, + }, +}; + +const bedrockAgentcoreInfo: AgentCreateInfo = { + agent_type: "bedrock_agentcore", + agent_type_display_name: "Bedrock AgentCore", + description: "Bedrock AgentCore runtimes", + logo_url: "/b.png", + use_a2a_form_fields: false, + litellm_params_template: { custom_llm_provider: "bedrock" }, + model_template: "bedrock/agentcore/{agent_runtime_arn}", + credential_fields: [ + { + key: "agent_runtime_arn", + label: "Agent Runtime ARN", + field_type: "text", + required: true, + include_in_litellm_params: false, + validation_pattern: "^arn:aws[a-zA-Z0-9-]*:bedrock-agentcore:[a-z0-9-]+:[0-9]{12}:runtime/.+$", + validation_message: + 'Enter the complete Bedrock AgentCore runtime ARN, including the runtime ID after "runtime/".', + }, + ], +}; + const setup = () => userEvent.setup({ pointerEventsCheck: PointerEventsCheckLevel.Never }); const renderView = () => render(); @@ -247,6 +281,69 @@ describe("AgentInfoView update payload", () => { }); }); + it("preserves the full AgentCore runtime ARN (including the resource id after runtime/) across an unedited save", async () => { + vi.mocked(networking.getAgentCreateMetadata).mockResolvedValue([bedrockAgentcoreInfo]); + vi.mocked(networking.getAgentInfo).mockResolvedValue(BEDROCK_AGENTCORE_AGENT as never); + const user = setup(); + renderView(); + await openEditor(user); + + expect(await screen.findByLabelText("Agent Runtime ARN")).toHaveValue(FULL_RUNTIME_ARN); + + await save(user); + + expect((patchedPayload().litellm_params as Record).model).toBe( + `bedrock/agentcore/${FULL_RUNTIME_ARN}`, + ); + }); + + it("blocks the save and shows a validation error when the Agent Runtime ARN is truncated", async () => { + vi.mocked(networking.getAgentCreateMetadata).mockResolvedValue([bedrockAgentcoreInfo]); + vi.mocked(networking.getAgentInfo).mockResolvedValue(BEDROCK_AGENTCORE_AGENT as never); + const user = setup(); + renderView(); + await openEditor(user); + + const arnField = await screen.findByLabelText("Agent Runtime ARN"); + await user.clear(arnField); + fireEvent.change(arnField, { + target: { value: "arn:aws:bedrock-agentcore:eu-central-1:123456789012:runtime" }, + }); + await user.click(screen.getByRole("button", { name: "Save Changes" })); + + expect( + await screen.findByText( + 'Enter the complete Bedrock AgentCore runtime ARN, including the runtime ID after "runtime/".', + ), + ).toBeInTheDocument(); + expect(networking.patchAgentCall).not.toHaveBeenCalled(); + }); + + it("renders and saves normally when a field's validation_pattern is not a valid regex", async () => { + const infoWithBadPattern: AgentCreateInfo = { + ...bedrockAgentcoreInfo, + credential_fields: [ + { + ...bedrockAgentcoreInfo.credential_fields[0], + validation_pattern: "(unterminated", + }, + ], + }; + vi.mocked(networking.getAgentCreateMetadata).mockResolvedValue([infoWithBadPattern]); + vi.mocked(networking.getAgentInfo).mockResolvedValue(BEDROCK_AGENTCORE_AGENT as never); + const user = setup(); + renderView(); + await openEditor(user); + + expect(await screen.findByLabelText("Agent Runtime ARN")).toHaveValue(FULL_RUNTIME_ARN); + + await save(user); + + expect((patchedPayload().litellm_params as Record).model).toBe( + `bedrock/agentcore/${FULL_RUNTIME_ARN}`, + ); + }); + it("reloads the agent and leaves edit mode when the edit is cancelled", async () => { const user = setup(); renderView(); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_type_utils.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_type_utils.test.ts new file mode 100644 index 00000000000..0c2d2500776 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_type_utils.test.ts @@ -0,0 +1,73 @@ +import { describe, it, expect } from "vitest"; +import { detectAgentType, extractModelTemplateValues, parseDynamicAgentForForm } from "./agent_type_utils"; +import type { AgentCreateInfo } from "@/components/networking"; +import type { Agent } from "@/components/agents/types"; + +const FULL_RUNTIME_ARN = "arn:aws:bedrock-agentcore:eu-central-1:123456789012:runtime/hosted_agent_4vm3i-BaTdfOELAs"; + +const bedrockAgentcoreInfo: AgentCreateInfo = { + agent_type: "bedrock_agentcore", + agent_type_display_name: "Bedrock AgentCore", + model_template: "bedrock/agentcore/{agent_runtime_arn}", + credential_fields: [ + { + key: "agent_runtime_arn", + label: "Agent Runtime ARN", + required: true, + include_in_litellm_params: false, + }, + ], +}; + +describe("extractModelTemplateValues", () => { + it("recovers a placeholder value that itself contains '/' (an AWS ARN resource path)", () => { + const values = extractModelTemplateValues( + "bedrock/agentcore/{agent_runtime_arn}", + `bedrock/agentcore/${FULL_RUNTIME_ARN}`, + ); + + expect(values.agent_runtime_arn).toBe(FULL_RUNTIME_ARN); + }); + + it("recovers a placeholder value with no '/' (single path segment)", () => { + const values = extractModelTemplateValues("langgraph/{assistant_id}", "langgraph/asst_1"); + + expect(values.assistant_id).toBe("asst_1"); + }); + + it("returns no match when the model does not fit the template", () => { + const values = extractModelTemplateValues("langgraph/{assistant_id}", "azure_ai/agents/asst_1"); + + expect(values).toEqual({}); + }); +}); + +describe("parseDynamicAgentForForm", () => { + it("preserves the full runtime ARN, including the resource id after 'runtime/', when populating the edit form", () => { + const agent = { + agent_id: "agent-1", + agent_name: "bedrock-agent", + agent_card_params: { description: "" }, + litellm_params: { + custom_llm_provider: "bedrock", + model: `bedrock/agentcore/${FULL_RUNTIME_ARN}`, + }, + } as unknown as Agent; + + const values = parseDynamicAgentForForm(agent, bedrockAgentcoreInfo); + + expect(values.agent_runtime_arn).toBe(FULL_RUNTIME_ARN); + }); +}); + +describe("detectAgentType", () => { + it("detects bedrock_agentcore agents from the model prefix", () => { + const agent = { + agent_id: "agent-1", + agent_name: "bedrock-agent", + litellm_params: { model: `bedrock/agentcore/${FULL_RUNTIME_ARN}` }, + } as unknown as Agent; + + expect(detectAgentType(agent)).toBe("bedrock_agentcore"); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_type_utils.ts b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_type_utils.ts index f91590c5732..506f355c519 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_type_utils.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_type_utils.ts @@ -25,6 +25,29 @@ export const detectAgentType = (agent: Agent): string => { return "a2a"; }; +const escapeRegExp = (segment: string): string => segment.replace(/[.*+?^${}()|[\]\\]/g, "\\$&"); + +/** + * Reverses a `model_template` (e.g. "bedrock/agentcore/{agent_runtime_arn}") against a stored + * `model` string to recover the placeholder values that produced it. Builds a regex from the + * template's literal segments rather than matching by split("/") position, because a + * placeholder's value can itself contain "/" (an AWS ARN's "runtime/" resource path, + * a Vertex AI reasoning engine's "projects/.../reasoningEngines/..." resource id) and would + * otherwise be cut off at the first one. + */ +export const extractModelTemplateValues = (template: string, model: string): Record => { + // Splitting on a regex with a capturing group interleaves the captured placeholder + // names between the surrounding literal segments, e.g. "a/{x}/b" -> ["a/", "x", "/b"]. + const parts = template.split(/\{([a-zA-Z0-9_]+)\}/g); + const fieldNames = parts.filter((_part, index) => index % 2 === 1); + const pattern = parts.map((part, index) => (index % 2 === 1 ? "(.+)" : escapeRegExp(part))).join(""); + + const match = model.match(new RegExp(`^${pattern}$`)); + if (!match) return {}; + + return Object.fromEntries(fieldNames.map((name, index) => [name, match[index + 1]])); +}; + /** * Parses agent data for dynamic form fields (non-A2A agents). * Extracts values from litellm_params based on the agent type metadata. @@ -35,24 +58,18 @@ export const parseDynamicAgentForForm = (agent: Agent, agentTypeInfo: AgentCreat description: agent.agent_card_params?.description || "", }; + const templateValues = + agentTypeInfo.model_template && agent.litellm_params?.model + ? extractModelTemplateValues(agentTypeInfo.model_template, agent.litellm_params.model) + : {}; + // Extract credential field values from litellm_params for (const field of agentTypeInfo.credential_fields) { if (field.include_in_litellm_params !== false) { values[field.key] = agent.litellm_params?.[field.key] || field.default_value || ""; - } else { - // For fields not in litellm_params (like agent_id), try to extract from model string - if (agentTypeInfo.model_template && agent.litellm_params?.model) { - const model = agent.litellm_params.model; - const templateParts = agentTypeInfo.model_template.split("/"); - const modelParts = model.split("/"); - - // Find the placeholder position and extract the value - templateParts.forEach((part, index) => { - if (part === `{${field.key}}` && modelParts[index]) { - values[field.key] = modelParts[index]; - } - }); - } + } else if (templateValues[field.key] !== undefined) { + // For fields not in litellm_params (like agent_runtime_arn), recover from the model string + values[field.key] = templateValues[field.key]; } } diff --git a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/dynamic_agent_form_fields.tsx b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/dynamic_agent_form_fields.tsx index 04a8b0df9d9..0f3355fe073 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/dynamic_agent_form_fields.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/dynamic_agent_form_fields.tsx @@ -24,58 +24,80 @@ interface DynamicAgentFormFieldsProps { export const unmountedDynamicFieldNames = (mountedPanels: readonly string[]): readonly string[] => mountedPanels.includes(AGENT_FORM_CONFIG.cost.key) ? [] : COST_FIELD_NAMES; -const CredentialField = ({ field }: { field: AgentCredentialFieldMetadata }) => ( - - {({ value, onChange, ref, ...control }) => { - const text = typeof value === "string" ? value : ""; - if (field.field_type === "password") { - return ( - - ); - } - if (field.field_type === "textarea") { - return ( -