From 9b5da78531c0f48f977690b946be7a7f2ad3c005 Mon Sep 17 00:00:00 2001 From: mateo Date: Fri, 16 Oct 2026 21:38:28 +0000 Subject: [PATCH 01/24] fix(router): drop supports_web_search: false deployments for Anthropic and Gemini web search tools Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/router_utils/common_utils.py | 31 +++++++++++++---- .../test_router_utils_common_utils.py | 33 +++++++++++++++++++ 2 files changed, 57 insertions(+), 7 deletions(-) diff --git a/litellm/router_utils/common_utils.py b/litellm/router_utils/common_utils.py index c4e4d9ccc0b..316d76e67b1 100644 --- a/litellm/router_utils/common_utils.py +++ b/litellm/router_utils/common_utils.py @@ -30,6 +30,17 @@ _V = TypeVar("_V") MID_STREAM_CONTINUATION_KWARG: Final = "_mid_stream_continuation" ROUTER_ONLY_CALL_KWARGS: Final = frozenset({"silent_model", "include_fallback_errors", MID_STREAM_CONTINUATION_KWARG}) _OBJECT_MAPPING_ADAPTER: Final[TypeAdapter[Mapping[object, object]]] = TypeAdapter(Mapping[object, object]) +_OBJECT_TUPLE_ADAPTER: Final[TypeAdapter[tuple[object, ...]]] = TypeAdapter(tuple[object, ...]) +_GEMINI_SEARCH_TOOL_KEYS: Final = frozenset( + { + "googleSearch", + "google_search", + "googleSearchRetrieval", + "google_search_retrieval", + "enterpriseWebSearch", + "enterprise_web_search", + } +) def without_router_only_kwargs(kwargs: Mapping[str, _V]) -> dict[str, _V]: @@ -329,6 +340,16 @@ def _deployment_supports_web_search(deployment: dict) -> bool: return True +def _is_web_search_tool(tool: object) -> bool: + if not isinstance(tool, Mapping): + return False + tool_mapping: Final = _OBJECT_MAPPING_ADAPTER.validate_python(tool) + tool_type: Final = tool_mapping.get("type") + if isinstance(tool_type, str) and (tool_type == "web_search" or tool_type.startswith("web_search_")): + return True + return any(key in _GEMINI_SEARCH_TOOL_KEYS for key in tool_mapping if isinstance(key, str)) + + def filter_web_search_deployments( healthy_deployments: list[dict] | dict, request_kwargs: dict | None = None, @@ -343,13 +364,9 @@ def filter_web_search_deployments( if isinstance(healthy_deployments, dict): return healthy_deployments - is_web_search_request = False - tools: Final = request_kwargs.get("tools") or [] - for tool in tools: - # These are the two websearch tools for OpenAI / Azure. - if tool.get("type") == "web_search" or tool.get("type") == "web_search_preview": - is_web_search_request = True - break + tools_value: Final[object] = request_kwargs.get("tools") or [] + tools: Final = _OBJECT_TUPLE_ADAPTER.validate_python(tools_value) + is_web_search_request: Final = any(_is_web_search_tool(tool) for tool in tools) if not is_web_search_request: return healthy_deployments diff --git a/tests/unit/router_utils/test_router_utils_common_utils.py b/tests/unit/router_utils/test_router_utils_common_utils.py index af8fca72252..0ee7b81ec19 100644 --- a/tests/unit/router_utils/test_router_utils_common_utils.py +++ b/tests/unit/router_utils/test_router_utils_common_utils.py @@ -295,6 +295,39 @@ class TestFilterWebSearchDeployments: result_ids = [d["model_info"]["id"] for d in result] assert "deployment-3" not in result_ids + @pytest.mark.parametrize( + "tool", + [ + {"type": "web_search_20250305", "name": "web_search"}, + {"type": "web_search_20260209", "name": "web_search"}, + {"type": "web_search_preview_2025_03_11"}, + {"googleSearch": {}}, + {"google_search": {}}, + {"googleSearchRetrieval": {}}, + {"enterpriseWebSearch": {}}, + ], + ) + def test_native_web_search_tools_filter_unsupported( + self, sample_deployments: list[dict[str, object]], tool: dict[str, object] + ) -> None: + result = filter_web_search_deployments(sample_deployments, {"tools": [tool]}) + assert result == sample_deployments[:2] + + @pytest.mark.parametrize( + "tool", + [ + {"type": "web_fetch_20250910", "name": "web_fetch"}, + {"type": "function", "function": {"name": "web_search"}}, + {"type": "function", "name": "googleSearch"}, + {"urlContext": {}}, + ], + ) + def test_non_web_search_tools_preserve_deployments( + self, sample_deployments: list[dict[str, object]], tool: dict[str, object] + ) -> None: + result = filter_web_search_deployments(sample_deployments, {"tools": [tool]}) + assert result == sample_deployments + def test_web_search_with_other_tools(self, sample_deployments): """Web search filtering works when mixed with other tools""" request_kwargs = { From 07546d5189301f8e0225558b89a8d0193c22eb87 Mon Sep 17 00:00:00 2001 From: mateo Date: Fri, 16 Oct 2026 21:54:16 +0000 Subject: [PATCH 02/24] test(router): annotate new web search filter test locals with Final Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/unit/router_utils/test_router_utils_common_utils.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/unit/router_utils/test_router_utils_common_utils.py b/tests/unit/router_utils/test_router_utils_common_utils.py index 0ee7b81ec19..6c52d43dfba 100644 --- a/tests/unit/router_utils/test_router_utils_common_utils.py +++ b/tests/unit/router_utils/test_router_utils_common_utils.py @@ -310,7 +310,7 @@ class TestFilterWebSearchDeployments: def test_native_web_search_tools_filter_unsupported( self, sample_deployments: list[dict[str, object]], tool: dict[str, object] ) -> None: - result = filter_web_search_deployments(sample_deployments, {"tools": [tool]}) + result: Final = filter_web_search_deployments(sample_deployments, {"tools": [tool]}) assert result == sample_deployments[:2] @pytest.mark.parametrize( @@ -325,7 +325,7 @@ class TestFilterWebSearchDeployments: def test_non_web_search_tools_preserve_deployments( self, sample_deployments: list[dict[str, object]], tool: dict[str, object] ) -> None: - result = filter_web_search_deployments(sample_deployments, {"tools": [tool]}) + result: Final = filter_web_search_deployments(sample_deployments, {"tools": [tool]}) assert result == sample_deployments def test_web_search_with_other_tools(self, sample_deployments): From 85a070a39382f322d9bd228a33cbf929ebb43d33 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 28 Oct 2026 23:12:37 -0700 Subject: [PATCH 03/24] fix(router): filter web search deployments on the sync pick, skip null Gemini search keys, and keep every deployment when none supports search --- litellm/router.py | 5 +- litellm/router_utils/common_utils.py | 34 ++++--- .../test_router_utils_common_utils.py | 93 +++++++++++++++++-- 3 files changed, 112 insertions(+), 20 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index fc888495e99..5df73364782 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -16416,7 +16416,10 @@ class Router: ) healthy_deployments = ( # pyright: ignore[reportUnknownVariableType] # rebind-ok: route filter narrows self._without_deployments_that_cannot_serve_route( - healthy_deployments, # pyright: ignore[reportUnknownArgumentType] # legacy selector type + filter_web_search_deployments( + healthy_deployments, # pyright: ignore[reportUnknownArgumentType] # legacy selector type + request_kwargs, + ), model, ) ) diff --git a/litellm/router_utils/common_utils.py b/litellm/router_utils/common_utils.py index 28279910876..ae381768023 100644 --- a/litellm/router_utils/common_utils.py +++ b/litellm/router_utils/common_utils.py @@ -13,9 +13,11 @@ import litellm from litellm._logging import verbose_logger, verbose_router_logger from litellm.constants import ROUTER_FALLBACK_ERROR_DETAIL_MAX_CHARS from litellm.exceptions import BadRequestError +from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap from litellm.litellm_core_utils.sensitive_data_masker import mask_sensitive_structure +from litellm.rust_bridge.public_call import optional_sequence from litellm.types.router import ( CredentialLiteLLMParams, Deployment, @@ -35,7 +37,6 @@ _V = TypeVar("_V") MID_STREAM_CONTINUATION_KWARG: Final = "_mid_stream_continuation" ROUTER_ONLY_CALL_KWARGS: Final = frozenset({"silent_model", "include_fallback_errors", MID_STREAM_CONTINUATION_KWARG}) _OBJECT_MAPPING_ADAPTER: Final[TypeAdapter[Mapping[object, object]]] = TypeAdapter(Mapping[object, object]) -_OBJECT_TUPLE_ADAPTER: Final[TypeAdapter[tuple[object, ...]]] = TypeAdapter(tuple[object, ...]) _GEMINI_SEARCH_TOOL_KEYS: Final = frozenset( { "googleSearch", @@ -46,6 +47,7 @@ _GEMINI_SEARCH_TOOL_KEYS: Final = frozenset( "enterprise_web_search", } ) +_WEB_SEARCH_TOOL_TYPES: Final = frozenset({"web_search", "google_search"}) _DEPLOYMENT_ROWS_ADAPTER: Final[TypeAdapter[tuple[Mapping[str, object], ...]]] = TypeAdapter( tuple[Mapping[str, object], ...] ) @@ -439,13 +441,19 @@ def _deployment_supports_web_search(deployment: dict) -> bool: def _is_web_search_tool(tool: object) -> bool: - if not isinstance(tool, Mapping): + tool_fields: Final = as_str_mapping(tool) + if tool_fields is None: return False - tool_mapping: Final = _OBJECT_MAPPING_ADAPTER.validate_python(tool) - tool_type: Final = tool_mapping.get("type") - if isinstance(tool_type, str) and (tool_type == "web_search" or tool_type.startswith("web_search_")): + tool_type: Final = tool_fields.get("type") + if isinstance(tool_type, str) and (tool_type in _WEB_SEARCH_TOOL_TYPES or tool_type.startswith("web_search_")): return True - return any(key in _GEMINI_SEARCH_TOOL_KEYS for key in tool_mapping if isinstance(key, str)) + return not _GEMINI_SEARCH_TOOL_KEYS.isdisjoint(tool_fields) and any( + tool_fields[key] is not None for key in _GEMINI_SEARCH_TOOL_KEYS.intersection(tool_fields) + ) + + +def _is_web_search_request(tools: object, web_search_options: object) -> bool: + return web_search_options is not None or any(_is_web_search_tool(tool) for tool in optional_sequence(tools) or ()) def filter_web_search_deployments( @@ -463,17 +471,17 @@ def filter_web_search_deployments( return healthy_deployments tools: Final[object] = request_kwargs.get("tools") - is_web_search_request: Final = isinstance(tools, (list, tuple)) and any( - _is_web_search_tool(tool) for tool in _OBJECT_TUPLE_ADAPTER.validate_python(tools) - ) - if not is_web_search_request: + web_search_options: Final[object] = request_kwargs.get("web_search_options") + if not _is_web_search_request(tools, web_search_options): return healthy_deployments # Filter out deployments that don't support web search final_deployments: Final = [d for d in healthy_deployments if _deployment_supports_web_search(d)] - if len(healthy_deployments) > 0 and len(final_deployments) == 0: - verbose_logger.warning("No deployments support web search for request") - return final_deployments + if final_deployments: + return final_deployments + if healthy_deployments: + verbose_logger.warning("No deployment supports web search for this request, so every deployment stays eligible") + return healthy_deployments # Credential params that only one provider family reads, paired with the providers diff --git a/tests/unit/router_utils/test_router_utils_common_utils.py b/tests/unit/router_utils/test_router_utils_common_utils.py index f9e710cff46..269c9cd9bd0 100644 --- a/tests/unit/router_utils/test_router_utils_common_utils.py +++ b/tests/unit/router_utils/test_router_utils_common_utils.py @@ -318,7 +318,10 @@ class TestFilterWebSearchDeployments: {"googleSearch": {}}, {"google_search": {}}, {"googleSearchRetrieval": {}}, + {"google_search_retrieval": {}}, {"enterpriseWebSearch": {}}, + {"enterprise_web_search": {}}, + {"type": "google_search"}, ], ) def test_native_web_search_tools_filter_unsupported( @@ -334,6 +337,8 @@ class TestFilterWebSearchDeployments: {"type": "function", "function": {"name": "web_search"}}, {"type": "function", "name": "googleSearch"}, {"urlContext": {}}, + {"googleSearch": None}, + {"google_search": None, "urlContext": {}}, ], ) def test_non_web_search_tools_preserve_deployments( @@ -365,15 +370,27 @@ class TestFilterWebSearchDeployments: result = filter_web_search_deployments(deployments, request_kwargs) assert len(result) == 2 - def test_no_deployments_support_web_search(self): - """When no deployments support web search, all are filtered out""" - deployments = [ + def test_no_deployment_supporting_web_search_keeps_every_deployment(self) -> None: + deployments: Final = [ {"model_info": {"id": "d1", "supports_web_search": False}}, {"model_info": {"id": "d2", "supports_web_search": False}}, ] - request_kwargs = {"tools": [{"type": "web_search"}]} - result = filter_web_search_deployments(deployments, request_kwargs) - assert len(result) == 0 + result: Final = filter_web_search_deployments(deployments, {"tools": [{"type": "web_search"}]}) + assert result == deployments + + def test_a_tool_that_is_not_an_object_is_no_web_search_request( + self, sample_deployments: list[dict[str, object]] + ) -> None: + result: Final = filter_web_search_deployments(sample_deployments, {"tools": ["web_search", None]}) + assert result == sample_deployments + + def test_web_search_options_filter_unsupported(self, sample_deployments: list[dict[str, object]]) -> None: + result: Final = filter_web_search_deployments(sample_deployments, {"web_search_options": {}}) + assert result == sample_deployments[:2] + + def test_null_web_search_options_preserve_deployments(self, sample_deployments: list[dict[str, object]]) -> None: + result: Final = filter_web_search_deployments(sample_deployments, {"web_search_options": None}) + assert result == sample_deployments def test_missing_config_defaults_to_supported(self): """Deployments without supports_web_search config default to True""" @@ -408,6 +425,70 @@ class TestFilterWebSearchDeployments: assert result == deployment + +def _router_with_a_deployment_that_disables_web_search(routing_strategy: str) -> Router: + return Router( + model_list=[ + { + "model_name": "search-model", + "litellm_params": {"model": "openai/gpt-x", "api_key": "fake-key"}, + "model_info": {"id": "no-search", "supports_web_search": False}, + }, + { + "model_name": "search-model", + "litellm_params": {"model": "openai/gpt-y", "api_key": "fake-key"}, + "model_info": {"id": "search"}, + }, + ], + routing_strategy=routing_strategy, + ) + + +_ANTHROPIC_WEB_SEARCH_REQUEST: Final = {"tools": [{"type": "web_search_20250305", "name": "web_search"}]} + + +@pytest.mark.parametrize("routing_strategy", ["simple-shuffle", "usage-based-routing"]) +def test_sync_pick_skips_deployments_without_web_search(routing_strategy: str) -> None: + router: Final = _router_with_a_deployment_that_disables_web_search(routing_strategy) + picks: Final = { + router.get_available_deployment(model="search-model", request_kwargs=dict(_ANTHROPIC_WEB_SEARCH_REQUEST))[ + "model_info" + ]["id"] + for _ in range(20) + } + assert picks == {"search"} + + +@pytest.mark.asyncio +async def test_web_search_request_still_routes_when_every_deployment_disables_web_search() -> None: + router: Final = Router( + model_list=[ + { + "model_name": "search-model", + "litellm_params": {"model": "openai/gpt-x", "api_key": "fake-key"}, + "model_info": {"id": "no-search", "supports_web_search": False}, + } + ] + ) + deployment: Final = await router.async_get_available_deployment( + model="search-model", request_kwargs=dict(_ANTHROPIC_WEB_SEARCH_REQUEST) + ) + assert deployment["model_info"]["id"] == "no-search" + + +@pytest.mark.asyncio +async def test_legacy_usage_based_async_pick_skips_deployments_without_web_search() -> None: + router: Final = _router_with_a_deployment_that_disables_web_search("usage-based-routing") + picks: Final = { + ( + await router.async_get_available_deployment( + model="search-model", request_kwargs=dict(_ANTHROPIC_WEB_SEARCH_REQUEST) + ) + )["model_info"]["id"] + for _ in range(20) + } + assert picks == {"search"} + def test_invalidate_model_group_info_cache(): """Test that _invalidate_model_group_info_cache clears the LRU cache.""" router = Router( From 7025ed6b89eade1fcd0ba21f7c99f7412b6a0de6 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 29 Oct 2026 00:17:17 -0700 Subject: [PATCH 04/24] fix(router): keep the empty deployment list when no deployment supports web search --- litellm/router_utils/common_utils.py | 8 ++--- .../test_router_utils_common_utils.py | 36 +++++++++++++------ 2 files changed, 28 insertions(+), 16 deletions(-) diff --git a/litellm/router_utils/common_utils.py b/litellm/router_utils/common_utils.py index ae381768023..4ae99518685 100644 --- a/litellm/router_utils/common_utils.py +++ b/litellm/router_utils/common_utils.py @@ -477,11 +477,9 @@ def filter_web_search_deployments( # Filter out deployments that don't support web search final_deployments: Final = [d for d in healthy_deployments if _deployment_supports_web_search(d)] - if final_deployments: - return final_deployments - if healthy_deployments: - verbose_logger.warning("No deployment supports web search for this request, so every deployment stays eligible") - return healthy_deployments + if healthy_deployments and not final_deployments: + verbose_logger.warning("No deployments support web search for request") + return final_deployments # Credential params that only one provider family reads, paired with the providers diff --git a/tests/unit/router_utils/test_router_utils_common_utils.py b/tests/unit/router_utils/test_router_utils_common_utils.py index 269c9cd9bd0..568cb4ee9e1 100644 --- a/tests/unit/router_utils/test_router_utils_common_utils.py +++ b/tests/unit/router_utils/test_router_utils_common_utils.py @@ -370,13 +370,21 @@ class TestFilterWebSearchDeployments: result = filter_web_search_deployments(deployments, request_kwargs) assert len(result) == 2 - def test_no_deployment_supporting_web_search_keeps_every_deployment(self) -> None: + @pytest.mark.parametrize( + "request_kwargs", + [ + {"tools": [{"type": "web_search"}]}, + {"tools": [{"type": "web_search_20250305", "name": "web_search"}]}, + {"tools": [{"googleSearch": {}}]}, + {"web_search_options": {}}, + ], + ) + def test_no_deployment_supporting_web_search_leaves_none(self, request_kwargs: dict[str, object]) -> None: deployments: Final = [ {"model_info": {"id": "d1", "supports_web_search": False}}, {"model_info": {"id": "d2", "supports_web_search": False}}, ] - result: Final = filter_web_search_deployments(deployments, {"tools": [{"type": "web_search"}]}) - assert result == deployments + assert filter_web_search_deployments(deployments, request_kwargs) == [] def test_a_tool_that_is_not_an_object_is_no_web_search_request( self, sample_deployments: list[dict[str, object]] @@ -460,20 +468,26 @@ def test_sync_pick_skips_deployments_without_web_search(routing_strategy: str) - @pytest.mark.asyncio -async def test_web_search_request_still_routes_when_every_deployment_disables_web_search() -> None: +async def test_web_search_request_to_a_group_without_web_search_falls_back_to_a_group_with_it() -> None: router: Final = Router( model_list=[ { - "model_name": "search-model", - "litellm_params": {"model": "openai/gpt-x", "api_key": "fake-key"}, + "model_name": "flagged-model", + "litellm_params": {"model": "openai/gpt-x", "api_key": "fake-key", "mock_response": "flagged"}, "model_info": {"id": "no-search", "supports_web_search": False}, - } - ] + }, + { + "model_name": "search-model", + "litellm_params": {"model": "openai/gpt-y", "api_key": "fake-key", "mock_response": "searched"}, + "model_info": {"id": "search"}, + }, + ], + fallbacks=[{"flagged-model": ["search-model"]}], ) - deployment: Final = await router.async_get_available_deployment( - model="search-model", request_kwargs=dict(_ANTHROPIC_WEB_SEARCH_REQUEST) + response: Final = await router.acompletion( + model="flagged-model", messages=[{"role": "user", "content": "hi"}], **_ANTHROPIC_WEB_SEARCH_REQUEST ) - assert deployment["model_info"]["id"] == "no-search" + assert response.choices[0].message.content == "searched" @pytest.mark.asyncio From 3271825b1a9cdaf73f8e8f40f0a222a203a4ed22 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 29 Oct 2026 03:09:03 -0700 Subject: [PATCH 05/24] fix(router): keep main's routing for web search when no deployment can search --- litellm/router.py | 1 + litellm/router_utils/common_utils.py | 29 ++++++-- .../test_router_utils_common_utils.py | 69 ++++++++++++++----- 3 files changed, 76 insertions(+), 23 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index 5df73364782..88032b2719c 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -16419,6 +16419,7 @@ class Router: filter_web_search_deployments( healthy_deployments, # pyright: ignore[reportUnknownArgumentType] # legacy selector type request_kwargs, + empty_for_openai_web_search=False, ), model, ) diff --git a/litellm/router_utils/common_utils.py b/litellm/router_utils/common_utils.py index 4ae99518685..f473933e7e2 100644 --- a/litellm/router_utils/common_utils.py +++ b/litellm/router_utils/common_utils.py @@ -48,6 +48,7 @@ _GEMINI_SEARCH_TOOL_KEYS: Final = frozenset( } ) _WEB_SEARCH_TOOL_TYPES: Final = frozenset({"web_search", "google_search"}) +_OPENAI_WEB_SEARCH_TOOL_TYPES: Final = frozenset({"web_search", "web_search_preview"}) _DEPLOYMENT_ROWS_ADAPTER: Final[TypeAdapter[tuple[Mapping[str, object], ...]]] = TypeAdapter( tuple[Mapping[str, object], ...] ) @@ -456,30 +457,48 @@ def _is_web_search_request(tools: object, web_search_options: object) -> bool: return web_search_options is not None or any(_is_web_search_tool(tool) for tool in optional_sequence(tools) or ()) +def _is_openai_web_search_tool(tool: object) -> bool: + tool_fields: Final = as_str_mapping(tool) + return ( + tool_fields is not None + and isinstance(tool_type := tool_fields.get("type"), str) + and tool_type in _OPENAI_WEB_SEARCH_TOOL_TYPES + ) + + +def _has_openai_web_search_tool(tools: object) -> bool: + tool_list: Final = optional_sequence(tools) + return tool_list is not None and any(_is_openai_web_search_tool(tool) for tool in tool_list) + + def filter_web_search_deployments( healthy_deployments: list[dict] | dict, request_kwargs: dict | None = None, + *, + empty_for_openai_web_search: bool = True, ) -> list[dict] | dict: """ If the request is websearch, filter out deployments that don't support web search """ - if request_kwargs is None: + request_fields: Final = as_str_mapping(request_kwargs) + if request_fields is None: return healthy_deployments # When a specific deployment was already chosen, it's returned as a dict # rather than a list - nothing to filter, just pass through if isinstance(healthy_deployments, dict): return healthy_deployments - tools: Final[object] = request_kwargs.get("tools") - web_search_options: Final[object] = request_kwargs.get("web_search_options") + tools: Final = request_fields.get("tools") + web_search_options: Final = request_fields.get("web_search_options") if not _is_web_search_request(tools, web_search_options): return healthy_deployments # Filter out deployments that don't support web search final_deployments: Final = [d for d in healthy_deployments if _deployment_supports_web_search(d)] - if healthy_deployments and not final_deployments: + empties_group: Final = empty_for_openai_web_search and _has_openai_web_search_tool(tools) + if healthy_deployments and not final_deployments and empties_group: verbose_logger.warning("No deployments support web search for request") - return final_deployments + return final_deployments if final_deployments or empties_group else healthy_deployments # Credential params that only one provider family reads, paired with the providers diff --git a/tests/unit/router_utils/test_router_utils_common_utils.py b/tests/unit/router_utils/test_router_utils_common_utils.py index 568cb4ee9e1..a9c9cf629a9 100644 --- a/tests/unit/router_utils/test_router_utils_common_utils.py +++ b/tests/unit/router_utils/test_router_utils_common_utils.py @@ -370,21 +370,31 @@ class TestFilterWebSearchDeployments: result = filter_web_search_deployments(deployments, request_kwargs) assert len(result) == 2 - @pytest.mark.parametrize( - "request_kwargs", - [ - {"tools": [{"type": "web_search"}]}, - {"tools": [{"type": "web_search_20250305", "name": "web_search"}]}, - {"tools": [{"googleSearch": {}}]}, - {"web_search_options": {}}, - ], - ) - def test_no_deployment_supporting_web_search_leaves_none(self, request_kwargs: dict[str, object]) -> None: + @pytest.mark.parametrize("tool_type", ["web_search", "web_search_preview"]) + def test_no_deployment_supporting_openai_web_search_leaves_none(self, tool_type: str) -> None: deployments: Final = [ {"model_info": {"id": "d1", "supports_web_search": False}}, {"model_info": {"id": "d2", "supports_web_search": False}}, ] - assert filter_web_search_deployments(deployments, request_kwargs) == [] + assert filter_web_search_deployments(deployments, {"tools": [{"type": tool_type}]}) == [] + + @pytest.mark.parametrize( + "request_kwargs", + [ + {"tools": [{"type": "web_search_20250305", "name": "web_search"}]}, + {"tools": [{"type": "web_search_2025_08_26"}]}, + {"tools": [{"googleSearch": {}}]}, + {"web_search_options": {}}, + ], + ) + def test_no_deployment_supporting_other_web_search_keeps_every_deployment( + self, request_kwargs: dict[str, object] + ) -> None: + deployments: Final = [ + {"model_info": {"id": "d1", "supports_web_search": False}}, + {"model_info": {"id": "d2", "supports_web_search": False}}, + ] + assert filter_web_search_deployments(deployments, request_kwargs) == deployments def test_a_tool_that_is_not_an_object_is_no_web_search_request( self, sample_deployments: list[dict[str, object]] @@ -433,7 +443,6 @@ class TestFilterWebSearchDeployments: assert result == deployment - def _router_with_a_deployment_that_disables_web_search(routing_strategy: str) -> Router: return Router( model_list=[ @@ -467,9 +476,8 @@ def test_sync_pick_skips_deployments_without_web_search(routing_strategy: str) - assert picks == {"search"} -@pytest.mark.asyncio -async def test_web_search_request_to_a_group_without_web_search_falls_back_to_a_group_with_it() -> None: - router: Final = Router( +def _router_with_a_group_that_disables_web_search(routing_strategy: str = "simple-shuffle") -> Router: + return Router( model_list=[ { "model_name": "flagged-model", @@ -483,11 +491,35 @@ async def test_web_search_request_to_a_group_without_web_search_falls_back_to_a_ }, ], fallbacks=[{"flagged-model": ["search-model"]}], + routing_strategy=routing_strategy, ) - response: Final = await router.acompletion( - model="flagged-model", messages=[{"role": "user", "content": "hi"}], **_ANTHROPIC_WEB_SEARCH_REQUEST + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("tools", "expected_content"), + [ + ([{"type": "web_search"}], "searched"), + ([{"type": "web_search_20250305", "name": "web_search"}], "flagged"), + ([{"googleSearch": dict[str, object]()}], "flagged"), + ], +) +async def test_web_search_to_a_group_without_web_search_routes_like_main( + tools: list[dict[str, object]], expected_content: str +) -> None: + response: Final = await _router_with_a_group_that_disables_web_search().acompletion( + model="flagged-model", messages=[{"role": "user", "content": "hi"}], tools=tools ) - assert response.choices[0].message.content == "searched" + assert response.choices[0].message.content == expected_content + + +@pytest.mark.parametrize("routing_strategy", ["simple-shuffle", "usage-based-routing"]) +def test_sync_pick_keeps_a_group_without_web_search_for_openai_web_search(routing_strategy: str) -> None: + router: Final = _router_with_a_group_that_disables_web_search(routing_strategy) + deployment: Final = router.get_available_deployment( + model="flagged-model", request_kwargs={"tools": [{"type": "web_search"}]} + ) + assert deployment["model_info"]["id"] == "no-search" @pytest.mark.asyncio @@ -503,6 +535,7 @@ async def test_legacy_usage_based_async_pick_skips_deployments_without_web_searc } assert picks == {"search"} + def test_invalidate_model_group_info_cache(): """Test that _invalidate_model_group_info_cache clears the LRU cache.""" router = Router( From 585ccf378ee5e5a133558b65dc5c0727d95ce982 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 29 Oct 2026 05:26:17 -0700 Subject: [PATCH 06/24] fix(router): prefer web search deployments only after the other router filters run --- litellm/router.py | 29 ++--- litellm/router_utils/common_utils.py | 61 ++++++---- .../test_router_utils_common_utils.py | 106 ++++++++++++++++-- 3 files changed, 149 insertions(+), 47 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index 88032b2719c..33f53d9b852 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -245,6 +245,7 @@ from litellm.router_utils.common_utils import ( format_no_fallback_group_message, get_request_team_id, is_proxy_admin_request, + prefer_web_search_deployments, provider_for_generic_call, request_user_api_key_auth, resolve_model_group_alias, @@ -15245,10 +15246,13 @@ class Router: ## ORDER FILTERING ## -> if user set 'order' in deployments, return deployments with lowest order (e.g. order=1 > order=2) _target_order: Final = (request_kwargs or {}).pop("_target_order", None) - batch_candidates: Final = self._without_declared_text_only_deployments( - cast(list[dict[str, object]], healthy_deployments), # cast-ok: router deployments are plain dicts - model, - messages, + batch_candidates: Final = prefer_web_search_deployments( + self._without_declared_text_only_deployments( + cast(list[dict[str, object]], healthy_deployments), # cast-ok: router deployments are plain dicts + model, + messages, + cast(Mapping[str, object] | None, request_kwargs), # cast-ok: legacy signature uses a bare dict + ), cast(Mapping[str, object] | None, request_kwargs), # cast-ok: legacy signature uses a bare dict ) embedding_batch_admitting: Final = deployments_admitting_embedding_batch( @@ -16416,11 +16420,7 @@ class Router: ) healthy_deployments = ( # pyright: ignore[reportUnknownVariableType] # rebind-ok: route filter narrows self._without_deployments_that_cannot_serve_route( - filter_web_search_deployments( - healthy_deployments, # pyright: ignore[reportUnknownArgumentType] # legacy selector type - request_kwargs, - empty_for_openai_web_search=False, - ), + healthy_deployments, # pyright: ignore[reportUnknownArgumentType] # legacy selector type model, ) ) @@ -16481,10 +16481,13 @@ class Router: ## ORDER FILTERING ## -> if user set 'order' in deployments, return deployments with lowest order (e.g. order=1 > order=2) _target_order: Final = (request_kwargs or {}).pop("_target_order", None) - batch_candidates: Final = self._without_declared_text_only_deployments( - cast(list[dict[str, object]], healthy_deployments), # cast-ok: router deployments are plain dicts - model, - messages, + batch_candidates: Final = prefer_web_search_deployments( + self._without_declared_text_only_deployments( + cast(list[dict[str, object]], healthy_deployments), # cast-ok: router deployments are plain dicts + model, + messages, + cast(Mapping[str, object] | None, request_kwargs), # cast-ok: legacy signature uses a bare dict + ), cast(Mapping[str, object] | None, request_kwargs), # cast-ok: legacy signature uses a bare dict ) embedding_batch_admitting: Final = deployments_admitting_embedding_batch( diff --git a/litellm/router_utils/common_utils.py b/litellm/router_utils/common_utils.py index f473933e7e2..052bbdccb05 100644 --- a/litellm/router_utils/common_utils.py +++ b/litellm/router_utils/common_utils.py @@ -457,48 +457,61 @@ def _is_web_search_request(tools: object, web_search_options: object) -> bool: return web_search_options is not None or any(_is_web_search_tool(tool) for tool in optional_sequence(tools) or ()) -def _is_openai_web_search_tool(tool: object) -> bool: - tool_fields: Final = as_str_mapping(tool) - return ( - tool_fields is not None - and isinstance(tool_type := tool_fields.get("type"), str) - and tool_type in _OPENAI_WEB_SEARCH_TOOL_TYPES - ) - - -def _has_openai_web_search_tool(tools: object) -> bool: - tool_list: Final = optional_sequence(tools) - return tool_list is not None and any(_is_openai_web_search_tool(tool) for tool in tool_list) - - def filter_web_search_deployments( healthy_deployments: list[dict] | dict, request_kwargs: dict | None = None, - *, - empty_for_openai_web_search: bool = True, ) -> list[dict] | dict: """ If the request is websearch, filter out deployments that don't support web search """ - request_fields: Final = as_str_mapping(request_kwargs) - if request_fields is None: + if request_kwargs is None: return healthy_deployments # When a specific deployment was already chosen, it's returned as a dict # rather than a list - nothing to filter, just pass through if isinstance(healthy_deployments, dict): return healthy_deployments - tools: Final = request_fields.get("tools") - web_search_options: Final = request_fields.get("web_search_options") - if not _is_web_search_request(tools, web_search_options): + tools: Final[object] = request_kwargs.get("tools") + is_web_search_request: Final = isinstance(tools, (list, tuple)) and any( + isinstance(tool, dict) + and isinstance(tool_type := tool.get("type"), str) + and tool_type in _OPENAI_WEB_SEARCH_TOOL_TYPES + for tool in tools + ) + if not is_web_search_request: return healthy_deployments # Filter out deployments that don't support web search final_deployments: Final = [d for d in healthy_deployments if _deployment_supports_web_search(d)] - empties_group: Final = empty_for_openai_web_search and _has_openai_web_search_tool(tools) - if healthy_deployments and not final_deployments and empties_group: + if len(healthy_deployments) > 0 and len(final_deployments) == 0: verbose_logger.warning("No deployments support web search for request") - return final_deployments if final_deployments or empties_group else healthy_deployments + return final_deployments + + +def _deployment_id(deployment: Mapping[str, object]) -> object: + model_info: Final = as_str_mapping(deployment.get("model_info")) + return None if model_info is None else model_info.get("id") + + +def prefer_web_search_deployments( + deployments: list[dict[str, object]], # mutable-ok: router deployment-list contract + request_kwargs: Mapping[str, object] | None, +) -> list[dict[str, object]]: # mutable-ok: returns a mutable deployment list + if request_kwargs is None or not _is_web_search_request( + request_kwargs.get("tools"), request_kwargs.get("web_search_options") + ): + return deployments + tried: Final = frozenset( + ( + *as_excluded_deployment_ids(request_kwargs.get("_excluded_deployment_ids")), + *as_retry_skipped_deployment_ids(request_kwargs.get("_retry_skipped_deployment_ids")), + ) + ) + capable: Final = [d for d in deployments if _deployment_supports_web_search(d)] + if any(_deployment_id(d) not in tried for d in capable): + return capable + verbose_router_logger.debug("No deployment left for this web search request supports web search, so all stay") + return deployments # Credential params that only one provider family reads, paired with the providers diff --git a/tests/unit/router_utils/test_router_utils_common_utils.py b/tests/unit/router_utils/test_router_utils_common_utils.py index a9c9cf629a9..9387aaf487f 100644 --- a/tests/unit/router_utils/test_router_utils_common_utils.py +++ b/tests/unit/router_utils/test_router_utils_common_utils.py @@ -15,6 +15,7 @@ from litellm.router_utils.common_utils import ( add_model_file_id_mappings, filter_team_based_models, filter_web_search_deployments, + prefer_web_search_deployments, format_content_policy_no_fallback_note, format_fallback_outcome_message, provider_for_generic_call, @@ -327,7 +328,7 @@ class TestFilterWebSearchDeployments: def test_native_web_search_tools_filter_unsupported( self, sample_deployments: list[dict[str, object]], tool: dict[str, object] ) -> None: - result: Final = filter_web_search_deployments(sample_deployments, {"tools": [tool]}) + result: Final = prefer_web_search_deployments(sample_deployments, {"tools": [tool]}) assert result == sample_deployments[:2] @pytest.mark.parametrize( @@ -344,7 +345,7 @@ class TestFilterWebSearchDeployments: def test_non_web_search_tools_preserve_deployments( self, sample_deployments: list[dict[str, object]], tool: dict[str, object] ) -> None: - result: Final = filter_web_search_deployments(sample_deployments, {"tools": [tool]}) + result: Final = prefer_web_search_deployments(sample_deployments, {"tools": [tool]}) assert result == sample_deployments def test_web_search_with_other_tools(self, sample_deployments): @@ -384,30 +385,29 @@ class TestFilterWebSearchDeployments: {"tools": [{"type": "web_search_20250305", "name": "web_search"}]}, {"tools": [{"type": "web_search_2025_08_26"}]}, {"tools": [{"googleSearch": {}}]}, + {"tools": [{"type": "web_search"}]}, {"web_search_options": {}}, ], ) - def test_no_deployment_supporting_other_web_search_keeps_every_deployment( - self, request_kwargs: dict[str, object] - ) -> None: - deployments: Final = [ + def test_no_deployment_supporting_web_search_keeps_every_deployment(self, request_kwargs: dict[str, object]) -> None: + deployments: Final[list[dict[str, object]]] = [ {"model_info": {"id": "d1", "supports_web_search": False}}, {"model_info": {"id": "d2", "supports_web_search": False}}, ] - assert filter_web_search_deployments(deployments, request_kwargs) == deployments + assert prefer_web_search_deployments(deployments, request_kwargs) == deployments def test_a_tool_that_is_not_an_object_is_no_web_search_request( self, sample_deployments: list[dict[str, object]] ) -> None: - result: Final = filter_web_search_deployments(sample_deployments, {"tools": ["web_search", None]}) + result: Final = prefer_web_search_deployments(sample_deployments, {"tools": ["web_search", None]}) assert result == sample_deployments def test_web_search_options_filter_unsupported(self, sample_deployments: list[dict[str, object]]) -> None: - result: Final = filter_web_search_deployments(sample_deployments, {"web_search_options": {}}) + result: Final = prefer_web_search_deployments(sample_deployments, {"web_search_options": {}}) assert result == sample_deployments[:2] def test_null_web_search_options_preserve_deployments(self, sample_deployments: list[dict[str, object]]) -> None: - result: Final = filter_web_search_deployments(sample_deployments, {"web_search_options": None}) + result: Final = prefer_web_search_deployments(sample_deployments, {"web_search_options": None}) assert result == sample_deployments def test_missing_config_defaults_to_supported(self): @@ -536,6 +536,92 @@ async def test_legacy_usage_based_async_pick_skips_deployments_without_web_searc assert picks == {"search"} +def _router_with_a_mixed_web_search_group(enable_tag_filtering: bool) -> Router: + return Router( + model_list=[ + { + "model_name": "mixed", + "litellm_params": {"model": "openai/gpt-x", "api_key": "fake-key", "tags": ["cheap"], "order": 1}, + "model_info": {"id": "no-search", "supports_web_search": False}, + }, + { + "model_name": "mixed", + "litellm_params": {"model": "openai/gpt-y", "api_key": "fake-key", "tags": ["premium"], "order": 2}, + "model_info": {"id": "search"}, + }, + ], + enable_tag_filtering=enable_tag_filtering, + ) + + +def _cool_down(router: Router, model_ids: tuple[str, ...]) -> None: + for model_id in model_ids: + router.cooldown_cache.add_deployment_to_cooldown( + model_id=model_id, original_exception=Exception("rate limited"), exception_status=429, cooldown_time=300 + ) + + +_LATER_FILTER_CASES: Final = pytest.mark.parametrize( + ("enable_tag_filtering", "cooled_down", "request_kwargs", "expected_id"), + [ + (False, (), _ANTHROPIC_WEB_SEARCH_REQUEST, "search"), + (False, (), {"web_search_options": {}}, "search"), + (False, ("search",), _ANTHROPIC_WEB_SEARCH_REQUEST, "no-search"), + (False, ("search",), {"web_search_options": {}}, "no-search"), + (False, (), {**_ANTHROPIC_WEB_SEARCH_REQUEST, "_excluded_deployment_ids": ["search"]}, "no-search"), + (False, (), {**_ANTHROPIC_WEB_SEARCH_REQUEST, "_retry_skipped_deployment_ids": ("search",)}, "no-search"), + ], + ids=["order", "order-options", "cooldown", "cooldown-options", "weighted-failover", "retry-skip"], +) + + +@pytest.mark.asyncio +@_LATER_FILTER_CASES +async def test_async_pick_prefers_web_search_until_later_filters_leave_none_that_can( + enable_tag_filtering: bool, cooled_down: tuple[str, ...], request_kwargs: dict[str, object], expected_id: str +) -> None: + router: Final = _router_with_a_mixed_web_search_group(enable_tag_filtering) + _cool_down(router, cooled_down) + deployment: Final = await router.async_get_available_deployment( + model="mixed", messages=[{"role": "user", "content": "hi"}], request_kwargs=dict(request_kwargs) + ) + assert deployment["model_info"]["id"] == expected_id + + +@_LATER_FILTER_CASES +def test_sync_pick_prefers_web_search_until_later_filters_leave_none_that_can( + enable_tag_filtering: bool, cooled_down: tuple[str, ...], request_kwargs: dict[str, object], expected_id: str +) -> None: + router: Final = _router_with_a_mixed_web_search_group(enable_tag_filtering) + _cool_down(router, cooled_down) + deployment: Final = router.get_available_deployment( + model="mixed", messages=[{"role": "user", "content": "hi"}], request_kwargs=dict(request_kwargs) + ) + assert deployment["model_info"]["id"] == expected_id + + +@pytest.mark.asyncio +async def test_tag_routing_to_a_deployment_without_web_search_keeps_it() -> None: + deployment: Final = await _router_with_a_mixed_web_search_group(True).async_get_available_deployment( + model="mixed", + messages=[{"role": "user", "content": "hi"}], + request_kwargs={"web_search_options": {}, "metadata": {"tags": ["cheap"]}}, + ) + assert deployment["model_info"]["id"] == "no-search" + + +@pytest.mark.asyncio +async def test_openai_web_search_still_finds_no_deployment_while_the_only_one_that_can_search_cools_down() -> None: + router: Final = _router_with_a_mixed_web_search_group(False) + _cool_down(router, ("search",)) + with pytest.raises(litellm.RouterRateLimitError): + await router.async_get_available_deployment( + model="mixed", + messages=[{"role": "user", "content": "hi"}], + request_kwargs={"tools": [{"type": "web_search"}]}, + ) + + def test_invalidate_model_group_info_cache(): """Test that _invalidate_model_group_info_cache clears the LRU cache.""" router = Router( From 48f20a081ece3762cb670208fc6ced1130639774 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 29 Oct 2026 06:57:58 -0700 Subject: [PATCH 07/24] fix(router): let a web search request skip a session or key pin on a deployment that cannot search --- .../deployment_affinity_check.py | 10 +- .../test_deployment_affinity_check.py | 98 +++++++++++++++++++ 2 files changed, 105 insertions(+), 3 deletions(-) diff --git a/litellm/router_utils/pre_call_checks/deployment_affinity_check.py b/litellm/router_utils/pre_call_checks/deployment_affinity_check.py index daf7d6bb345..62272e751b0 100644 --- a/litellm/router_utils/pre_call_checks/deployment_affinity_check.py +++ b/litellm/router_utils/pre_call_checks/deployment_affinity_check.py @@ -31,7 +31,7 @@ from litellm.caching.dual_cache import DualCache from litellm.constants import SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY, SESSION_ID_GENERATED_METADATA_KEY from litellm.integrations.custom_logger import CustomLogger, Span from litellm.responses.utils import ResponsesAPIRequestUtils -from litellm.router_utils.common_utils import deployment_ids_failed_in_request +from litellm.router_utils.common_utils import deployment_ids_failed_in_request, prefer_web_search_deployments from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import CallTypes @@ -433,6 +433,10 @@ class DeploymentAffinityCheck(CustomLogger): failed_deployment_ids: Final = deployment_ids_failed_in_request( cast(Mapping[str, object], request_kwargs) # cast-ok: request kwargs keys are param names ) + pin_candidates: Final = prefer_web_search_deployments( + cast(list[dict[str, object]], healthy_deployments), # cast-ok: router deployments are plain dicts + request_kwargs, + ) user_key: Final = ( self.get_user_key_from_request_kwargs(request_kwargs=request_kwargs) if (session_affinity_active or enable_user_key) @@ -464,7 +468,7 @@ class DeploymentAffinityCheck(CustomLogger): if session_model_id and session_model_id not in failed_deployment_ids: session_deployment: Final = self._find_deployment_by_model_id( - healthy_deployments=typed_healthy_deployments, + healthy_deployments=pin_candidates, model_id=session_model_id, ) if session_deployment is not None: @@ -501,7 +505,7 @@ class DeploymentAffinityCheck(CustomLogger): return typed_healthy_deployments deployment = self._find_deployment_by_model_id( - healthy_deployments=typed_healthy_deployments, + healthy_deployments=pin_candidates, model_id=model_id, ) if deployment is None: diff --git a/tests/unit/router_utils/pre_call_checks/test_deployment_affinity_check.py b/tests/unit/router_utils/pre_call_checks/test_deployment_affinity_check.py index fa75f980cc0..799acc5d5cc 100644 --- a/tests/unit/router_utils/pre_call_checks/test_deployment_affinity_check.py +++ b/tests/unit/router_utils/pre_call_checks/test_deployment_affinity_check.py @@ -1299,3 +1299,101 @@ async def test_previous_response_id_pin_holds_even_after_its_deployment_failed() }, parent_otel_span=None, ) == [healthy_deployments[0]] + + +def _mixed_web_search_router(optional_pre_call_checks: list[str]) -> litellm.Router: # mutable-ok: Router signature + return litellm.Router( + model_list=[ + { + "model_name": "mixed", + "litellm_params": {"model": "openai/gpt-x", "api_key": "fake-key", "mock_response": "hi"}, + "model_info": {"id": "no-search", "supports_web_search": False}, + }, + { + "model_name": "mixed", + "litellm_params": {"model": "openai/gpt-y", "api_key": "fake-key", "mock_response": "hi"}, + "model_info": {"id": "search"}, + }, + ], + optional_pre_call_checks=optional_pre_call_checks, + ) + + +_CALLER_METADATA: Final = {"user_api_key_hash": "caller-1", "session_id": "session-1"} + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "web_search_request", + [{"web_search_options": {}}, {"tools": [{"type": "web_search_20250305", "name": "web_search"}]}], + ids=["web-search-options", "web-search-tool"], +) +@pytest.mark.parametrize( + ("optional_pre_call_checks", "pin_key"), + [ + ( + ["deployment_affinity"], + DeploymentAffinityCheck.get_affinity_cache_key(model_group="mixed", user_key="caller-1"), + ), + ( + ["session_affinity"], + DeploymentAffinityCheck.get_session_affinity_cache_key( + model_group="mixed", session_id="session-1", user_key="caller-1" + ), + ), + ], + ids=["user-key-pin", "session-pin"], +) +async def test_pin_on_a_deployment_without_web_search_yields_to_a_web_search_request( + optional_pre_call_checks: list[str], # mutable-ok: Router signature + pin_key: str, + web_search_request: dict[str, object], # mutable-ok: request kwargs +): + router: Final = _mixed_web_search_router(optional_pre_call_checks) + await router.cache.async_set_cache(pin_key, {"model_id": "no-search"}, ttl=3600) + messages: Final = [{"role": "user", "content": "hi"}] + + plain: Final = await router.acompletion(model="mixed", messages=messages, metadata=dict(_CALLER_METADATA)) + searching: Final = await router.acompletion( + model="mixed", messages=messages, metadata=dict(_CALLER_METADATA), **web_search_request + ) + + assert plain._hidden_params["model_id"] == "no-search" + assert searching._hidden_params["model_id"] == "search" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("search_supported", "tried"), + [((False, False), ()), ((False, True), ("search",))], + ids=["no-deployment-can-search", "capable-deployment-already-tried"], +) +async def test_pin_holds_for_a_web_search_request_when_no_untried_deployment_can_search( + search_supported: tuple[bool, bool], tried: tuple[str, ...] +): + cache: Final = AsyncMock() + cache.async_get_cache = AsyncMock(return_value={"model_id": "no-search"}) + callback: Final = DeploymentAffinityCheck( + cache=cache, ttl_seconds=60, enable_user_key_affinity=True, enable_responses_api_affinity=False + ) + healthy_deployments: Final = [ + { + "model_name": "mixed", + "litellm_params": {"model": "openai/gpt-x"}, + "model_info": {"id": model_id, "supports_web_search": supported}, + } + for model_id, supported in zip(("no-search", "search"), search_supported) + ] + + filtered: Final = await callback.async_filter_deployments( + model="mixed", + healthy_deployments=healthy_deployments, + messages=None, + request_kwargs={ + "metadata": {"user_api_key_hash": "caller-1"}, + "web_search_options": {}, + "_excluded_deployment_ids": list(tried), + }, + ) + + assert [d["model_info"]["id"] for d in filtered] == ["no-search"] From 5b2ccc8e585abf6c59629ac51df4a51dfb0262af Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 29 Oct 2026 08:51:01 -0700 Subject: [PATCH 08/24] fix(router): keep the web search preference inside each order tier --- litellm/router_utils/common_utils.py | 18 ++- .../test_router_utils_common_utils.py | 136 ++++++++++++++++-- 2 files changed, 134 insertions(+), 20 deletions(-) diff --git a/litellm/router_utils/common_utils.py b/litellm/router_utils/common_utils.py index 052bbdccb05..dd82f4e7862 100644 --- a/litellm/router_utils/common_utils.py +++ b/litellm/router_utils/common_utils.py @@ -507,11 +507,19 @@ def prefer_web_search_deployments( *as_retry_skipped_deployment_ids(request_kwargs.get("_retry_skipped_deployment_ids")), ) ) - capable: Final = [d for d in deployments if _deployment_supports_web_search(d)] - if any(_deployment_id(d) not in tried for d in capable): - return capable - verbose_router_logger.debug("No deployment left for this web search request supports web search, so all stay") - return deployments + untried: Final = (d for d in deployments if _deployment_id(d) not in tried) + orders_with_an_untried_capable_deployment: Final = frozenset( + litellm.utils.get_deployment_order(d) for d in untried if _deployment_supports_web_search(d) + ) + if not orders_with_an_untried_capable_deployment: + verbose_router_logger.debug("No deployment left for this web search request supports web search, so all stay") + return deployments + return [ + d + for d in deployments + if _deployment_supports_web_search(d) + or litellm.utils.get_deployment_order(d) not in orders_with_an_untried_capable_deployment + ] # Credential params that only one provider family reads, paired with the providers diff --git a/tests/unit/router_utils/test_router_utils_common_utils.py b/tests/unit/router_utils/test_router_utils_common_utils.py index 9387aaf487f..273973c4847 100644 --- a/tests/unit/router_utils/test_router_utils_common_utils.py +++ b/tests/unit/router_utils/test_router_utils_common_utils.py @@ -536,17 +536,33 @@ async def test_legacy_usage_based_async_pick_skips_deployments_without_web_searc assert picks == {"search"} -def _router_with_a_mixed_web_search_group(enable_tag_filtering: bool) -> Router: +def _order_param(order: int | None) -> dict[str, int]: + return {} if order is None else {"order": order} + + +def _router_with_a_mixed_web_search_group( + enable_tag_filtering: bool, orders: tuple[int | None, int | None] = (None, None) +) -> Router: return Router( model_list=[ { "model_name": "mixed", - "litellm_params": {"model": "openai/gpt-x", "api_key": "fake-key", "tags": ["cheap"], "order": 1}, + "litellm_params": { + "model": "openai/gpt-x", + "api_key": "fake-key", + "tags": ["cheap"], + **_order_param(orders[0]), + }, "model_info": {"id": "no-search", "supports_web_search": False}, }, { "model_name": "mixed", - "litellm_params": {"model": "openai/gpt-y", "api_key": "fake-key", "tags": ["premium"], "order": 2}, + "litellm_params": { + "model": "openai/gpt-y", + "api_key": "fake-key", + "tags": ["premium"], + **_order_param(orders[1]), + }, "model_info": {"id": "search"}, }, ], @@ -562,25 +578,39 @@ def _cool_down(router: Router, model_ids: tuple[str, ...]) -> None: _LATER_FILTER_CASES: Final = pytest.mark.parametrize( - ("enable_tag_filtering", "cooled_down", "request_kwargs", "expected_id"), + ("orders", "cooled_down", "request_kwargs", "expected_id"), [ - (False, (), _ANTHROPIC_WEB_SEARCH_REQUEST, "search"), - (False, (), {"web_search_options": {}}, "search"), - (False, ("search",), _ANTHROPIC_WEB_SEARCH_REQUEST, "no-search"), - (False, ("search",), {"web_search_options": {}}, "no-search"), - (False, (), {**_ANTHROPIC_WEB_SEARCH_REQUEST, "_excluded_deployment_ids": ["search"]}, "no-search"), - (False, (), {**_ANTHROPIC_WEB_SEARCH_REQUEST, "_retry_skipped_deployment_ids": ("search",)}, "no-search"), + ((None, None), (), _ANTHROPIC_WEB_SEARCH_REQUEST, "search"), + ((None, None), (), {"web_search_options": {}}, "search"), + ((1, 2), (), _ANTHROPIC_WEB_SEARCH_REQUEST, "no-search"), + ((1, 2), (), {"web_search_options": {}}, "no-search"), + ((None, None), ("search",), _ANTHROPIC_WEB_SEARCH_REQUEST, "no-search"), + ((None, None), ("search",), {"web_search_options": {}}, "no-search"), + ((None, None), (), {**_ANTHROPIC_WEB_SEARCH_REQUEST, "_excluded_deployment_ids": ["search"]}, "no-search"), + ((None, None), (), {**_ANTHROPIC_WEB_SEARCH_REQUEST, "_retry_skipped_deployment_ids": ("search",)}, "no-search"), + ], + ids=[ + "same-tier", + "same-tier-options", + "lower-order-tier", + "lower-order-tier-options", + "cooldown", + "cooldown-options", + "weighted-failover", + "retry-skip", ], - ids=["order", "order-options", "cooldown", "cooldown-options", "weighted-failover", "retry-skip"], ) @pytest.mark.asyncio @_LATER_FILTER_CASES async def test_async_pick_prefers_web_search_until_later_filters_leave_none_that_can( - enable_tag_filtering: bool, cooled_down: tuple[str, ...], request_kwargs: dict[str, object], expected_id: str + orders: tuple[int | None, int | None], + cooled_down: tuple[str, ...], + request_kwargs: dict[str, object], + expected_id: str, ) -> None: - router: Final = _router_with_a_mixed_web_search_group(enable_tag_filtering) + router: Final = _router_with_a_mixed_web_search_group(False, orders) _cool_down(router, cooled_down) deployment: Final = await router.async_get_available_deployment( model="mixed", messages=[{"role": "user", "content": "hi"}], request_kwargs=dict(request_kwargs) @@ -590,9 +620,12 @@ async def test_async_pick_prefers_web_search_until_later_filters_leave_none_that @_LATER_FILTER_CASES def test_sync_pick_prefers_web_search_until_later_filters_leave_none_that_can( - enable_tag_filtering: bool, cooled_down: tuple[str, ...], request_kwargs: dict[str, object], expected_id: str + orders: tuple[int | None, int | None], + cooled_down: tuple[str, ...], + request_kwargs: dict[str, object], + expected_id: str, ) -> None: - router: Final = _router_with_a_mixed_web_search_group(enable_tag_filtering) + router: Final = _router_with_a_mixed_web_search_group(False, orders) _cool_down(router, cooled_down) deployment: Final = router.get_available_deployment( model="mixed", messages=[{"role": "user", "content": "hi"}], request_kwargs=dict(request_kwargs) @@ -600,6 +633,79 @@ def test_sync_pick_prefers_web_search_until_later_filters_leave_none_that_can( assert deployment["model_info"]["id"] == expected_id +_WEB_SEARCH_REQUESTS: Final = pytest.mark.parametrize( + "web_search_request", + [_ANTHROPIC_WEB_SEARCH_REQUEST, {"web_search_options": {}}], + ids=["web-search-tool", "web-search-options"], +) + + +@pytest.mark.asyncio +@_WEB_SEARCH_REQUESTS +async def test_order_fallback_answers_a_web_search_request_from_the_next_tier_when_the_capable_one_fails( + web_search_request: dict[str, object], +) -> None: + router: Final = Router( + model_list=[ + { + "model_name": "ordered", + "litellm_params": { + "model": "openai/gpt-x", + "api_key": "fake-key", + "mock_response": Exception("connection error"), + "order": 1, + }, + "model_info": {"id": "search"}, + }, + { + "model_name": "ordered", + "litellm_params": {"model": "openai/gpt-y", "api_key": "fake-key", "mock_response": "hi", "order": 2}, + "model_info": {"id": "no-search", "supports_web_search": False}, + }, + ], + num_retries=0, + ) + response: Final = await router.acompletion( + model="ordered", messages=[{"role": "user", "content": "hi"}], **web_search_request + ) + assert response._hidden_params["model_id"] == "no-search" + + +@pytest.mark.asyncio +@_WEB_SEARCH_REQUESTS +async def test_web_search_request_prefers_the_capable_deployment_inside_its_order_tier( + web_search_request: dict[str, object], +) -> None: + router: Final = Router( + model_list=[ + { + "model_name": "ordered", + "litellm_params": {"model": "openai/gpt-x", "api_key": "fake-key", "order": 1}, + "model_info": {"id": "no-search", "supports_web_search": False}, + }, + { + "model_name": "ordered", + "litellm_params": {"model": "openai/gpt-y", "api_key": "fake-key", "order": 1}, + "model_info": {"id": "search"}, + }, + { + "model_name": "ordered", + "litellm_params": {"model": "openai/gpt-z", "api_key": "fake-key", "order": 2}, + "model_info": {"id": "later-no-search", "supports_web_search": False}, + }, + ], + ) + picks: Final = { + ( + await router.async_get_available_deployment( + model="ordered", messages=[{"role": "user", "content": "hi"}], request_kwargs=dict(web_search_request) + ) + )["model_info"]["id"] + for _ in range(20) + } + assert picks == {"search"} + + @pytest.mark.asyncio async def test_tag_routing_to_a_deployment_without_web_search_keeps_it() -> None: deployment: Final = await _router_with_a_mixed_web_search_group(True).async_get_available_deployment( From a179118b3eb0e62f3251ba11544e2b9bf64c78ea Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 29 Oct 2026 13:13:54 -0700 Subject: [PATCH 09/24] fix(router): prefer web search deployments at the strategy pick and fall back when none has capacity --- litellm/router.py | 73 +++++---- litellm/router_utils/common_utils.py | 65 +++++++- .../test_router_utils_common_utils.py | 154 +++++++++++++++++- 3 files changed, 252 insertions(+), 40 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index 33f53d9b852..16eadfe159f 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -235,6 +235,7 @@ from litellm.router_utils.clientside_credential_handler import ( ) from litellm.router_utils.common_utils import ( MID_STREAM_CONTINUATION_KWARG, + apick_preferring_web_search, as_retry_skipped_deployment_ids, deployments_the_caller_may_use, filter_team_based_models, @@ -245,6 +246,7 @@ from litellm.router_utils.common_utils import ( format_no_fallback_group_message, get_request_team_id, is_proxy_admin_request, + pick_preferring_web_search, prefer_web_search_deployments, provider_for_generic_call, request_user_api_key_auth, @@ -15246,13 +15248,10 @@ class Router: ## ORDER FILTERING ## -> if user set 'order' in deployments, return deployments with lowest order (e.g. order=1 > order=2) _target_order: Final = (request_kwargs or {}).pop("_target_order", None) - batch_candidates: Final = prefer_web_search_deployments( - self._without_declared_text_only_deployments( - cast(list[dict[str, object]], healthy_deployments), # cast-ok: router deployments are plain dicts - model, - messages, - cast(Mapping[str, object] | None, request_kwargs), # cast-ok: legacy signature uses a bare dict - ), + batch_candidates: Final = self._without_declared_text_only_deployments( + cast(list[dict[str, object]], healthy_deployments), # cast-ok: router deployments are plain dicts + model, + messages, cast(Mapping[str, object] | None, request_kwargs), # cast-ok: legacy signature uses a bare dict ) embedding_batch_admitting: Final = deployments_admitting_embedding_batch( @@ -15535,26 +15534,31 @@ class Router: return healthy_deployments[0] start_time: Final = time.time() + routing_kwargs: Final = cast(Mapping[str, object], request_kwargs) # cast-ok: str-keyed request kwargs if strategy == "simple-shuffle": shuffled: Final = simple_shuffle( resolve_model_alias=self.get_model_from_alias, - healthy_deployments=healthy_deployments, + healthy_deployments=prefer_web_search_deployments(healthy_deployments, routing_kwargs), model=model, - request_kwargs=request_kwargs, + request_kwargs=routing_kwargs, ) phase_event(_DEPLOYMENT_SELECTED_EVENT, pick_attributes) return shuffled with PrefetchedUsage.scoped( routing_read_batch.prefetched_usage if routing_read_batch is not None else None ): - deployment: Final = await self._select_deployment_async( - strategy=strategy, - selector=strategy_selector, - model=model, - healthy_deployments=healthy_deployments, - messages=messages, - input=input, - request_kwargs=request_kwargs, + deployment: Final = await apick_preferring_web_search( + lambda candidates: self._select_deployment_async( + strategy=strategy, + selector=strategy_selector, + model=model, + healthy_deployments=candidates, + messages=messages, + input=input, + request_kwargs=request_kwargs, + ), + cast(list[dict[str, object]], healthy_deployments), # cast-ok: router deployments are plain dicts + routing_kwargs, ) if deployment is None: exception: Final = await async_raise_no_deployment_exception( @@ -16481,13 +16485,10 @@ class Router: ## ORDER FILTERING ## -> if user set 'order' in deployments, return deployments with lowest order (e.g. order=1 > order=2) _target_order: Final = (request_kwargs or {}).pop("_target_order", None) - batch_candidates: Final = prefer_web_search_deployments( - self._without_declared_text_only_deployments( - cast(list[dict[str, object]], healthy_deployments), # cast-ok: router deployments are plain dicts - model, - messages, - cast(Mapping[str, object] | None, request_kwargs), # cast-ok: legacy signature uses a bare dict - ), + batch_candidates: Final = self._without_declared_text_only_deployments( + cast(list[dict[str, object]], healthy_deployments), # cast-ok: router deployments are plain dicts + model, + messages, cast(Mapping[str, object] | None, request_kwargs), # cast-ok: legacy signature uses a bare dict ) embedding_batch_admitting: Final = deployments_admitting_embedding_batch( @@ -16575,20 +16576,24 @@ class Router: ############## Check 'weight' param set for weighted pick ################# shuffled: Final = simple_shuffle( resolve_model_alias=self.get_model_from_alias, - healthy_deployments=healthy_deployments, + healthy_deployments=prefer_web_search_deployments(healthy_deployments, routing_kwargs), model=model, - request_kwargs=request_kwargs, + request_kwargs=routing_kwargs, ) phase_event(_DEPLOYMENT_SELECTED_EVENT, pick_attributes) return shuffled - deployment: Final = self._select_deployment_sync( - strategy=strategy, - selector=strategy_selector, - model=model, - healthy_deployments=healthy_deployments, - messages=messages, - input=input, - request_kwargs=request_kwargs, + deployment: Final = pick_preferring_web_search( + lambda candidates: self._select_deployment_sync( + strategy=strategy, + selector=strategy_selector, + model=model, + healthy_deployments=candidates, + messages=messages, + input=input, + request_kwargs=request_kwargs, + ), + cast(list[dict[str, object]], healthy_deployments), # cast-ok: router deployments are plain dicts + routing_kwargs, ) if deployment is None: diff --git a/litellm/router_utils/common_utils.py b/litellm/router_utils/common_utils.py index dd82f4e7862..5b3f923d936 100644 --- a/litellm/router_utils/common_utils.py +++ b/litellm/router_utils/common_utils.py @@ -1,6 +1,6 @@ import hashlib import json -from collections.abc import Mapping, Sequence +from collections.abc import Awaitable, Callable, Mapping, Sequence from types import MappingProxyType from typing import TYPE_CHECKING, Final, TypeVar @@ -33,6 +33,8 @@ from litellm.types.utils import ( ) _V = TypeVar("_V") +_Picked: Final = TypeVar("_Picked") +_DeploymentT: Final = TypeVar("_DeploymentT", bound=dict[str, object]) MID_STREAM_CONTINUATION_KWARG: Final = "_mid_stream_continuation" ROUTER_ONLY_CALL_KWARGS: Final = frozenset({"silent_model", "include_fallback_errors", MID_STREAM_CONTINUATION_KWARG}) @@ -53,6 +55,9 @@ _DEPLOYMENT_ROWS_ADAPTER: Final[TypeAdapter[tuple[Mapping[str, object], ...]]] = tuple[Mapping[str, object], ...] ) _DEPLOYMENT_IDS_ADAPTER: Final = TypeAdapter(tuple[str, ...]) +_RETRY_BREADCRUMBS_ADAPTER: Final[TypeAdapter[tuple[Mapping[str, object], ...]]] = TypeAdapter( + tuple[Mapping[str, object], ...] +) _FALLBACK_ENTRIES_ADAPTER: Final[TypeAdapter[tuple[Mapping[str, object], ...]]] = TypeAdapter( tuple[Mapping[str, object], ...] ) @@ -493,10 +498,31 @@ def _deployment_id(deployment: Mapping[str, object]) -> object: return None if model_info is None else model_info.get("id") +def _retry_breadcrumbs(metadata: object) -> tuple[Mapping[str, object], ...]: + fields: Final = as_str_mapping(metadata) + breadcrumbs: Final = None if fields is None else fields.get("previous_models") + if breadcrumbs is None: + return () + try: + return _RETRY_BREADCRUMBS_ADAPTER.validate_python(breadcrumbs) + except ValidationError: + return () + + +def _deployment_ids_retried_in_request(request_kwargs: Mapping[str, object]) -> tuple[str, ...]: + breadcrumbs: Final = ( + *_retry_breadcrumbs(request_kwargs.get("metadata")), + *_retry_breadcrumbs(request_kwargs.get("litellm_metadata")), + ) + return tuple( + deployment_id for breadcrumb in breadcrumbs if isinstance(deployment_id := breadcrumb.get("deployment_id"), str) + ) + + def prefer_web_search_deployments( - deployments: list[dict[str, object]], # mutable-ok: router deployment-list contract + deployments: list[_DeploymentT], # mutable-ok: router deployment-list contract request_kwargs: Mapping[str, object] | None, -) -> list[dict[str, object]]: # mutable-ok: returns a mutable deployment list +) -> list[_DeploymentT]: # mutable-ok: returns a mutable deployment list if request_kwargs is None or not _is_web_search_request( request_kwargs.get("tools"), request_kwargs.get("web_search_options") ): @@ -505,10 +531,11 @@ def prefer_web_search_deployments( ( *as_excluded_deployment_ids(request_kwargs.get("_excluded_deployment_ids")), *as_retry_skipped_deployment_ids(request_kwargs.get("_retry_skipped_deployment_ids")), + *_deployment_ids_retried_in_request(request_kwargs), ) ) untried: Final = (d for d in deployments if _deployment_id(d) not in tried) - orders_with_an_untried_capable_deployment: Final = frozenset( + orders_with_an_untried_capable_deployment: Final = tuple( litellm.utils.get_deployment_order(d) for d in untried if _deployment_supports_web_search(d) ) if not orders_with_an_untried_capable_deployment: @@ -522,6 +549,36 @@ def prefer_web_search_deployments( ] +def pick_preferring_web_search( + pick: Callable[[list[_DeploymentT]], _Picked | None], # mutable-ok: router deployment-list contract + deployments: list[_DeploymentT], # mutable-ok: router deployment-list contract + request_kwargs: Mapping[str, object] | None, +) -> _Picked | None: + preferred: Final = prefer_web_search_deployments(deployments, request_kwargs) + if len(preferred) == len(deployments): + return pick(deployments) + try: + from_preferred: Final = pick(preferred) + except (litellm.RateLimitError, ValueError): + return pick(deployments) + return from_preferred if from_preferred is not None else pick(deployments) + + +async def apick_preferring_web_search( + pick: Callable[[list[_DeploymentT]], Awaitable[_Picked | None]], # mutable-ok: router deployment-list contract + deployments: list[_DeploymentT], # mutable-ok: router deployment-list contract + request_kwargs: Mapping[str, object] | None, +) -> _Picked | None: + preferred: Final = prefer_web_search_deployments(deployments, request_kwargs) + if len(preferred) == len(deployments): + return await pick(deployments) + try: + from_preferred: Final = await pick(preferred) + except (litellm.RateLimitError, ValueError): + return await pick(deployments) + return from_preferred if from_preferred is not None else await pick(deployments) + + # Credential params that only one provider family reads, paired with the providers # that read them. A deployment carrying them while resolving elsewhere is almost # always a missing route prefix: `model: claude-sonnet-5` with `aws_region_name` diff --git a/tests/unit/router_utils/test_router_utils_common_utils.py b/tests/unit/router_utils/test_router_utils_common_utils.py index 273973c4847..0049793a354 100644 --- a/tests/unit/router_utils/test_router_utils_common_utils.py +++ b/tests/unit/router_utils/test_router_utils_common_utils.py @@ -1,3 +1,4 @@ +import asyncio import logging from itertools import chain from typing import Dict, Final, List, Optional, Union @@ -15,6 +16,7 @@ from litellm.router_utils.common_utils import ( add_model_file_id_mappings, filter_team_based_models, filter_web_search_deployments, + pick_preferring_web_search, prefer_web_search_deployments, format_content_policy_no_fallback_note, format_fallback_outcome_message, @@ -523,8 +525,11 @@ def test_sync_pick_keeps_a_group_without_web_search_for_openai_web_search(routin @pytest.mark.asyncio -async def test_legacy_usage_based_async_pick_skips_deployments_without_web_search() -> None: - router: Final = _router_with_a_deployment_that_disables_web_search("usage-based-routing") +@pytest.mark.parametrize( + "routing_strategy", ["usage-based-routing", "usage-based-routing-v2", "least-busy", "latency-based-routing"] +) +async def test_async_strategy_pick_skips_deployments_without_web_search(routing_strategy: str) -> None: + router: Final = _router_with_a_deployment_that_disables_web_search(routing_strategy) picks: Final = { ( await router.async_get_available_deployment( @@ -706,6 +711,151 @@ async def test_web_search_request_prefers_the_capable_deployment_inside_its_orde assert picks == {"search"} +@pytest.mark.asyncio +@pytest.mark.parametrize("metadata_field", ["metadata", "litellm_metadata"]) +async def test_web_search_preference_gives_way_once_the_capable_deployment_failed_this_request( + metadata_field: str, +) -> None: + router: Final = _router_with_a_mixed_web_search_group(False) + picks: Final = { + ( + await router.async_get_available_deployment( + model="mixed", + messages=[{"role": "user", "content": "hi"}], + request_kwargs={ + **_ANTHROPIC_WEB_SEARCH_REQUEST, + metadata_field: {"previous_models": ({"deployment_id": "search", "exception_type": "Timeout"},)}, + }, + ) + )["model_info"]["id"] + for _ in range(40) + } + assert picks == {"search", "no-search"} + + +@pytest.mark.asyncio +@_WEB_SEARCH_REQUESTS +async def test_retry_after_a_retryable_failure_on_the_capable_deployment_reaches_the_other_one( + web_search_request: dict[str, object], +) -> None: + router: Final = Router( + model_list=[ + { + "model_name": "mixed", + "litellm_params": { + "model": "openai/gpt-x", + "api_key": "fake-key", + "mock_response": Exception("connection error"), + "weight": 1, + }, + "model_info": {"id": "search"}, + }, + { + "model_name": "mixed", + "litellm_params": {"model": "openai/gpt-y", "api_key": "fake-key", "mock_response": "hi", "weight": 10**6}, + "model_info": {"id": "no-search", "supports_web_search": False}, + }, + ], + num_retries=1, + disable_cooldowns=True, + ) + response: Final = await router.acompletion( + model="mixed", messages=[{"role": "user", "content": "hi"}], **web_search_request + ) + assert response._hidden_params["model_id"] == "no-search" + + +def _rate_limited_web_search_group(routing_strategy: str) -> Router: + return Router( + model_list=[ + { + "model_name": "mixed", + "litellm_params": {"model": "openai/gpt-x", "api_key": "fake-key", "mock_response": "hi", "rpm": 1}, + "model_info": {"id": "search"}, + }, + { + "model_name": "mixed", + "litellm_params": {"model": "openai/gpt-y", "api_key": "fake-key", "mock_response": "hi", "rpm": 1000}, + "model_info": {"id": "no-search", "supports_web_search": False}, + }, + ], + routing_strategy=routing_strategy, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("routing_strategy", ["usage-based-routing-v2", "cost-based-routing"]) +async def test_web_search_request_goes_to_the_other_deployment_once_the_capable_one_is_at_its_rpm_limit( + routing_strategy: str, +) -> None: + router: Final = _rate_limited_web_search_group(routing_strategy) + served_by: Final = [] # mutable-ok: collects one id per sequential request + for _ in range(3): + response = await router.acompletion( + model="mixed", messages=[{"role": "user", "content": "hi"}], **_ANTHROPIC_WEB_SEARCH_REQUEST + ) + served_by.append(response._hidden_params["model_id"]) + await asyncio.sleep(0.2) + assert served_by == ["search", "no-search", "no-search"] + + +def test_sync_web_search_request_goes_to_the_other_deployment_once_the_capable_one_is_at_its_rpm_limit() -> None: + router: Final = _rate_limited_web_search_group("usage-based-routing-v2") + served_by: Final = tuple( + router.completion( + model="mixed", messages=[{"role": "user", "content": "hi"}], **_ANTHROPIC_WEB_SEARCH_REQUEST + )._hidden_params["model_id"] + for _ in range(3) + ) + assert served_by == ("search", "no-search", "no-search") + + +def test_sync_pick_falls_back_to_every_deployment_when_the_strategy_finds_none_that_can_search() -> None: + deployments: Final = [ + {"model_info": {"id": "search"}}, + {"model_info": {"id": "no-search", "supports_web_search": False}}, + ] + picked: Final = pick_preferring_web_search( + lambda candidates: next((d for d in candidates if d["model_info"]["id"] == "no-search"), None), + deployments, + _ANTHROPIC_WEB_SEARCH_REQUEST, + ) + assert picked == deployments[1] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("search_order", "no_search_order"), + [([1], [1]), ({"tier": 1}, None)], + ids=["list-order", "dict-order"], +) +async def test_web_search_request_routes_on_a_group_whose_order_is_not_a_number( + search_order: object, no_search_order: object +) -> None: + router: Final = Router( + model_list=[ + { + "model_name": "mixed", + "litellm_params": {"model": "openai/gpt-x", "api_key": "fake-key"}, + "model_info": {"id": "search", "order": search_order}, + }, + { + "model_name": "mixed", + "litellm_params": {"model": "openai/gpt-y", "api_key": "fake-key"}, + "model_info": { + "id": "no-search", + "supports_web_search": False, + **({} if no_search_order is None else {"order": no_search_order}), + }, + }, + ], + ) + deployment: Final = await router.async_get_available_deployment( + model="mixed", messages=[{"role": "user", "content": "hi"}], request_kwargs=dict(_ANTHROPIC_WEB_SEARCH_REQUEST) + ) + assert deployment["model_info"]["id"] == "search" + + @pytest.mark.asyncio async def test_tag_routing_to_a_deployment_without_web_search_keeps_it() -> None: deployment: Final = await _router_with_a_mixed_web_search_group(True).async_get_available_deployment( From 27ec8d2adf840aaa2e9c7451a372053b03b64273 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 29 Oct 2026 17:04:26 -0700 Subject: [PATCH 10/24] fix(router): count a failed deployment as tried when the caller supplies the key With caller-supplied credentials the retry record's deployment_id is the generated id, so the web search preference never saw the configured deployment that failed and sent every retry back to it. Retry records now also carry the configured id. --- litellm/router.py | 13 ++++++++++--- litellm/router_utils/common_utils.py | 4 +++- litellm/types/router.py | 3 ++- .../test_router_utils_common_utils.py | 15 ++++++++++++--- 4 files changed, 27 insertions(+), 8 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index 16eadfe159f..ec0692a937f 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -119,6 +119,7 @@ from litellm.litellm_core_utils.batch_utils import ( ) from litellm.litellm_core_utils.completion_timeout import CompletionTimeout from litellm.litellm_core_utils.core_helpers import ( + as_str_mapping, coerce_token_limit, get_litellm_metadata_from_kwargs, get_metadata_variable_name_from_kwargs, @@ -9880,14 +9881,15 @@ class Router: _metadata_var: Final = "litellm_metadata" if "litellm_metadata" in kwargs else "metadata" request_metadata: Final = cast("Mapping[str, object]", kwargs[_metadata_var]) # cast-ok: untyped kwargs model_group: Final = kwargs.get("model") - model_info: Final = request_metadata.get("model_info") - deployment_id: Final = model_info.get("id") if isinstance(model_info, Mapping) else None + model_info: Final = as_str_mapping(request_metadata.get("model_info")) or {} + deployment_id: Final = model_info.get("id") + configured_deployment_id: Final = model_info.get("original_model_id") base_model: Final = get_base_model_from_litellm_call_metadata( metadata=request_metadata ) or self._deployment_params_base_model(deployment_id if isinstance(deployment_id, str) else None) deployment_model: Final = request_metadata.get("deployment") attempted_retries: Final = request_metadata.get("attempted_retries") - attempt_record: Final[RetryAttemptRecord] = { + attempt_fields: Final[RetryAttemptRecord] = { "model_group": model_group if isinstance(model_group, str) else None, "deployment_id": deployment_id if isinstance(deployment_id, str) else None, "model": deployment_model if isinstance(deployment_model, str) else None, @@ -9896,6 +9898,11 @@ class Router: "exception_string": str(e), "attempted_retries": attempted_retries if type(attempted_retries) is int else None, } + attempt_record: Final[RetryAttemptRecord] = ( + {**attempt_fields, "configured_deployment_id": configured_deployment_id} + if isinstance(configured_deployment_id, str) + else attempt_fields + ) earlier_breadcrumbs: Final = request_metadata.get("previous_models") kept_breadcrumbs: Final[tuple[object, ...]] = ( tuple(earlier_breadcrumbs)[-(RETRY_BREADCRUMB_LIMIT - 1) :] diff --git a/litellm/router_utils/common_utils.py b/litellm/router_utils/common_utils.py index 5b3f923d936..34e12c1f751 100644 --- a/litellm/router_utils/common_utils.py +++ b/litellm/router_utils/common_utils.py @@ -515,7 +515,9 @@ def _deployment_ids_retried_in_request(request_kwargs: Mapping[str, object]) -> *_retry_breadcrumbs(request_kwargs.get("litellm_metadata")), ) return tuple( - deployment_id for breadcrumb in breadcrumbs if isinstance(deployment_id := breadcrumb.get("deployment_id"), str) + deployment_id + for breadcrumb in breadcrumbs + if isinstance(deployment_id := breadcrumb.get("configured_deployment_id", breadcrumb.get("deployment_id")), str) ) diff --git a/litellm/types/router.py b/litellm/types/router.py index f8de2128a80..75c4eea9547 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -23,7 +23,7 @@ from zoneinfo import ZoneInfo, ZoneInfoNotFoundError import httpx from openai import Timeout as SDKTimeout from pydantic import BeforeValidator, ConfigDict, Field, JsonValue, StrictInt, field_validator, model_validator -from typing_extensions import Protocol, ReadOnly, Required, TypedDict, runtime_checkable +from typing_extensions import NotRequired, Protocol, ReadOnly, Required, TypedDict, runtime_checkable from litellm._logging import verbose_logger from litellm._uuid import uuid @@ -1197,6 +1197,7 @@ class RetryAttemptRecord(TypedDict): exception_type: ReadOnly[str] exception_string: ReadOnly[str] attempted_retries: ReadOnly[int | None] + configured_deployment_id: ReadOnly[NotRequired[str]] VALID_LITELLM_ENVIRONMENTS = [ diff --git a/tests/unit/router_utils/test_router_utils_common_utils.py b/tests/unit/router_utils/test_router_utils_common_utils.py index 0049793a354..a28731afad1 100644 --- a/tests/unit/router_utils/test_router_utils_common_utils.py +++ b/tests/unit/router_utils/test_router_utils_common_utils.py @@ -735,8 +735,12 @@ async def test_web_search_preference_gives_way_once_the_capable_deployment_faile @pytest.mark.asyncio @_WEB_SEARCH_REQUESTS +@pytest.mark.parametrize( + "caller_credentials", [{}, {"api_key": "caller-key"}], ids=["configured-credentials", "caller-credentials"] +) async def test_retry_after_a_retryable_failure_on_the_capable_deployment_reaches_the_other_one( web_search_request: dict[str, object], + caller_credentials: dict[str, str], ) -> None: router: Final = Router( model_list=[ @@ -752,7 +756,12 @@ async def test_retry_after_a_retryable_failure_on_the_capable_deployment_reaches }, { "model_name": "mixed", - "litellm_params": {"model": "openai/gpt-y", "api_key": "fake-key", "mock_response": "hi", "weight": 10**6}, + "litellm_params": { + "model": "openai/gpt-y", + "api_key": "fake-key", + "mock_response": "from no-search", + "weight": 10**6, + }, "model_info": {"id": "no-search", "supports_web_search": False}, }, ], @@ -760,9 +769,9 @@ async def test_retry_after_a_retryable_failure_on_the_capable_deployment_reaches disable_cooldowns=True, ) response: Final = await router.acompletion( - model="mixed", messages=[{"role": "user", "content": "hi"}], **web_search_request + model="mixed", messages=[{"role": "user", "content": "hi"}], **web_search_request, **caller_credentials ) - assert response._hidden_params["model_id"] == "no-search" + assert response.choices[0].message.content == "from no-search" def _rate_limited_web_search_group(routing_strategy: str) -> Router: From 46aed6565f746a82d854b7c423ae071ea9da2fea Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 29 Oct 2026 18:12:10 -0700 Subject: [PATCH 11/24] fix(router): let the web search preference step aside once retry history was trimmed --- litellm/router_utils/common_utils.py | 11 +++ .../test_router_utils_common_utils.py | 84 +++++++++++++++++++ 2 files changed, 95 insertions(+) diff --git a/litellm/router_utils/common_utils.py b/litellm/router_utils/common_utils.py index 34e12c1f751..c2f70503cc2 100644 --- a/litellm/router_utils/common_utils.py +++ b/litellm/router_utils/common_utils.py @@ -509,6 +509,12 @@ def _retry_breadcrumbs(metadata: object) -> tuple[Mapping[str, object], ...]: return () +def _retry_history_was_trimmed(metadata: object) -> bool: + fields: Final = as_str_mapping(metadata) + retry_count: Final = None if fields is None else fields.get("request_retry_count") + return type(retry_count) is int and retry_count > len(_retry_breadcrumbs(metadata)) + + def _deployment_ids_retried_in_request(request_kwargs: Mapping[str, object]) -> tuple[str, ...]: breadcrumbs: Final = ( *_retry_breadcrumbs(request_kwargs.get("metadata")), @@ -529,6 +535,11 @@ def prefer_web_search_deployments( request_kwargs.get("tools"), request_kwargs.get("web_search_options") ): return deployments + if _retry_history_was_trimmed(request_kwargs.get("metadata")) or _retry_history_was_trimmed( + request_kwargs.get("litellm_metadata") + ): + verbose_router_logger.debug("This request's retry history was trimmed, so every deployment stays") + return deployments tried: Final = frozenset( ( *as_excluded_deployment_ids(request_kwargs.get("_excluded_deployment_ids")), diff --git a/tests/unit/router_utils/test_router_utils_common_utils.py b/tests/unit/router_utils/test_router_utils_common_utils.py index a28731afad1..4e76fff5aad 100644 --- a/tests/unit/router_utils/test_router_utils_common_utils.py +++ b/tests/unit/router_utils/test_router_utils_common_utils.py @@ -774,6 +774,90 @@ async def test_retry_after_a_retryable_failure_on_the_capable_deployment_reaches assert response.choices[0].message.content == "from no-search" +@pytest.mark.asyncio +@pytest.mark.parametrize("metadata_field", ["metadata", "litellm_metadata"]) +@pytest.mark.parametrize(("request_retry_count", "flagged_reachable"), [(1, False), (2, True)], ids=["kept", "trimmed"]) +async def test_web_search_preference_steps_aside_only_once_the_retry_history_was_trimmed( + metadata_field: str, request_retry_count: int, flagged_reachable: bool +) -> None: + router: Final = Router( + model_list=[ + *( + { + "model_name": "mixed", + "litellm_params": {"model": "openai/gpt-x", "api_key": "fake-key", "mock_response": "hi"}, + "model_info": {"id": deployment_id}, + } + for deployment_id in ("search-a", "search-b") + ), + { + "model_name": "mixed", + "litellm_params": { + "model": "openai/gpt-y", + "api_key": "fake-key", + "mock_response": "hi", + "weight": 10**6, + }, + "model_info": {"id": "no-search", "supports_web_search": False}, + }, + ], + ) + picks: Final = { + ( + await router.async_get_available_deployment( + model="mixed", + messages=[{"role": "user", "content": "hi"}], + request_kwargs={ + **_ANTHROPIC_WEB_SEARCH_REQUEST, + metadata_field: { + "previous_models": ({"deployment_id": "search-a", "exception_type": "Timeout"},), + "request_retry_count": request_retry_count, + }, + }, + ) + )["model_info"]["id"] + for _ in range(20) + } + assert ("no-search" in picks) is flagged_reachable + + +@pytest.mark.asyncio +async def test_retries_past_the_kept_retry_history_reach_the_deployment_that_cannot_search() -> None: + capable_deployments: Final = [ + { + "model_name": "mixed", + "litellm_params": { + "model": "openai/gpt-x", + "api_key": "fake-key", + "mock_response": Exception("connection error"), + }, + "model_info": {"id": f"search-{index}"}, + } + for index in range(5) + ] + router: Final = Router( + model_list=[ + *capable_deployments, + { + "model_name": "mixed", + "litellm_params": { + "model": "openai/gpt-y", + "api_key": "fake-key", + "mock_response": "from no-search", + "weight": 10**6, + }, + "model_info": {"id": "no-search", "supports_web_search": False}, + }, + ], + num_retries=5, + disable_cooldowns=True, + ) + response: Final = await router.acompletion( + model="mixed", messages=[{"role": "user", "content": "hi"}], **_ANTHROPIC_WEB_SEARCH_REQUEST + ) + assert response.choices[0].message.content == "from no-search" + + def _rate_limited_web_search_group(routing_strategy: str) -> Router: return Router( model_list=[ From 6b7163176a18b4e735d29d382e770a77da94679a Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 29 Oct 2026 23:28:36 -0700 Subject: [PATCH 12/24] fix(router): rotate web search retries through untried capable deployments and respect enforced rate limits --- litellm/router.py | 21 +++- litellm/router_utils/common_utils.py | 38 +++++-- .../deployment_affinity_check.py | 37 +++--- .../pre_call_checks/model_rate_limit_check.py | 55 ++++++++- .../test_deployment_affinity_check.py | 52 +++++++++ .../test_router_utils_common_utils.py | 107 +++++++++++++++++- .../test_enforce_model_rate_limits.py | 23 ++++ 7 files changed, 306 insertions(+), 27 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index 133bbcf8b8c..c6171ee7d2e 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -275,6 +275,7 @@ from litellm.router_utils.common_utils import ( truncate_fallback_error_detail, warn_on_ignored_model_capability_flags, warn_on_provider_credential_mismatch, + web_search_capable_deployments, with_admin_api_key_secret_reference, without_model_capability_flags, without_router_only_kwargs, @@ -10382,6 +10383,9 @@ class Router(DeprecatedRouterMembers): healthy_deployments = self._filter_blocked_deployments(healthy_deployments) return healthy_deployments, _all_deployments + def _model_rate_limit_checks(self) -> tuple[ModelRateLimitingCheck, ...]: + return tuple(callback for callback in self._routing_callbacks() if isinstance(callback, ModelRateLimitingCheck)) + def routing_strategy_pre_call_checks(self, deployment: dict): """ Mimics 'async_routing_strategy_pre_call_checks' @@ -15794,10 +15798,24 @@ class Router(DeprecatedRouterMembers): start_time: Final = time.time() routing_kwargs: Final = cast(Mapping[str, object], request_kwargs) # cast-ok: str-keyed request kwargs + web_search_capable: Final = web_search_capable_deployments( + request_kwargs=routing_kwargs, + deployments=cast(list[dict[str, object]], healthy_deployments), # cast-ok: str-keyed dicts + ) + at_rate_limit: Final = frozenset( + chain.from_iterable( + [ + await check.async_deployment_ids_at_limit(web_search_capable, parent_otel_span) + for check in self._model_rate_limit_checks() + ] + ) + ) if strategy == "simple-shuffle": shuffled: Final = simple_shuffle( resolve_model_alias=self._core_get_model_from_alias, - healthy_deployments=prefer_web_search_deployments(healthy_deployments, routing_kwargs), + healthy_deployments=prefer_web_search_deployments( + healthy_deployments, routing_kwargs, at_rate_limit + ), model=model, request_kwargs=routing_kwargs, ) @@ -15818,6 +15836,7 @@ class Router(DeprecatedRouterMembers): ), cast(list[dict[str, object]], healthy_deployments), # cast-ok: router deployments are plain dicts routing_kwargs, + at_rate_limit, ) if deployment is None: exception: Final = await async_raise_no_deployment_exception( diff --git a/litellm/router_utils/common_utils.py b/litellm/router_utils/common_utils.py index c2f70503cc2..775cd17afcb 100644 --- a/litellm/router_utils/common_utils.py +++ b/litellm/router_utils/common_utils.py @@ -1,6 +1,7 @@ import hashlib import json from collections.abc import Awaitable, Callable, Mapping, Sequence +from itertools import chain from types import MappingProxyType from typing import TYPE_CHECKING, Final, TypeVar @@ -459,7 +460,19 @@ def _is_web_search_tool(tool: object) -> bool: def _is_web_search_request(tools: object, web_search_options: object) -> bool: - return web_search_options is not None or any(_is_web_search_tool(tool) for tool in optional_sequence(tools) or ()) + return as_str_mapping(web_search_options) is not None or any( + _is_web_search_tool(tool) for tool in optional_sequence(tools) or () + ) + + +def web_search_capable_deployments( + deployments: Sequence[_DeploymentT], request_kwargs: Mapping[str, object] | None +) -> tuple[_DeploymentT, ...]: + if request_kwargs is None or not _is_web_search_request( + request_kwargs.get("tools"), request_kwargs.get("web_search_options") + ): + return () + return tuple(d for d in deployments if _deployment_supports_web_search(d)) def filter_web_search_deployments( @@ -515,21 +528,26 @@ def _retry_history_was_trimmed(metadata: object) -> bool: return type(retry_count) is int and retry_count > len(_retry_breadcrumbs(metadata)) +def _breadcrumb_deployment_ids(breadcrumb: Mapping[str, object]) -> tuple[str, ...]: + return tuple( + deployment_id + for deployment_id in (breadcrumb.get("deployment_id"), breadcrumb.get("configured_deployment_id")) + if isinstance(deployment_id, str) + ) + + def _deployment_ids_retried_in_request(request_kwargs: Mapping[str, object]) -> tuple[str, ...]: breadcrumbs: Final = ( *_retry_breadcrumbs(request_kwargs.get("metadata")), *_retry_breadcrumbs(request_kwargs.get("litellm_metadata")), ) - return tuple( - deployment_id - for breadcrumb in breadcrumbs - if isinstance(deployment_id := breadcrumb.get("configured_deployment_id", breadcrumb.get("deployment_id")), str) - ) + return tuple(chain.from_iterable(map(_breadcrumb_deployment_ids, breadcrumbs))) def prefer_web_search_deployments( deployments: list[_DeploymentT], # mutable-ok: router deployment-list contract request_kwargs: Mapping[str, object] | None, + at_rate_limit: frozenset[str] = frozenset(), ) -> list[_DeploymentT]: # mutable-ok: returns a mutable deployment list if request_kwargs is None or not _is_web_search_request( request_kwargs.get("tools"), request_kwargs.get("web_search_options") @@ -545,6 +563,7 @@ def prefer_web_search_deployments( *as_excluded_deployment_ids(request_kwargs.get("_excluded_deployment_ids")), *as_retry_skipped_deployment_ids(request_kwargs.get("_retry_skipped_deployment_ids")), *_deployment_ids_retried_in_request(request_kwargs), + *at_rate_limit, ) ) untried: Final = (d for d in deployments if _deployment_id(d) not in tried) @@ -557,8 +576,8 @@ def prefer_web_search_deployments( return [ d for d in deployments - if _deployment_supports_web_search(d) - or litellm.utils.get_deployment_order(d) not in orders_with_an_untried_capable_deployment + if litellm.utils.get_deployment_order(d) not in orders_with_an_untried_capable_deployment + or (_deployment_supports_web_search(d) and _deployment_id(d) not in tried) ] @@ -581,8 +600,9 @@ async def apick_preferring_web_search( pick: Callable[[list[_DeploymentT]], Awaitable[_Picked | None]], # mutable-ok: router deployment-list contract deployments: list[_DeploymentT], # mutable-ok: router deployment-list contract request_kwargs: Mapping[str, object] | None, + at_rate_limit: frozenset[str] = frozenset(), ) -> _Picked | None: - preferred: Final = prefer_web_search_deployments(deployments, request_kwargs) + preferred: Final = prefer_web_search_deployments(deployments, request_kwargs, at_rate_limit) if len(preferred) == len(deployments): return await pick(deployments) try: diff --git a/litellm/router_utils/pre_call_checks/deployment_affinity_check.py b/litellm/router_utils/pre_call_checks/deployment_affinity_check.py index 62272e751b0..21f9b0a7c1b 100644 --- a/litellm/router_utils/pre_call_checks/deployment_affinity_check.py +++ b/litellm/router_utils/pre_call_checks/deployment_affinity_check.py @@ -14,7 +14,7 @@ where routing to a consistent deployment is still beneficial. import hashlib from collections.abc import Callable, Mapping, Sequence -from typing import Final, cast +from typing import Final, TypeVar, cast from pydantic import TypeAdapter from typing_extensions import ReadOnly, TypedDict @@ -35,6 +35,8 @@ from litellm.router_utils.common_utils import deployment_ids_failed_in_request, from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import CallTypes +_DeploymentT: Final = TypeVar("_DeploymentT", bound=dict[str, object]) + class DeploymentAffinityCacheValue(TypedDict, closed=True): model_id: ReadOnly[str] @@ -366,7 +368,7 @@ class DeploymentAffinityCheck(CustomLogger): return self._pinned_model_id(winner) @staticmethod - def _find_deployment_by_model_id(healthy_deployments: list[dict], model_id: str) -> dict | None: + def _find_deployment_by_model_id(healthy_deployments: Sequence[_DeploymentT], model_id: str) -> _DeploymentT | None: for deployment in healthy_deployments: model_info = deployment.get("model_info") if not isinstance(model_info, dict): @@ -391,7 +393,7 @@ class DeploymentAffinityCheck(CustomLogger): 2. cached API-key deployment affinity """ request_kwargs = request_kwargs or {} - typed_healthy_deployments: Final = cast(list[dict], healthy_deployments) + typed_healthy_deployments: Final = cast(list[dict[str, object]], healthy_deployments) # cast-ok: router dicts if request_kwargs.get("_target_order") is not None: return typed_healthy_deployments @@ -467,17 +469,16 @@ class DeploymentAffinityCheck(CustomLogger): session_model_id = session_cache_result if session_model_id and session_model_id not in failed_deployment_ids: - session_deployment: Final = self._find_deployment_by_model_id( - healthy_deployments=pin_candidates, - model_id=session_model_id, + session_deployments: Final = self._pinned_deployments( + typed_healthy_deployments, pin_candidates, session_model_id ) - if session_deployment is not None: + if session_deployments is not None: verbose_router_logger.debug( "DeploymentAffinityCheck: session-id affinity hit -> deployment=%s session_id=%s", session_model_id, session_id, ) - return [session_deployment] + return list(session_deployments) else: verbose_router_logger.debug( "DeploymentAffinityCheck: session-id pinned deployment=%s not found in healthy_deployments", @@ -504,11 +505,8 @@ class DeploymentAffinityCheck(CustomLogger): if not model_id or model_id in failed_deployment_ids: return typed_healthy_deployments - deployment = self._find_deployment_by_model_id( - healthy_deployments=pin_candidates, - model_id=model_id, - ) - if deployment is None: + pinned_deployments: Final = self._pinned_deployments(typed_healthy_deployments, pin_candidates, model_id) + if pinned_deployments is None: verbose_router_logger.debug( "DeploymentAffinityCheck: pinned deployment=%s not found in healthy_deployments", model_id, @@ -520,7 +518,18 @@ class DeploymentAffinityCheck(CustomLogger): model_id, self._shorten_for_logs(self._hash_user_key(user_key)), ) - return [deployment] + return list(pinned_deployments) + + @classmethod + def _pinned_deployments( + cls, healthy_deployments: Sequence[_DeploymentT], pin_candidates: Sequence[_DeploymentT], model_id: str + ) -> tuple[_DeploymentT, ...] | None: + pinned: Final = cls._find_deployment_by_model_id(healthy_deployments=healthy_deployments, model_id=model_id) + if pinned is None: + return None + if cls._find_deployment_by_model_id(healthy_deployments=pin_candidates, model_id=model_id) is not None: + return (pinned,) + return (*pin_candidates, pinned) async def _read_shared_session_pin(self, cache_key: str, ttl_seconds: int, parent_otel_span: Span | None) -> object: redis_cache: Final = self.cache.redis_cache diff --git a/litellm/router_utils/pre_call_checks/model_rate_limit_check.py b/litellm/router_utils/pre_call_checks/model_rate_limit_check.py index f4b41211c6e..8cc47627d99 100644 --- a/litellm/router_utils/pre_call_checks/model_rate_limit_check.py +++ b/litellm/router_utils/pre_call_checks/model_rate_limit_check.py @@ -12,12 +12,13 @@ is logged the first time such a deployment is seen. """ import contextlib -from collections.abc import Callable, Mapping +from collections.abc import Callable, Mapping, Sequence from datetime import datetime, timedelta +from itertools import chain from typing import TYPE_CHECKING, Any, Final, Literal import httpx -from pydantic import BaseModel, ConfigDict +from pydantic import BaseModel, ConfigDict, ValidationError import litellm from litellm._internal_context import with_service_target @@ -117,6 +118,19 @@ class _DailyLimitDeployment(BaseModel): ) +def _deployment_model_id(deployment: Mapping[str, object]) -> object: + model_info: Final = deployment.get("model_info") + return model_info.get("id") if isinstance(model_info, Mapping) else None + + +def _counters_with_limits(model_id: str, counters: Sequence[tuple[str, object]]) -> tuple[tuple[str, str, int], ...]: + return tuple((model_id, key, limit) for key, limit in counters if type(limit) is int) + + +def _reached(current: object, limit: int) -> bool: + return isinstance(current, (int, float)) and current >= limit + + def _reset_offset_from_kwargs(kwargs: Mapping[str, object]) -> timedelta: litellm_params: Final = kwargs.get("litellm_params") if litellm_params is None: @@ -299,6 +313,43 @@ class ModelRateLimitingCheck(CustomLogger): except RedisCircuitBreakerOpenError: return local_tpm + def _limit_counters( + self, + deployment: dict[str, object], # mutable-ok: the limit and key readers take router deployment dicts + now: datetime, + ) -> tuple[tuple[str, str, int], ...]: + model_id: Final = _deployment_model_id(deployment) + if not isinstance(model_id, str): + return () + tpm_limit, rpm_limit = self._get_deployment_limits(deployment) + tpm_key, rpm_key = self._get_cache_keys(deployment, now.strftime("%H-%M")) + try: + daily: Final = _DailyLimitDeployment.model_validate(deployment) + except ValidationError: + return _counters_with_limits(model_id, ((tpm_key, tpm_limit), (rpm_key, rpm_limit))) + tpd_limit, rpd_limit = daily.limits() + tpd_key, rpd_key = daily.cache_keys(now) + return _counters_with_limits( + model_id, ((tpm_key, tpm_limit), (rpm_key, rpm_limit), (tpd_key, tpd_limit), (rpd_key, rpd_limit)) + ) + + async def async_deployment_ids_at_limit( + self, + deployments: Sequence[dict[str, object]], # mutable-ok: _limit_counters reads router deployment dicts + parent_otel_span: Span | None, + ) -> frozenset[str]: + now: Final = self.clock() + counters: Final = tuple( + chain.from_iterable(self._limit_counters(deployment, now) for deployment in deployments) + ) + return frozenset( + [ + model_id + for model_id, key, limit in counters + if _reached(await self._async_get_current_tpm(key, limit, parent_otel_span), limit) + ] + ) + @with_service_target(ROUTER_USAGE_TARGET) def pre_call_check(self, deployment: dict) -> dict | None: """ diff --git a/tests/unit/router_utils/pre_call_checks/test_deployment_affinity_check.py b/tests/unit/router_utils/pre_call_checks/test_deployment_affinity_check.py index 799acc5d5cc..94823593fe2 100644 --- a/tests/unit/router_utils/pre_call_checks/test_deployment_affinity_check.py +++ b/tests/unit/router_utils/pre_call_checks/test_deployment_affinity_check.py @@ -1397,3 +1397,55 @@ async def test_pin_holds_for_a_web_search_request_when_no_untried_deployment_can ) assert [d["model_info"]["id"] for d in filtered] == ["no-search"] + + +@pytest.mark.asyncio +async def test_key_pin_holds_for_a_web_search_request_when_tag_filtering_removes_every_deployment_that_can_search(): + router: Final = litellm.Router( + model_list=[ + { + "model_name": "mixed", + "litellm_params": { + "model": "openai/gpt-x", + "api_key": "fake-key", + "mock_response": "hi", + "tags": ["cheap"], + }, + "model_info": {"id": model_id, "supports_web_search": False}, + } + for model_id in ("flagged-a", "flagged-b") + ] + + [ + { + "model_name": "mixed", + "litellm_params": { + "model": "openai/gpt-y", + "api_key": "fake-key", + "mock_response": "hi", + "tags": ["premium"], + }, + "model_info": {"id": "search"}, + } + ], + optional_pre_call_checks=["deployment_affinity"], + enable_tag_filtering=True, + ) + await router.cache.async_set_cache( + DeploymentAffinityCheck.get_affinity_cache_key(model_group="mixed", user_key="caller-1"), + {"model_id": "flagged-a"}, + ttl=3600, + ) + + served_by: Final = { + ( + await router.acompletion( + model="mixed", + messages=[{"role": "user", "content": "hi"}], + metadata={"user_api_key_hash": "caller-1", "tags": ["cheap"]}, + web_search_options={}, + ) + )._hidden_params["model_id"] + for _ in range(20) + } + + assert served_by == {"flagged-a"} diff --git a/tests/unit/router_utils/test_router_utils_common_utils.py b/tests/unit/router_utils/test_router_utils_common_utils.py index 4e76fff5aad..d9d1c4458b2 100644 --- a/tests/unit/router_utils/test_router_utils_common_utils.py +++ b/tests/unit/router_utils/test_router_utils_common_utils.py @@ -738,9 +738,13 @@ async def test_web_search_preference_gives_way_once_the_capable_deployment_faile @pytest.mark.parametrize( "caller_credentials", [{}, {"api_key": "caller-key"}], ids=["configured-credentials", "caller-credentials"] ) +@pytest.mark.parametrize( + "operator_model_info", [{}, {"original_model_id": "set-by-operator"}], ids=["plain", "operator-original-model-id"] +) async def test_retry_after_a_retryable_failure_on_the_capable_deployment_reaches_the_other_one( web_search_request: dict[str, object], caller_credentials: dict[str, str], + operator_model_info: dict[str, str], ) -> None: router: Final = Router( model_list=[ @@ -752,7 +756,7 @@ async def test_retry_after_a_retryable_failure_on_the_capable_deployment_reaches "mock_response": Exception("connection error"), "weight": 1, }, - "model_info": {"id": "search"}, + "model_info": {"id": "search", **operator_model_info}, }, { "model_name": "mixed", @@ -858,6 +862,107 @@ async def test_retries_past_the_kept_retry_history_reach_the_deployment_that_can assert response.choices[0].message.content == "from no-search" +@pytest.mark.asyncio +async def test_retry_after_a_capable_deployment_fails_goes_to_another_capable_one_before_any_other() -> None: + router: Final = Router( + model_list=[ + { + "model_name": "mixed", + "litellm_params": { + "model": "openai/gpt-x", + "api_key": "fake-key", + "mock_response": Exception("connection error"), + "weight": 10**6, + }, + "model_info": {"id": "dead-search"}, + }, + { + "model_name": "mixed", + "litellm_params": {"model": "openai/gpt-x", "api_key": "fake-key", "mock_response": "from live-search"}, + "model_info": {"id": "live-search"}, + }, + { + "model_name": "mixed", + "litellm_params": {"model": "openai/gpt-y", "api_key": "fake-key", "mock_response": "from no-search"}, + "model_info": {"id": "no-search", "supports_web_search": False}, + }, + ], + num_retries=1, + disable_cooldowns=True, + ) + replies: Final = tuple( + [ + ( + await router.acompletion( + model="mixed", messages=[{"role": "user", "content": "hi"}], **_ANTHROPIC_WEB_SEARCH_REQUEST + ) + ) + .choices[0] + .message.content + for _ in range(5) + ] + ) + assert replies == ("from live-search",) * 5 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("web_search_options", [False, "", 0, []], ids=["false", "empty-string", "zero", "list"]) +async def test_web_search_options_that_is_not_an_object_routes_like_a_plain_request( + web_search_options: object, +) -> None: + router: Final = _router_with_a_mixed_web_search_group(False) + picks: Final = { + ( + await router.async_get_available_deployment( + model="mixed", + messages=[{"role": "user", "content": "hi"}], + request_kwargs={"web_search_options": web_search_options}, + ) + )["model_info"]["id"] + for _ in range(40) + } + assert picks == {"search", "no-search"} + + +def _web_search_group_with_model_rate_limits() -> Router: + return Router( + model_list=[ + { + "model_name": "mixed", + "litellm_params": {"model": "openai/gpt-x", "api_key": "fake-key", "mock_response": "hi", "rpm": 1}, + "model_info": {"id": "search"}, + }, + { + "model_name": "mixed", + "litellm_params": { + "model": "openai/gpt-y", + "api_key": "fake-key", + "mock_response": "hi", + "rpm": 10**6, + }, + "model_info": {"id": "no-search", "supports_web_search": False}, + }, + ], + optional_pre_call_checks=["enforce_model_rate_limits"], + ) + + +@pytest.mark.asyncio +async def test_web_search_request_goes_to_the_other_deployment_once_the_capable_one_hits_its_enforced_rpm() -> None: + router: Final = _web_search_group_with_model_rate_limits() + served_by: Final = tuple( + [ + ( + await router.acompletion( + model="mixed", messages=[{"role": "user", "content": "hi"}], **_ANTHROPIC_WEB_SEARCH_REQUEST + ) + )._hidden_params["model_id"] + for _ in range(3) + ] + ) + assert served_by == ("search", "no-search", "no-search") + + def _rate_limited_web_search_group(routing_strategy: str) -> Router: return Router( model_list=[ diff --git a/tests/unit/test_router/test_enforce_model_rate_limits.py b/tests/unit/test_router/test_enforce_model_rate_limits.py index ff3ffc0db0c..c152b9d4be8 100644 --- a/tests/unit/test_router/test_enforce_model_rate_limits.py +++ b/tests/unit/test_router/test_enforce_model_rate_limits.py @@ -539,6 +539,29 @@ class TestDailyLimits: assert "TPD limit=100, current usage=120" in str(exc_info.value) + @pytest.mark.asyncio + async def test_deployments_at_their_rpm_or_rpd_limit_are_reported_and_ones_below_are_not(self): + minute_spent = { + "litellm_params": {"model": "gpt-4.1-mini", "rpm": 1}, + "model_info": {"id": "minute-test-id"}, + "model_name": "daily-model", + } + below_limits = { + "litellm_params": {"model": "gpt-4.1-mini", "rpm": 2, "rpd": 2}, + "model_info": {"id": "below-test-id"}, + "model_name": "daily-model", + } + check = ModelRateLimitingCheck(dual_cache=DualCache(), clock=Clock(DAY_ONE)) + + await _async_pre_call(check, RPD_DEPLOYMENT) + await _async_pre_call(check, RPD_DEPLOYMENT) + await _async_pre_call(check, minute_spent) + await _async_pre_call(check, below_limits) + + assert await check.async_deployment_ids_at_limit( + [RPD_DEPLOYMENT, minute_spent, below_limits], None + ) == frozenset({"daily-test-id", "minute-test-id"}) + def test_pre_call_check_minute_rejected_request_does_not_spend_rpd(self): clock = Clock(DAY_ONE) check = ModelRateLimitingCheck(dual_cache=DualCache(), clock=clock) From 386b4206b4cfa153a06dc70376839fd05e9c77f2 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 30 Oct 2026 00:55:11 -0700 Subject: [PATCH 13/24] perf(router): read web search deployments' rate limit counters concurrently --- .../pre_call_checks/model_rate_limit_check.py | 10 +++--- .../test_enforce_model_rate_limits.py | 34 +++++++++++++++++++ 2 files changed, 39 insertions(+), 5 deletions(-) diff --git a/litellm/router_utils/pre_call_checks/model_rate_limit_check.py b/litellm/router_utils/pre_call_checks/model_rate_limit_check.py index 8cc47627d99..596483fd258 100644 --- a/litellm/router_utils/pre_call_checks/model_rate_limit_check.py +++ b/litellm/router_utils/pre_call_checks/model_rate_limit_check.py @@ -11,6 +11,7 @@ When a deployment sets both itpm/otpm and tpm/rpm, both are enforced. A warning is logged the first time such a deployment is seen. """ +import asyncio import contextlib from collections.abc import Callable, Mapping, Sequence from datetime import datetime, timedelta @@ -342,12 +343,11 @@ class ModelRateLimitingCheck(CustomLogger): counters: Final = tuple( chain.from_iterable(self._limit_counters(deployment, now) for deployment in deployments) ) + usages: Final = await asyncio.gather( + *(self._async_get_current_tpm(key, limit, parent_otel_span) for _, key, limit in counters) + ) return frozenset( - [ - model_id - for model_id, key, limit in counters - if _reached(await self._async_get_current_tpm(key, limit, parent_otel_span), limit) - ] + model_id for (model_id, _, limit), usage in zip(counters, usages, strict=True) if _reached(usage, limit) ) @with_service_target(ROUTER_USAGE_TARGET) diff --git a/tests/unit/test_router/test_enforce_model_rate_limits.py b/tests/unit/test_router/test_enforce_model_rate_limits.py index c152b9d4be8..8b794ace01b 100644 --- a/tests/unit/test_router/test_enforce_model_rate_limits.py +++ b/tests/unit/test_router/test_enforce_model_rate_limits.py @@ -483,6 +483,21 @@ class Clock: return self.now +class _ReadsThatFinishOnceAllStarted: + def __init__(self, expected: int, usage: int): + self.expected = expected + self.usage = usage + self.started = 0 + self.all_started = asyncio.Event() + + async def async_get_cache(self, key: str, parent_otel_span: object = None) -> int: + self.started += 1 + if self.started == self.expected: + self.all_started.set() + await self.all_started.wait() + return self.usage + + def _pre_call(check: ModelRateLimitingCheck, deployment: dict[str, object]) -> object: return check.pre_call_check(deployment) @@ -562,6 +577,25 @@ class TestDailyLimits: [RPD_DEPLOYMENT, minute_spent, below_limits], None ) == frozenset({"daily-test-id", "minute-test-id"}) + @pytest.mark.asyncio + async def test_deployment_limit_reads_wait_on_redis_together_rather_than_one_by_one(self): + deployments = [ + { + "litellm_params": {"model": "gpt-4.1-mini", "rpm": 5}, + "model_info": {"id": f"concurrent-{index}"}, + "model_name": "daily-model", + } + for index in range(3) + ] + reads = _ReadsThatFinishOnceAllStarted(expected=len(deployments), usage=5) + redis_cache = MagicMock() + redis_cache.async_get_cache = reads.async_get_cache + check = ModelRateLimitingCheck(dual_cache=DualCache(redis_cache=redis_cache), clock=Clock(DAY_ONE)) + + at_limit = await asyncio.wait_for(check.async_deployment_ids_at_limit(deployments, None), timeout=5) + + assert at_limit == frozenset({"concurrent-0", "concurrent-1", "concurrent-2"}) + def test_pre_call_check_minute_rejected_request_does_not_spend_rpd(self): clock = Clock(DAY_ONE) check = ModelRateLimitingCheck(dual_cache=DualCache(), clock=clock) From 1ab8a4e3c012a52698591e246e148a21722e32be Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 30 Oct 2026 02:11:41 -0700 Subject: [PATCH 14/24] fix(router): treat a zero request limit as reached and check enforced limits on the sync web search pick --- litellm/router.py | 12 +++- litellm/router_utils/common_utils.py | 3 +- .../pre_call_checks/model_rate_limit_check.py | 60 +++++++++++++++---- .../test_router_utils_common_utils.py | 44 ++++++++++++++ .../test_enforce_model_rate_limits.py | 50 ++++++++++++++++ 5 files changed, 154 insertions(+), 15 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index c6171ee7d2e..530722046cd 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -16862,12 +16862,21 @@ class Router(DeprecatedRouterMembers): ), ) + web_search_capable: Final = web_search_capable_deployments( + request_kwargs=routing_kwargs, + deployments=cast(list[dict[str, object]], healthy_deployments), # cast-ok: str-keyed dicts + ) + at_rate_limit: Final = frozenset( + chain.from_iterable( + check.deployment_ids_at_limit(web_search_capable) for check in self._model_rate_limit_checks() + ) + ) if strategy == "simple-shuffle": # if users pass rpm or tpm, we do a random weighted pick - based on rpm/tpm ############## Check 'weight' param set for weighted pick ################# shuffled: Final = simple_shuffle( resolve_model_alias=self._core_get_model_from_alias, - healthy_deployments=prefer_web_search_deployments(healthy_deployments, routing_kwargs), + healthy_deployments=prefer_web_search_deployments(healthy_deployments, routing_kwargs, at_rate_limit), model=model, request_kwargs=routing_kwargs, ) @@ -16885,6 +16894,7 @@ class Router(DeprecatedRouterMembers): ), cast(list[dict[str, object]], healthy_deployments), # cast-ok: router deployments are plain dicts routing_kwargs, + at_rate_limit, ) if deployment is None: diff --git a/litellm/router_utils/common_utils.py b/litellm/router_utils/common_utils.py index 775cd17afcb..49463fd4922 100644 --- a/litellm/router_utils/common_utils.py +++ b/litellm/router_utils/common_utils.py @@ -585,8 +585,9 @@ def pick_preferring_web_search( pick: Callable[[list[_DeploymentT]], _Picked | None], # mutable-ok: router deployment-list contract deployments: list[_DeploymentT], # mutable-ok: router deployment-list contract request_kwargs: Mapping[str, object] | None, + at_rate_limit: frozenset[str] = frozenset(), ) -> _Picked | None: - preferred: Final = prefer_web_search_deployments(deployments, request_kwargs) + preferred: Final = prefer_web_search_deployments(deployments, request_kwargs, at_rate_limit) if len(preferred) == len(deployments): return pick(deployments) try: diff --git a/litellm/router_utils/pre_call_checks/model_rate_limit_check.py b/litellm/router_utils/pre_call_checks/model_rate_limit_check.py index 596483fd258..ecb8b9a5937 100644 --- a/litellm/router_utils/pre_call_checks/model_rate_limit_check.py +++ b/litellm/router_utils/pre_call_checks/model_rate_limit_check.py @@ -14,6 +14,7 @@ is logged the first time such a deployment is seen. import asyncio import contextlib from collections.abc import Callable, Mapping, Sequence +from dataclasses import dataclass from datetime import datetime, timedelta from itertools import chain from typing import TYPE_CHECKING, Any, Final, Literal @@ -124,12 +125,25 @@ def _deployment_model_id(deployment: Mapping[str, object]) -> object: return model_info.get("id") if isinstance(model_info, Mapping) else None -def _counters_with_limits(model_id: str, counters: Sequence[tuple[str, object]]) -> tuple[tuple[str, str, int], ...]: - return tuple((model_id, key, limit) for key, limit in counters if type(limit) is int) +@dataclass(frozen=True, slots=True) +class _LimitCounter: + model_id: str + key: str + limit: int + counts_requests: bool + + def reached(self, usage: object) -> bool: + if isinstance(usage, (int, float)): + return usage >= self.limit + return self.counts_requests and self.limit <= 0 -def _reached(current: object, limit: int) -> bool: - return isinstance(current, (int, float)) and current >= limit +def _counters_with_limits(model_id: str, counters: Sequence[tuple[str, object, bool]]) -> tuple[_LimitCounter, ...]: + return tuple( + _LimitCounter(model_id, key, limit, counts_requests) + for key, limit, counts_requests in counters + if type(limit) is int + ) def _reset_offset_from_kwargs(kwargs: Mapping[str, object]) -> timedelta: @@ -318,7 +332,7 @@ class ModelRateLimitingCheck(CustomLogger): self, deployment: dict[str, object], # mutable-ok: the limit and key readers take router deployment dicts now: datetime, - ) -> tuple[tuple[str, str, int], ...]: + ) -> tuple[_LimitCounter, ...]: model_id: Final = _deployment_model_id(deployment) if not isinstance(model_id, str): return () @@ -327,11 +341,34 @@ class ModelRateLimitingCheck(CustomLogger): try: daily: Final = _DailyLimitDeployment.model_validate(deployment) except ValidationError: - return _counters_with_limits(model_id, ((tpm_key, tpm_limit), (rpm_key, rpm_limit))) + return _counters_with_limits(model_id, ((tpm_key, tpm_limit, False), (rpm_key, rpm_limit, True))) tpd_limit, rpd_limit = daily.limits() tpd_key, rpd_key = daily.cache_keys(now) return _counters_with_limits( - model_id, ((tpm_key, tpm_limit), (rpm_key, rpm_limit), (tpd_key, tpd_limit), (rpd_key, rpd_limit)) + model_id, + ( + (tpm_key, tpm_limit, False), + (rpm_key, rpm_limit, True), + (tpd_key, tpd_limit, False), + (rpd_key, rpd_limit, True), + ), + ) + + def _counters_of( + self, + deployments: Sequence[dict[str, object]], # mutable-ok: _limit_counters reads router deployment dicts + ) -> tuple[_LimitCounter, ...]: + now: Final = self.clock() + return tuple(chain.from_iterable(self._limit_counters(deployment, now) for deployment in deployments)) + + def deployment_ids_at_limit( + self, + deployments: Sequence[dict[str, object]], # mutable-ok: _limit_counters reads router deployment dicts + ) -> frozenset[str]: + return frozenset( + counter.model_id + for counter in self._counters_of(deployments) + if counter.reached(self._get_current_tpm(counter.key, counter.limit)) ) async def async_deployment_ids_at_limit( @@ -339,15 +376,12 @@ class ModelRateLimitingCheck(CustomLogger): deployments: Sequence[dict[str, object]], # mutable-ok: _limit_counters reads router deployment dicts parent_otel_span: Span | None, ) -> frozenset[str]: - now: Final = self.clock() - counters: Final = tuple( - chain.from_iterable(self._limit_counters(deployment, now) for deployment in deployments) - ) + counters: Final = self._counters_of(deployments) usages: Final = await asyncio.gather( - *(self._async_get_current_tpm(key, limit, parent_otel_span) for _, key, limit in counters) + *(self._async_get_current_tpm(counter.key, counter.limit, parent_otel_span) for counter in counters) ) return frozenset( - model_id for (model_id, _, limit), usage in zip(counters, usages, strict=True) if _reached(usage, limit) + counter.model_id for counter, usage in zip(counters, usages, strict=True) if counter.reached(usage) ) @with_service_target(ROUTER_USAGE_TARGET) diff --git a/tests/unit/router_utils/test_router_utils_common_utils.py b/tests/unit/router_utils/test_router_utils_common_utils.py index d9d1c4458b2..3e67725074c 100644 --- a/tests/unit/router_utils/test_router_utils_common_utils.py +++ b/tests/unit/router_utils/test_router_utils_common_utils.py @@ -963,6 +963,50 @@ async def test_web_search_request_goes_to_the_other_deployment_once_the_capable_ assert served_by == ("search", "no-search", "no-search") +def test_sync_web_search_request_goes_to_the_other_deployment_once_the_capable_one_hits_its_enforced_rpm() -> None: + router: Final = _web_search_group_with_model_rate_limits() + served_by: Final = tuple( + router.completion( + model="mixed", messages=[{"role": "user", "content": "hi"}], num_retries=0, **_ANTHROPIC_WEB_SEARCH_REQUEST + )._hidden_params["model_id"] + for _ in range(3) + ) + assert served_by == ("search", "no-search", "no-search") + + +@pytest.mark.asyncio +async def test_web_search_request_never_goes_to_a_capable_deployment_whose_enforced_rpm_is_zero() -> None: + router: Final = Router( + model_list=[ + { + "model_name": "mixed", + "litellm_params": {"model": "openai/gpt-x", "api_key": "fake-key", "mock_response": "hi", "rpm": 0}, + "model_info": {"id": "search"}, + }, + { + "model_name": "mixed", + "litellm_params": {"model": "openai/gpt-y", "api_key": "fake-key", "mock_response": "hi", "rpm": 1000}, + "model_info": {"id": "no-search", "supports_web_search": False}, + }, + ], + optional_pre_call_checks=["enforce_model_rate_limits"], + ) + served_by: Final = tuple( + [ + ( + await router.acompletion( + model="mixed", + messages=[{"role": "user", "content": "hi"}], + num_retries=0, + **_ANTHROPIC_WEB_SEARCH_REQUEST, + ) + )._hidden_params["model_id"] + for _ in range(5) + ] + ) + assert served_by == ("no-search",) * 5 + + def _rate_limited_web_search_group(routing_strategy: str) -> Router: return Router( model_list=[ diff --git a/tests/unit/test_router/test_enforce_model_rate_limits.py b/tests/unit/test_router/test_enforce_model_rate_limits.py index 8b794ace01b..706a3ebe01e 100644 --- a/tests/unit/test_router/test_enforce_model_rate_limits.py +++ b/tests/unit/test_router/test_enforce_model_rate_limits.py @@ -506,6 +506,33 @@ async def _async_pre_call(check: ModelRateLimitingCheck, deployment: dict[str, o return await check.async_pre_call_check(deployment) +def _rejects(check: ModelRateLimitingCheck, deployment: dict[str, object]) -> bool: + try: + _pre_call(check, deployment) + except litellm.RateLimitError: + return True + return False + + +async def _async_rejects(check: ModelRateLimitingCheck, deployment: dict[str, object]) -> bool: + try: + await _async_pre_call(check, deployment) + except litellm.RateLimitError: + return True + return False + + +_FRESH_LIMITS: Final = ({"rpm": 0}, {"rpd": 0}, {"tpm": 0}, {"tpd": 0}, {"rpm": 1, "tpm": 1}) + + +def _deployment_with(limits: dict[str, int]) -> dict[str, object]: + return { + "litellm_params": {"model": "gpt-4.1-mini", **limits}, + "model_info": {"id": "limit-test-id"}, + "model_name": "daily-model", + } + + async def _async_log_success(check: ModelRateLimitingCheck, kwargs: dict[str, object]) -> None: await check.async_log_success_event(kwargs, None, None, None) @@ -596,6 +623,29 @@ class TestDailyLimits: assert at_limit == frozenset({"concurrent-0", "concurrent-1", "concurrent-2"}) + @pytest.mark.parametrize("limits", _FRESH_LIMITS, ids=lambda limits: ",".join(limits)) + def test_a_deployment_is_reported_at_its_limit_exactly_when_its_next_request_is_rejected( + self, limits: dict[str, int] + ): + deployment = _deployment_with(limits) + check = ModelRateLimitingCheck(dual_cache=DualCache(), clock=Clock(DAY_ONE)) + + reported = "limit-test-id" in check.deployment_ids_at_limit([deployment]) + + assert reported == _rejects(check, deployment) + + @pytest.mark.asyncio + @pytest.mark.parametrize("limits", _FRESH_LIMITS, ids=lambda limits: ",".join(limits)) + async def test_async_a_deployment_is_reported_at_its_limit_exactly_when_its_next_request_is_rejected( + self, limits: dict[str, int] + ): + deployment = _deployment_with(limits) + check = ModelRateLimitingCheck(dual_cache=DualCache(), clock=Clock(DAY_ONE)) + + reported = "limit-test-id" in await check.async_deployment_ids_at_limit([deployment], None) + + assert reported == await _async_rejects(check, deployment) + def test_pre_call_check_minute_rejected_request_does_not_spend_rpd(self): clock = Clock(DAY_ONE) check = ModelRateLimitingCheck(dual_cache=DualCache(), clock=clock) From e7420e12db54328e8f661aabb86ccf222966688d Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 30 Oct 2026 03:12:03 -0700 Subject: [PATCH 15/24] fix(router): read float and bool rate limits at the web search pick the way enforcement does --- .../pre_call_checks/model_rate_limit_check.py | 8 ++--- .../test_router_utils_common_utils.py | 36 +++++++++++++++---- .../test_enforce_model_rate_limits.py | 13 +++---- 3 files changed, 40 insertions(+), 17 deletions(-) diff --git a/litellm/router_utils/pre_call_checks/model_rate_limit_check.py b/litellm/router_utils/pre_call_checks/model_rate_limit_check.py index ecb8b9a5937..bf556386a7f 100644 --- a/litellm/router_utils/pre_call_checks/model_rate_limit_check.py +++ b/litellm/router_utils/pre_call_checks/model_rate_limit_check.py @@ -129,7 +129,7 @@ def _deployment_model_id(deployment: Mapping[str, object]) -> object: class _LimitCounter: model_id: str key: str - limit: int + limit: float counts_requests: bool def reached(self, usage: object) -> bool: @@ -142,7 +142,7 @@ def _counters_with_limits(model_id: str, counters: Sequence[tuple[str, object, b return tuple( _LimitCounter(model_id, key, limit, counts_requests) for key, limit, counts_requests in counters - if type(limit) is int + if isinstance(limit, (int, float)) ) @@ -307,7 +307,7 @@ class ModelRateLimitingCheck(CustomLogger): return tpm_key, rpm_key @with_service_target(ROUTER_USAGE_TARGET) - def _get_current_tpm(self, tpm_key: str, tpm_limit: int) -> int | None: + def _get_current_tpm(self, tpm_key: str, tpm_limit: float) -> int | None: local_tpm: Final = self.dual_cache.get_cache(key=tpm_key, local_only=True) redis_cache: Final = self.dual_cache.redis_cache if redis_cache is None or (local_tpm is not None and local_tpm >= tpm_limit): @@ -318,7 +318,7 @@ class ModelRateLimitingCheck(CustomLogger): return local_tpm @with_service_target(ROUTER_USAGE_TARGET) - async def _async_get_current_tpm(self, tpm_key: str, tpm_limit: int, parent_otel_span: Span | None) -> int | None: + async def _async_get_current_tpm(self, tpm_key: str, tpm_limit: float, parent_otel_span: Span | None) -> int | None: local_tpm: Final = await self.dual_cache.async_get_cache(key=tpm_key, local_only=True) redis_cache: Final = self.dual_cache.redis_cache if redis_cache is None or (local_tpm is not None and local_tpm >= tpm_limit): diff --git a/tests/unit/router_utils/test_router_utils_common_utils.py b/tests/unit/router_utils/test_router_utils_common_utils.py index 3e67725074c..b3394b98590 100644 --- a/tests/unit/router_utils/test_router_utils_common_utils.py +++ b/tests/unit/router_utils/test_router_utils_common_utils.py @@ -1,6 +1,8 @@ import asyncio import logging +from collections.abc import Mapping from itertools import chain +from types import MappingProxyType from typing import Dict, Final, List, Optional, Union from unittest.mock import Mock @@ -924,13 +926,27 @@ async def test_web_search_options_that_is_not_an_object_routes_like_a_plain_requ assert picks == {"search", "no-search"} -def _web_search_group_with_model_rate_limits() -> Router: +_SEARCH_RPM_OF_ONE: Final = ( + pytest.param({"litellm_params": {"rpm": 1}}, id="int-rpm-in-litellm-params"), + pytest.param({"model_info": {"rpm": 1.0}}, id="float-rpm-in-model-info"), + pytest.param({"model_info": {"rpm": True}}, id="bool-rpm-in-model-info"), +) + + +def _web_search_group_with_model_rate_limits( + search_rpm: Mapping[str, Mapping[str, object]] = MappingProxyType({"litellm_params": {"rpm": 1}}), +) -> Router: return Router( model_list=[ { "model_name": "mixed", - "litellm_params": {"model": "openai/gpt-x", "api_key": "fake-key", "mock_response": "hi", "rpm": 1}, - "model_info": {"id": "search"}, + "litellm_params": { + "model": "openai/gpt-x", + "api_key": "fake-key", + "mock_response": "hi", + **search_rpm.get("litellm_params", {}), + }, + "model_info": {"id": "search", **search_rpm.get("model_info", {})}, }, { "model_name": "mixed", @@ -948,8 +964,11 @@ def _web_search_group_with_model_rate_limits() -> Router: @pytest.mark.asyncio -async def test_web_search_request_goes_to_the_other_deployment_once_the_capable_one_hits_its_enforced_rpm() -> None: - router: Final = _web_search_group_with_model_rate_limits() +@pytest.mark.parametrize("search_rpm", _SEARCH_RPM_OF_ONE) +async def test_web_search_request_goes_to_the_other_deployment_once_the_capable_one_hits_its_enforced_rpm( + search_rpm: Mapping[str, Mapping[str, object]], +) -> None: + router: Final = _web_search_group_with_model_rate_limits(search_rpm) served_by: Final = tuple( [ ( @@ -963,8 +982,11 @@ async def test_web_search_request_goes_to_the_other_deployment_once_the_capable_ assert served_by == ("search", "no-search", "no-search") -def test_sync_web_search_request_goes_to_the_other_deployment_once_the_capable_one_hits_its_enforced_rpm() -> None: - router: Final = _web_search_group_with_model_rate_limits() +@pytest.mark.parametrize("search_rpm", _SEARCH_RPM_OF_ONE) +def test_sync_web_search_request_goes_to_the_other_deployment_once_the_capable_one_hits_its_enforced_rpm( + search_rpm: Mapping[str, Mapping[str, object]], +) -> None: + router: Final = _web_search_group_with_model_rate_limits(search_rpm) served_by: Final = tuple( router.completion( model="mixed", messages=[{"role": "user", "content": "hi"}], num_retries=0, **_ANTHROPIC_WEB_SEARCH_REQUEST diff --git a/tests/unit/test_router/test_enforce_model_rate_limits.py b/tests/unit/test_router/test_enforce_model_rate_limits.py index 706a3ebe01e..01af8c43598 100644 --- a/tests/unit/test_router/test_enforce_model_rate_limits.py +++ b/tests/unit/test_router/test_enforce_model_rate_limits.py @@ -6,6 +6,7 @@ regardless of the routing strategy being used. """ import asyncio +from collections.abc import Mapping from datetime import datetime, timedelta, timezone from typing import Final from unittest.mock import AsyncMock, MagicMock @@ -522,10 +523,10 @@ async def _async_rejects(check: ModelRateLimitingCheck, deployment: dict[str, ob return False -_FRESH_LIMITS: Final = ({"rpm": 0}, {"rpd": 0}, {"tpm": 0}, {"tpd": 0}, {"rpm": 1, "tpm": 1}) +_FRESH_LIMITS: Final = ({"rpm": 0}, {"rpm": 0.0}, {"rpm": False}, {"rpd": 0}, {"tpm": 0}, {"tpd": 0}, {"rpm": 1, "tpm": 1}) -def _deployment_with(limits: dict[str, int]) -> dict[str, object]: +def _deployment_with(limits: Mapping[str, float]) -> dict[str, object]: return { "litellm_params": {"model": "gpt-4.1-mini", **limits}, "model_info": {"id": "limit-test-id"}, @@ -623,9 +624,9 @@ class TestDailyLimits: assert at_limit == frozenset({"concurrent-0", "concurrent-1", "concurrent-2"}) - @pytest.mark.parametrize("limits", _FRESH_LIMITS, ids=lambda limits: ",".join(limits)) + @pytest.mark.parametrize("limits", _FRESH_LIMITS, ids=lambda limits: ",".join(f"{kind}={value!r}" for kind, value in limits.items())) def test_a_deployment_is_reported_at_its_limit_exactly_when_its_next_request_is_rejected( - self, limits: dict[str, int] + self, limits: Mapping[str, float] ): deployment = _deployment_with(limits) check = ModelRateLimitingCheck(dual_cache=DualCache(), clock=Clock(DAY_ONE)) @@ -635,9 +636,9 @@ class TestDailyLimits: assert reported == _rejects(check, deployment) @pytest.mark.asyncio - @pytest.mark.parametrize("limits", _FRESH_LIMITS, ids=lambda limits: ",".join(limits)) + @pytest.mark.parametrize("limits", _FRESH_LIMITS, ids=lambda limits: ",".join(f"{kind}={value!r}" for kind, value in limits.items())) async def test_async_a_deployment_is_reported_at_its_limit_exactly_when_its_next_request_is_rejected( - self, limits: dict[str, int] + self, limits: Mapping[str, float] ): deployment = _deployment_with(limits) check = ModelRateLimitingCheck(dual_cache=DualCache(), clock=Clock(DAY_ONE)) From 146e4c506b43e2a53269aeb3cf46e8d354b03af3 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 30 Oct 2026 04:18:13 -0700 Subject: [PATCH 16/24] perf(router): read sync search pick rate limit counters from Redis in one round trip --- .../pre_call_checks/model_rate_limit_check.py | 31 +++++++++-- .../test_enforce_model_rate_limits.py | 51 ++++++++++++++++++- 2 files changed, 78 insertions(+), 4 deletions(-) diff --git a/litellm/router_utils/pre_call_checks/model_rate_limit_check.py b/litellm/router_utils/pre_call_checks/model_rate_limit_check.py index bf556386a7f..303c88ba70a 100644 --- a/litellm/router_utils/pre_call_checks/model_rate_limit_check.py +++ b/litellm/router_utils/pre_call_checks/model_rate_limit_check.py @@ -361,14 +361,39 @@ class ModelRateLimitingCheck(CustomLogger): now: Final = self.clock() return tuple(chain.from_iterable(self._limit_counters(deployment, now) for deployment in deployments)) + def _local_usage(self, key: str) -> object: + return self.dual_cache.get_cache(key=key, local_only=True) + + def _redis_usages(self, keys: list[str]) -> Mapping[str, object] | None: + redis_cache: Final = self.dual_cache.redis_cache + if redis_cache is None or not keys: + return None + try: + return redis_cache.batch_get_cache(key_list=keys) + except RedisCircuitBreakerOpenError: + return None + + @with_service_target(ROUTER_USAGE_TARGET) + def _current_usages(self, counters: Sequence[_LimitCounter]) -> tuple[object, ...]: + local_usages: Final = tuple(self._local_usage(counter.key) for counter in counters) + redis_usages: Final = self._redis_usages( + [counter.key for counter, usage in zip(counters, local_usages, strict=True) if not counter.reached(usage)] + ) + if redis_usages is None: + return local_usages + return tuple( + usage if counter.reached(usage) else redis_usages.get(counter.key) + for counter, usage in zip(counters, local_usages, strict=True) + ) + def deployment_ids_at_limit( self, deployments: Sequence[dict[str, object]], # mutable-ok: _limit_counters reads router deployment dicts ) -> frozenset[str]: + counters: Final = self._counters_of(deployments) + usages: Final = self._current_usages(counters) return frozenset( - counter.model_id - for counter in self._counters_of(deployments) - if counter.reached(self._get_current_tpm(counter.key, counter.limit)) + counter.model_id for counter, usage in zip(counters, usages, strict=True) if counter.reached(usage) ) async def async_deployment_ids_at_limit( diff --git a/tests/unit/test_router/test_enforce_model_rate_limits.py b/tests/unit/test_router/test_enforce_model_rate_limits.py index 01af8c43598..28ea12bfaaf 100644 --- a/tests/unit/test_router/test_enforce_model_rate_limits.py +++ b/tests/unit/test_router/test_enforce_model_rate_limits.py @@ -6,7 +6,7 @@ regardless of the routing strategy being used. """ import asyncio -from collections.abc import Mapping +from collections.abc import Callable, Mapping from datetime import datetime, timedelta, timezone from typing import Final from unittest.mock import AsyncMock, MagicMock @@ -526,6 +526,13 @@ async def _async_rejects(check: ModelRateLimitingCheck, deployment: dict[str, ob _FRESH_LIMITS: Final = ({"rpm": 0}, {"rpm": 0.0}, {"rpm": False}, {"rpd": 0}, {"tpm": 0}, {"tpd": 0}, {"rpm": 1, "tpm": 1}) +def _redis_batch_answering(usage: int) -> Callable[[list[str]], dict[str, int]]: + def batch_get_cache(key_list: list[str]) -> dict[str, int]: + return dict.fromkeys(key_list, usage) + + return batch_get_cache + + def _deployment_with(limits: Mapping[str, float]) -> dict[str, object]: return { "litellm_params": {"model": "gpt-4.1-mini", **limits}, @@ -624,6 +631,48 @@ class TestDailyLimits: assert at_limit == frozenset({"concurrent-0", "concurrent-1", "concurrent-2"}) + def test_sync_deployment_limit_reads_reach_redis_in_one_round_trip(self): + deployments: list[dict[str, object]] = [ + { + "litellm_params": {"model": "gpt-4.1-mini", "rpm": 5, "rpd": 50}, + "model_info": {"id": f"batched-{index}"}, + "model_name": "daily-model", + } + for index in range(3) + ] + redis_cache = MagicMock() + redis_cache.batch_get_cache.side_effect = _redis_batch_answering(5) + check = ModelRateLimitingCheck(dual_cache=DualCache(redis_cache=redis_cache), clock=Clock(DAY_ONE)) + + at_limit = check.deployment_ids_at_limit(deployments) + + assert at_limit == frozenset({"batched-0", "batched-1", "batched-2"}) + assert redis_cache.batch_get_cache.call_count == 1 + assert len(redis_cache.batch_get_cache.call_args.kwargs["key_list"]) == 6 + + @pytest.mark.parametrize( + "redis_batch_get", + [_redis_batch_answering(0), RedisCircuitBreakerOpenError()], + ids=["redis-behind-this-worker", "redis-circuit-open"], + ) + def test_sync_a_deployment_this_worker_counted_to_its_limit_stays_at_its_limit( + self, redis_batch_get: Callable[[list[str]], dict[str, int]] | Exception + ): + spent: dict[str, object] = { + "litellm_params": {"model": "gpt-4.1-mini", "rpm": 1}, + "model_info": {"id": "spent-test-id"}, + "model_name": "daily-model", + } + redis_cache = MagicMock() + redis_cache.batch_get_cache.side_effect = redis_batch_get + dual_cache = DualCache(redis_cache=redis_cache) + check = ModelRateLimitingCheck(dual_cache=dual_cache, clock=Clock(DAY_ONE)) + _, rpm_key = check._get_cache_keys(spent, DAY_ONE.strftime("%H-%M")) + dual_cache.set_cache(key=rpm_key, value=1, local_only=True) + + assert check.deployment_ids_at_limit([spent, _deployment_with({"rpd": 2})]) == frozenset({"spent-test-id"}) + assert rpm_key not in redis_cache.batch_get_cache.call_args.kwargs["key_list"] + @pytest.mark.parametrize("limits", _FRESH_LIMITS, ids=lambda limits: ",".join(f"{kind}={value!r}" for kind, value in limits.items())) def test_a_deployment_is_reported_at_its_limit_exactly_when_its_next_request_is_rejected( self, limits: Mapping[str, float] From 0033626a0977ae057cef641c71a96df2b5b96900 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 30 Oct 2026 07:43:21 -0700 Subject: [PATCH 17/24] fix(router): count itpm and otpm limits when the web search pick skips deployments at their limit --- .../io_token_rate_limit_check.py | 10 ++++++ .../pre_call_checks/model_rate_limit_check.py | 31 ++++++++--------- .../test_router_utils_common_utils.py | 7 ++-- .../test_enforce_model_rate_limits.py | 33 ++++++++++++++++++- 4 files changed, 63 insertions(+), 18 deletions(-) diff --git a/litellm/router_utils/pre_call_checks/io_token_rate_limit_check.py b/litellm/router_utils/pre_call_checks/io_token_rate_limit_check.py index 7580c7dadd3..6531555855f 100644 --- a/litellm/router_utils/pre_call_checks/io_token_rate_limit_check.py +++ b/litellm/router_utils/pre_call_checks/io_token_rate_limit_check.py @@ -129,6 +129,16 @@ def _get_cache_keys(deployment: dict, current_minute: str) -> tuple[str, str] | return itpm_key, otpm_key +def io_token_limits_by_cache_key( + deployment: dict[str, object], # mutable-ok: the io token readers take router deployment dicts + current_minute: str, +) -> tuple[tuple[str, int | None], ...]: + cache_keys: Final = _get_cache_keys(deployment, current_minute) + if cache_keys is None: + return () + return tuple(zip(cache_keys, get_deployment_io_token_limits(deployment), strict=True)) + + def _served_model(deployment: object) -> str: params: Final = deployment.get("litellm_params") if is_str_keyed_mapping(deployment) else None served: Final = params.get("model") if is_str_keyed_mapping(params) else None diff --git a/litellm/router_utils/pre_call_checks/model_rate_limit_check.py b/litellm/router_utils/pre_call_checks/model_rate_limit_check.py index eb2dacbc67a..4dd98895e94 100644 --- a/litellm/router_utils/pre_call_checks/model_rate_limit_check.py +++ b/litellm/router_utils/pre_call_checks/model_rate_limit_check.py @@ -34,6 +34,7 @@ from litellm.router_utils.pre_call_checks.io_token_rate_limit_check import ( async_io_token_refund_failure, deployment_has_io_token_limits, get_io_token_rate_limit_request_kwargs, + io_token_limits_by_cache_key, io_token_pre_call_check, io_token_reconcile_success, io_token_refund_failure, @@ -129,18 +130,18 @@ class _LimitCounter: model_id: str key: str limit: float - counts_requests: bool + increments_first: bool def reached(self, usage: object) -> bool: - if isinstance(usage, (int, float)): - return usage >= self.limit - return self.counts_requests and self.limit <= 0 + if self.increments_first: + return (usage if isinstance(usage, (int, float)) else 0) + 1 > self.limit + return isinstance(usage, (int, float)) and usage >= self.limit def _counters_with_limits(model_id: str, counters: Sequence[tuple[str, object, bool]]) -> tuple[_LimitCounter, ...]: return tuple( - _LimitCounter(model_id, key, limit, counts_requests) - for key, limit, counts_requests in counters + _LimitCounter(model_id, key, limit, increments_first) + for key, limit, increments_first in counters if isinstance(limit, (int, float)) ) @@ -336,21 +337,21 @@ class ModelRateLimitingCheck(CustomLogger): if not isinstance(model_id, str): return () tpm_limit, rpm_limit = self._get_deployment_limits(deployment) - tpm_key, rpm_key = self._get_cache_keys(deployment, now.strftime("%H-%M")) + current_minute: Final = now.strftime("%H-%M") + tpm_key, rpm_key = self._get_cache_keys(deployment, current_minute) + minute_counters: Final = ( + (tpm_key, tpm_limit, False), + (rpm_key, rpm_limit, True), + *((key, limit, True) for key, limit in io_token_limits_by_cache_key(deployment, current_minute)), + ) try: daily: Final = _DailyLimitDeployment.model_validate(deployment) except ValidationError: - return _counters_with_limits(model_id, ((tpm_key, tpm_limit, False), (rpm_key, rpm_limit, True))) + return _counters_with_limits(model_id, minute_counters) tpd_limit, rpd_limit = daily.limits() tpd_key, rpd_key = daily.cache_keys(now) return _counters_with_limits( - model_id, - ( - (tpm_key, tpm_limit, False), - (rpm_key, rpm_limit, True), - (tpd_key, tpd_limit, False), - (rpd_key, rpd_limit, True), - ), + model_id, (*minute_counters, (tpd_key, tpd_limit, False), (rpd_key, rpd_limit, True)) ) def _counters_of( diff --git a/tests/unit/router_utils/test_router_utils_common_utils.py b/tests/unit/router_utils/test_router_utils_common_utils.py index b3394b98590..bd174ddf7e1 100644 --- a/tests/unit/router_utils/test_router_utils_common_utils.py +++ b/tests/unit/router_utils/test_router_utils_common_utils.py @@ -997,12 +997,15 @@ def test_sync_web_search_request_goes_to_the_other_deployment_once_the_capable_o @pytest.mark.asyncio -async def test_web_search_request_never_goes_to_a_capable_deployment_whose_enforced_rpm_is_zero() -> None: +@pytest.mark.parametrize("zero_limit", ["rpm", "itpm", "otpm"]) +async def test_web_search_request_never_goes_to_a_capable_deployment_whose_enforced_limit_is_zero( + zero_limit: str, +) -> None: router: Final = Router( model_list=[ { "model_name": "mixed", - "litellm_params": {"model": "openai/gpt-x", "api_key": "fake-key", "mock_response": "hi", "rpm": 0}, + "litellm_params": {"model": "openai/gpt-x", "api_key": "fake-key", "mock_response": "hi", zero_limit: 0}, "model_info": {"id": "search"}, }, { diff --git a/tests/unit/test_router/test_enforce_model_rate_limits.py b/tests/unit/test_router/test_enforce_model_rate_limits.py index 28ea12bfaaf..6c452f193fb 100644 --- a/tests/unit/test_router/test_enforce_model_rate_limits.py +++ b/tests/unit/test_router/test_enforce_model_rate_limits.py @@ -18,6 +18,7 @@ import litellm from litellm import Router from litellm.caching.dual_cache import DualCache from litellm.caching.redis_cache import RedisCircuitBreakerOpenError +from litellm.router_utils.pre_call_checks.io_token_rate_limit_check import io_token_limits_by_cache_key from litellm.router_utils.pre_call_checks.model_rate_limit_check import ( ModelRateLimitingCheck, ) @@ -523,7 +524,18 @@ async def _async_rejects(check: ModelRateLimitingCheck, deployment: dict[str, ob return False -_FRESH_LIMITS: Final = ({"rpm": 0}, {"rpm": 0.0}, {"rpm": False}, {"rpd": 0}, {"tpm": 0}, {"tpd": 0}, {"rpm": 1, "tpm": 1}) +_FRESH_LIMITS: Final = ( + {"rpm": 0}, + {"rpm": 0.0}, + {"rpm": 0.5}, + {"rpm": False}, + {"rpd": 0}, + {"tpm": 0}, + {"tpd": 0}, + {"itpm": 0}, + {"otpm": 0}, + {"rpm": 1, "tpm": 1}, +) def _redis_batch_answering(usage: int) -> Callable[[list[str]], dict[str, int]]: @@ -673,6 +685,25 @@ class TestDailyLimits: assert check.deployment_ids_at_limit([spent, _deployment_with({"rpd": 2})]) == frozenset({"spent-test-id"}) assert rpm_key not in redis_cache.batch_get_cache.call_args.kwargs["key_list"] + @pytest.mark.asyncio + @pytest.mark.parametrize("limit_kind", ["itpm", "otpm"]) + @pytest.mark.parametrize(("used", "reported"), [(99, False), (100, True)], ids=["room-left", "used-up"]) + async def test_a_deployment_is_reported_at_its_token_limit_once_no_token_is_left( + self, limit_kind: str, used: int, reported: bool + ): + deployment = _deployment_with({limit_kind: 100}) + dual_cache = DualCache() + check = ModelRateLimitingCheck(dual_cache=dual_cache, clock=Clock(DAY_ONE)) + token_key = next( + key + for key, limit in io_token_limits_by_cache_key(deployment, DAY_ONE.strftime("%H-%M")) + if limit is not None + ) + dual_cache.set_cache(key=token_key, value=used, local_only=True) + + assert ("limit-test-id" in check.deployment_ids_at_limit([deployment])) is reported + assert ("limit-test-id" in await check.async_deployment_ids_at_limit([deployment], None)) is reported + @pytest.mark.parametrize("limits", _FRESH_LIMITS, ids=lambda limits: ",".join(f"{kind}={value!r}" for kind, value in limits.items())) def test_a_deployment_is_reported_at_its_limit_exactly_when_its_next_request_is_rejected( self, limits: Mapping[str, float] From 3ed9d0548d5e1dc5d244b589bb85575dcbbf6216 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 30 Oct 2026 10:36:23 -0700 Subject: [PATCH 18/24] fix(router): let a pin without web search yield to a search deployment in an earlier order tier --- litellm/router_utils/common_utils.py | 33 +++++-- .../deployment_affinity_check.py | 24 +++-- .../test_deployment_affinity_check.py | 97 ++++++++++++++++++- 3 files changed, 136 insertions(+), 18 deletions(-) diff --git a/litellm/router_utils/common_utils.py b/litellm/router_utils/common_utils.py index 0c6a41bd499..d0bac860829 100644 --- a/litellm/router_utils/common_utils.py +++ b/litellm/router_utils/common_utils.py @@ -545,21 +545,19 @@ def _deployment_ids_retried_in_request(request_kwargs: Mapping[str, object]) -> return tuple(chain.from_iterable(map(_breadcrumb_deployment_ids, breadcrumbs))) -def prefer_web_search_deployments( - deployments: list[_DeploymentT], # mutable-ok: router deployment-list contract - request_kwargs: Mapping[str, object] | None, - at_rate_limit: frozenset[str] = frozenset(), -) -> list[_DeploymentT]: # mutable-ok: returns a mutable deployment list +def _web_search_tried_deployment_ids( + request_kwargs: Mapping[str, object] | None, at_rate_limit: frozenset[str] +) -> frozenset[str] | None: if request_kwargs is None or not _is_web_search_request( request_kwargs.get("tools"), request_kwargs.get("web_search_options") ): - return deployments + return None if _retry_history_was_trimmed(request_kwargs.get("metadata")) or _retry_history_was_trimmed( request_kwargs.get("litellm_metadata") ): verbose_router_logger.debug("This request's retry history was trimmed, so every deployment stays") - return deployments - tried: Final = frozenset( + return None + return frozenset( ( *as_excluded_deployment_ids(request_kwargs.get("_excluded_deployment_ids")), *as_retry_skipped_deployment_ids(request_kwargs.get("_retry_skipped_deployment_ids")), @@ -567,6 +565,16 @@ def prefer_web_search_deployments( *at_rate_limit, ) ) + + +def prefer_web_search_deployments( + deployments: list[_DeploymentT], # mutable-ok: router deployment-list contract + request_kwargs: Mapping[str, object] | None, + at_rate_limit: frozenset[str] = frozenset(), +) -> list[_DeploymentT]: # mutable-ok: returns a mutable deployment list + tried: Final = _web_search_tried_deployment_ids(request_kwargs, at_rate_limit) + if tried is None: + return deployments untried: Final = (d for d in deployments if _deployment_id(d) not in tried) orders_with_an_untried_capable_deployment: Final = tuple( litellm.utils.get_deployment_order(d) for d in untried if _deployment_supports_web_search(d) @@ -582,6 +590,15 @@ def prefer_web_search_deployments( ] +def web_search_alternatives_to_pin( + pinned: _DeploymentT, deployments: Sequence[_DeploymentT], request_kwargs: Mapping[str, object] +) -> tuple[_DeploymentT, ...]: + tried: Final = _web_search_tried_deployment_ids(request_kwargs, frozenset()) + if tried is None or _deployment_supports_web_search(pinned): + return () + return tuple(d for d in deployments if _deployment_supports_web_search(d) and _deployment_id(d) not in tried) + + def pick_preferring_web_search( pick: Callable[[list[_DeploymentT]], _Picked | None], # mutable-ok: router deployment-list contract deployments: list[_DeploymentT], # mutable-ok: router deployment-list contract diff --git a/litellm/router_utils/pre_call_checks/deployment_affinity_check.py b/litellm/router_utils/pre_call_checks/deployment_affinity_check.py index 21f9b0a7c1b..fca532f9736 100644 --- a/litellm/router_utils/pre_call_checks/deployment_affinity_check.py +++ b/litellm/router_utils/pre_call_checks/deployment_affinity_check.py @@ -31,7 +31,11 @@ from litellm.caching.dual_cache import DualCache from litellm.constants import SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY, SESSION_ID_GENERATED_METADATA_KEY from litellm.integrations.custom_logger import CustomLogger, Span from litellm.responses.utils import ResponsesAPIRequestUtils -from litellm.router_utils.common_utils import deployment_ids_failed_in_request, prefer_web_search_deployments +from litellm.router_utils.common_utils import ( + deployment_ids_failed_in_request, + prefer_web_search_deployments, + web_search_alternatives_to_pin, +) from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import CallTypes @@ -470,7 +474,7 @@ class DeploymentAffinityCheck(CustomLogger): if session_model_id and session_model_id not in failed_deployment_ids: session_deployments: Final = self._pinned_deployments( - typed_healthy_deployments, pin_candidates, session_model_id + typed_healthy_deployments, pin_candidates, session_model_id, request_kwargs ) if session_deployments is not None: verbose_router_logger.debug( @@ -505,7 +509,9 @@ class DeploymentAffinityCheck(CustomLogger): if not model_id or model_id in failed_deployment_ids: return typed_healthy_deployments - pinned_deployments: Final = self._pinned_deployments(typed_healthy_deployments, pin_candidates, model_id) + pinned_deployments: Final = self._pinned_deployments( + typed_healthy_deployments, pin_candidates, model_id, request_kwargs + ) if pinned_deployments is None: verbose_router_logger.debug( "DeploymentAffinityCheck: pinned deployment=%s not found in healthy_deployments", @@ -522,14 +528,18 @@ class DeploymentAffinityCheck(CustomLogger): @classmethod def _pinned_deployments( - cls, healthy_deployments: Sequence[_DeploymentT], pin_candidates: Sequence[_DeploymentT], model_id: str + cls, + healthy_deployments: Sequence[_DeploymentT], + pin_candidates: Sequence[_DeploymentT], + model_id: str, + request_kwargs: Mapping[str, object], ) -> tuple[_DeploymentT, ...] | None: pinned: Final = cls._find_deployment_by_model_id(healthy_deployments=healthy_deployments, model_id=model_id) if pinned is None: return None - if cls._find_deployment_by_model_id(healthy_deployments=pin_candidates, model_id=model_id) is not None: - return (pinned,) - return (*pin_candidates, pinned) + if cls._find_deployment_by_model_id(healthy_deployments=pin_candidates, model_id=model_id) is None: + return (*pin_candidates, pinned) + return (*web_search_alternatives_to_pin(pinned, healthy_deployments, request_kwargs), pinned) async def _read_shared_session_pin(self, cache_key: str, ttl_seconds: int, parent_otel_span: Span | None) -> object: redis_cache: Final = self.cache.redis_cache diff --git a/tests/unit/router_utils/pre_call_checks/test_deployment_affinity_check.py b/tests/unit/router_utils/pre_call_checks/test_deployment_affinity_check.py index 94823593fe2..ae0c11bd802 100644 --- a/tests/unit/router_utils/pre_call_checks/test_deployment_affinity_check.py +++ b/tests/unit/router_utils/pre_call_checks/test_deployment_affinity_check.py @@ -1301,17 +1301,31 @@ async def test_previous_response_id_pin_holds_even_after_its_deployment_failed() ) == [healthy_deployments[0]] -def _mixed_web_search_router(optional_pre_call_checks: list[str]) -> litellm.Router: # mutable-ok: Router signature +def _mixed_web_search_router( + optional_pre_call_checks: list[str], # mutable-ok: Router signature + orders: tuple[int | None, int | None] = (None, None), +) -> litellm.Router: + no_search_order, search_order = orders return litellm.Router( model_list=[ { "model_name": "mixed", - "litellm_params": {"model": "openai/gpt-x", "api_key": "fake-key", "mock_response": "hi"}, + "litellm_params": { + "model": "openai/gpt-x", + "api_key": "fake-key", + "mock_response": "hi", + **({} if no_search_order is None else {"order": no_search_order}), + }, "model_info": {"id": "no-search", "supports_web_search": False}, }, { "model_name": "mixed", - "litellm_params": {"model": "openai/gpt-y", "api_key": "fake-key", "mock_response": "hi"}, + "litellm_params": { + "model": "openai/gpt-y", + "api_key": "fake-key", + "mock_response": "hi", + **({} if search_order is None else {"order": search_order}), + }, "model_info": {"id": "search"}, }, ], @@ -1362,6 +1376,83 @@ async def test_pin_on_a_deployment_without_web_search_yields_to_a_web_search_req assert searching._hidden_params["model_id"] == "search" +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("orders", "searched_on"), + [((2, 1), "search"), ((None, 1), "search"), ((1, 2), "no-search"), ((1, None), "no-search")], + ids=["search-tier-first", "only-search-tier-ordered", "pinned-tier-first", "only-pinned-tier-ordered"], +) +@pytest.mark.parametrize( + ("optional_pre_call_checks", "pin_key"), + [ + ( + ["deployment_affinity"], + DeploymentAffinityCheck.get_affinity_cache_key(model_group="mixed", user_key="caller-1"), + ), + ( + ["session_affinity"], + DeploymentAffinityCheck.get_session_affinity_cache_key( + model_group="mixed", session_id="session-1", user_key="caller-1" + ), + ), + ], + ids=["user-key-pin", "session-pin"], +) +async def test_pin_without_web_search_yields_to_a_web_search_request_only_when_search_sits_in_an_earlier_order_tier( + optional_pre_call_checks: list[str], # mutable-ok: Router signature + pin_key: str, + orders: tuple[int | None, int | None], + searched_on: str, +): + router: Final = _mixed_web_search_router(optional_pre_call_checks, orders) + await router.cache.async_set_cache(pin_key, {"model_id": "no-search"}, ttl=3600) + messages: Final = [{"role": "user", "content": "hi"}] + + plain: Final = await router.acompletion(model="mixed", messages=messages, metadata=dict(_CALLER_METADATA)) + searching: Final = await router.acompletion( + model="mixed", messages=messages, metadata=dict(_CALLER_METADATA), web_search_options={} + ) + + assert plain._hidden_params["model_id"] == "no-search" + assert searching._hidden_params["model_id"] == searched_on + + +@pytest.mark.asyncio +async def test_session_pin_on_a_later_order_tier_holds_for_a_web_search_request_when_it_can_search(): + router: Final = litellm.Router( + model_list=[ + { + "model_name": "tiered", + "litellm_params": { + "model": "openai/gpt-x", + "api_key": "fake-key", + "mock_response": "hi", + "order": order, + }, + "model_info": {"id": model_id}, + } + for model_id, order in (("search-first-tier", 1), ("search-second-tier", 2)) + ], + optional_pre_call_checks=["session_affinity"], + ) + await router.cache.async_set_cache( + DeploymentAffinityCheck.get_session_affinity_cache_key( + model_group="tiered", session_id="session-1", user_key="caller-1" + ), + {"model_id": "search-second-tier"}, + ttl=3600, + ) + + searching: Final = await router.acompletion( + model="tiered", + messages=[{"role": "user", "content": "hi"}], + metadata=dict(_CALLER_METADATA), + web_search_options={}, + ) + + assert searching._hidden_params["model_id"] == "search-second-tier" + + @pytest.mark.asyncio @pytest.mark.parametrize( ("search_supported", "tried"), From 1332c4d0997d60aea13fc8c30cd87b495afa8b4c Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 30 Oct 2026 18:06:21 -0700 Subject: [PATCH 19/24] fix(router): keep a pin when the search deployment it would yield to is at its rate limit --- litellm/router.py | 23 ++++--- litellm/router_utils/common_utils.py | 7 ++- .../deployment_affinity_check.py | 49 ++++++++++----- .../test_deployment_affinity_check.py | 62 +++++++++++++++++++ 4 files changed, 114 insertions(+), 27 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index f887f962492..73efa33f477 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -3157,6 +3157,7 @@ class Router(DeprecatedRouterMembers): model_group_affinity_config=self.model_group_affinity_config, is_priority_group=self._is_priority_routing_group, session_affinity_repin_on_failover=self.session_affinity_repin_on_failover, + deployment_ids_at_limit=self._async_deployment_ids_at_rate_limit, ) self.optional_callbacks.append(affinity_callback) litellm.logging_callback_manager.add_litellm_callback(affinity_callback) @@ -3204,6 +3205,7 @@ class Router(DeprecatedRouterMembers): model_group_affinity_config=self.model_group_affinity_config, is_priority_group=self._is_priority_routing_group, session_affinity_repin_on_failover=self.session_affinity_repin_on_failover, + deployment_ids_at_limit=self._async_deployment_ids_at_rate_limit, ) self.optional_callbacks.append(affinity_callback) litellm.logging_callback_manager.add_litellm_callback(affinity_callback) @@ -10403,6 +10405,18 @@ class Router(DeprecatedRouterMembers): def _model_rate_limit_checks(self) -> tuple[ModelRateLimitingCheck, ...]: return tuple(callback for callback in self._routing_callbacks() if isinstance(callback, ModelRateLimitingCheck)) + async def _async_deployment_ids_at_rate_limit( + self, deployments: Sequence[dict[str, object]], parent_otel_span: Span | None + ) -> frozenset[str]: + return frozenset( + chain.from_iterable( + [ + await check.async_deployment_ids_at_limit(deployments, parent_otel_span) + for check in self._model_rate_limit_checks() + ] + ) + ) + def routing_strategy_pre_call_checks(self, deployment: dict): """ Mimics 'async_routing_strategy_pre_call_checks' @@ -15828,14 +15842,7 @@ class Router(DeprecatedRouterMembers): request_kwargs=routing_kwargs, deployments=cast(list[dict[str, object]], healthy_deployments), # cast-ok: str-keyed dicts ) - at_rate_limit: Final = frozenset( - chain.from_iterable( - [ - await check.async_deployment_ids_at_limit(web_search_capable, parent_otel_span) - for check in self._model_rate_limit_checks() - ] - ) - ) + at_rate_limit: Final = await self._async_deployment_ids_at_rate_limit(web_search_capable, parent_otel_span) if strategy == "simple-shuffle": shuffled: Final = simple_shuffle( resolve_model_alias=self._core_get_model_from_alias, diff --git a/litellm/router_utils/common_utils.py b/litellm/router_utils/common_utils.py index d0bac860829..0b228824eb2 100644 --- a/litellm/router_utils/common_utils.py +++ b/litellm/router_utils/common_utils.py @@ -591,9 +591,12 @@ def prefer_web_search_deployments( def web_search_alternatives_to_pin( - pinned: _DeploymentT, deployments: Sequence[_DeploymentT], request_kwargs: Mapping[str, object] + pinned: _DeploymentT, + deployments: Sequence[_DeploymentT], + request_kwargs: Mapping[str, object], + at_rate_limit: frozenset[str] = frozenset(), ) -> tuple[_DeploymentT, ...]: - tried: Final = _web_search_tried_deployment_ids(request_kwargs, frozenset()) + tried: Final = _web_search_tried_deployment_ids(request_kwargs, at_rate_limit) if tried is None or _deployment_supports_web_search(pinned): return () return tuple(d for d in deployments if _deployment_supports_web_search(d) and _deployment_id(d) not in tried) diff --git a/litellm/router_utils/pre_call_checks/deployment_affinity_check.py b/litellm/router_utils/pre_call_checks/deployment_affinity_check.py index fca532f9736..c72d05087ec 100644 --- a/litellm/router_utils/pre_call_checks/deployment_affinity_check.py +++ b/litellm/router_utils/pre_call_checks/deployment_affinity_check.py @@ -13,7 +13,7 @@ where routing to a consistent deployment is still beneficial. """ import hashlib -from collections.abc import Callable, Mapping, Sequence +from collections.abc import Awaitable, Callable, Mapping, Sequence from typing import Final, TypeVar, cast from pydantic import TypeAdapter @@ -35,6 +35,7 @@ from litellm.router_utils.common_utils import ( deployment_ids_failed_in_request, prefer_web_search_deployments, web_search_alternatives_to_pin, + web_search_capable_deployments, ) from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import CallTypes @@ -116,6 +117,8 @@ class DeploymentAffinityCheck(CustomLogger): model_group_affinity_config: dict[str, list[str]] | None = None, is_priority_group: Callable[[str], bool] | None = None, session_affinity_repin_on_failover: bool = False, + deployment_ids_at_limit: Callable[[Sequence[dict[str, object]], Span | None], Awaitable[frozenset[str]]] + | None = None, ): super().__init__() self.cache = cache @@ -126,6 +129,7 @@ class DeploymentAffinityCheck(CustomLogger): self.model_group_affinity_config: dict[str, list[str]] = model_group_affinity_config or {} self.is_priority_group = is_priority_group self.session_affinity_repin_on_failover = session_affinity_repin_on_failover + self.deployment_ids_at_limit = deployment_ids_at_limit def _get_effective_flags(self, model_group: str) -> tuple[bool, bool, bool]: """ @@ -439,10 +443,6 @@ class DeploymentAffinityCheck(CustomLogger): failed_deployment_ids: Final = deployment_ids_failed_in_request( cast(Mapping[str, object], request_kwargs) # cast-ok: request kwargs keys are param names ) - pin_candidates: Final = prefer_web_search_deployments( - cast(list[dict[str, object]], healthy_deployments), # cast-ok: router deployments are plain dicts - request_kwargs, - ) user_key: Final = ( self.get_user_key_from_request_kwargs(request_kwargs=request_kwargs) if (session_affinity_active or enable_user_key) @@ -473,8 +473,8 @@ class DeploymentAffinityCheck(CustomLogger): session_model_id = session_cache_result if session_model_id and session_model_id not in failed_deployment_ids: - session_deployments: Final = self._pinned_deployments( - typed_healthy_deployments, pin_candidates, session_model_id, request_kwargs + session_deployments: Final = await self._pinned_deployments( + typed_healthy_deployments, session_model_id, request_kwargs, parent_otel_span ) if session_deployments is not None: verbose_router_logger.debug( @@ -509,8 +509,8 @@ class DeploymentAffinityCheck(CustomLogger): if not model_id or model_id in failed_deployment_ids: return typed_healthy_deployments - pinned_deployments: Final = self._pinned_deployments( - typed_healthy_deployments, pin_candidates, model_id, request_kwargs + pinned_deployments: Final = await self._pinned_deployments( + typed_healthy_deployments, model_id, request_kwargs, parent_otel_span ) if pinned_deployments is None: verbose_router_logger.debug( @@ -526,20 +526,35 @@ class DeploymentAffinityCheck(CustomLogger): ) return list(pinned_deployments) - @classmethod - def _pinned_deployments( - cls, - healthy_deployments: Sequence[_DeploymentT], - pin_candidates: Sequence[_DeploymentT], + async def _pinned_deployments( + self, + healthy_deployments: list[_DeploymentT], # mutable-ok: router deployment-list contract model_id: str, request_kwargs: Mapping[str, object], + parent_otel_span: Span | None, ) -> tuple[_DeploymentT, ...] | None: - pinned: Final = cls._find_deployment_by_model_id(healthy_deployments=healthy_deployments, model_id=model_id) + pinned: Final = self._find_deployment_by_model_id(healthy_deployments=healthy_deployments, model_id=model_id) if pinned is None: return None - if cls._find_deployment_by_model_id(healthy_deployments=pin_candidates, model_id=model_id) is None: + at_rate_limit: Final = await self._web_search_deployment_ids_at_limit( + healthy_deployments, request_kwargs, parent_otel_span + ) + pin_candidates: Final = prefer_web_search_deployments(healthy_deployments, request_kwargs, at_rate_limit) + if self._find_deployment_by_model_id(healthy_deployments=pin_candidates, model_id=model_id) is None: return (*pin_candidates, pinned) - return (*web_search_alternatives_to_pin(pinned, healthy_deployments, request_kwargs), pinned) + return (*web_search_alternatives_to_pin(pinned, healthy_deployments, request_kwargs, at_rate_limit), pinned) + + async def _web_search_deployment_ids_at_limit( + self, + deployments: Sequence[_DeploymentT], + request_kwargs: Mapping[str, object], + parent_otel_span: Span | None, + ) -> frozenset[str]: + if self.deployment_ids_at_limit is None: + return frozenset() + return await self.deployment_ids_at_limit( + web_search_capable_deployments(deployments, request_kwargs), parent_otel_span + ) async def _read_shared_session_pin(self, cache_key: str, ttl_seconds: int, parent_otel_span: Span | None) -> object: redis_cache: Final = self.cache.redis_cache diff --git a/tests/unit/router_utils/pre_call_checks/test_deployment_affinity_check.py b/tests/unit/router_utils/pre_call_checks/test_deployment_affinity_check.py index ae0c11bd802..12881f38919 100644 --- a/tests/unit/router_utils/pre_call_checks/test_deployment_affinity_check.py +++ b/tests/unit/router_utils/pre_call_checks/test_deployment_affinity_check.py @@ -2,6 +2,7 @@ import asyncio import itertools import json from collections.abc import Sequence +from datetime import datetime, timezone from typing import Final from unittest.mock import AsyncMock, patch @@ -1304,6 +1305,8 @@ async def test_previous_response_id_pin_holds_even_after_its_deployment_failed() def _mixed_web_search_router( optional_pre_call_checks: list[str], # mutable-ok: Router signature orders: tuple[int | None, int | None] = (None, None), + search_rpm: int | None = None, + model_group_affinity_config: dict[str, list[str]] | None = None, # mutable-ok: Router signature ) -> litellm.Router: no_search_order, search_order = orders return litellm.Router( @@ -1325,11 +1328,14 @@ def _mixed_web_search_router( "api_key": "fake-key", "mock_response": "hi", **({} if search_order is None else {"order": search_order}), + **({} if search_rpm is None else {"rpm": search_rpm}), }, "model_info": {"id": "search"}, }, ], optional_pre_call_checks=optional_pre_call_checks, + model_group_affinity_config=model_group_affinity_config, + num_retries=0, ) @@ -1417,6 +1423,62 @@ async def test_pin_without_web_search_yields_to_a_web_search_request_only_when_s assert searching._hidden_params["model_id"] == searched_on +@pytest.mark.asyncio +@pytest.mark.parametrize( + "orders", [(None, 1), (2, 1), (None, None)], ids=["pin-unordered", "pin-in-later-tier", "same-tier"] +) +@pytest.mark.parametrize( + ("affinity_check", "per_group", "pin_key"), + [ + ( + "deployment_affinity", + False, + DeploymentAffinityCheck.get_affinity_cache_key(model_group="mixed", user_key="caller-1"), + ), + ( + "session_affinity", + False, + DeploymentAffinityCheck.get_session_affinity_cache_key( + model_group="mixed", session_id="session-1", user_key="caller-1" + ), + ), + ( + "deployment_affinity", + True, + DeploymentAffinityCheck.get_affinity_cache_key(model_group="mixed", user_key="caller-1"), + ), + ], + ids=["user-key-pin", "session-pin", "per-group-user-key-pin"], +) +async def test_pin_without_web_search_holds_for_a_web_search_request_once_the_earlier_search_tier_hits_its_rpm( + affinity_check: str, + per_group: bool, + pin_key: str, + orders: tuple[int | None, int | None], +): + router: Final = _mixed_web_search_router( + ["enforce_model_rate_limits"] if per_group else [affinity_check, "enforce_model_rate_limits"], + orders, + search_rpm=1, + model_group_affinity_config={"mixed": [affinity_check]} if per_group else None, + ) + await router.cache.async_set_cache(pin_key, {"model_id": "no-search"}, ttl=3600) + messages: Final = [{"role": "user", "content": "hi"}] + + with patch( + "litellm.router_utils.pre_call_checks.model_rate_limit_check.get_utc_datetime", + return_value=datetime(2026, 1, 1, 12, 30, 15, tzinfo=timezone.utc), + ): + served: Final = [ + await router.acompletion( + model="mixed", messages=messages, metadata=dict(_CALLER_METADATA), web_search_options={} + ) + for _ in range(4) + ] + + assert [response._hidden_params["model_id"] for response in served] == ["search", *["no-search"] * 3] + + @pytest.mark.asyncio async def test_session_pin_on_a_later_order_tier_holds_for_a_web_search_request_when_it_can_search(): router: Final = litellm.Router( From 0f18edba7d8b9ed4390bdc895a5742649b3e1032 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 30 Oct 2026 18:55:49 -0700 Subject: [PATCH 20/24] fix(router): let a pin yield only to search deployments that are untried and under their limits --- litellm/router_utils/common_utils.py | 2 +- .../deployment_affinity_check.py | 4 - .../test_deployment_affinity_check.py | 94 +++++++++++++++++++ 3 files changed, 95 insertions(+), 5 deletions(-) diff --git a/litellm/router_utils/common_utils.py b/litellm/router_utils/common_utils.py index 0b228824eb2..36adc86cec6 100644 --- a/litellm/router_utils/common_utils.py +++ b/litellm/router_utils/common_utils.py @@ -597,7 +597,7 @@ def web_search_alternatives_to_pin( at_rate_limit: frozenset[str] = frozenset(), ) -> tuple[_DeploymentT, ...]: tried: Final = _web_search_tried_deployment_ids(request_kwargs, at_rate_limit) - if tried is None or _deployment_supports_web_search(pinned): + if tried is None or (_deployment_supports_web_search(pinned) and _deployment_id(pinned) not in tried): return () return tuple(d for d in deployments if _deployment_supports_web_search(d) and _deployment_id(d) not in tried) diff --git a/litellm/router_utils/pre_call_checks/deployment_affinity_check.py b/litellm/router_utils/pre_call_checks/deployment_affinity_check.py index c72d05087ec..3aa59faacce 100644 --- a/litellm/router_utils/pre_call_checks/deployment_affinity_check.py +++ b/litellm/router_utils/pre_call_checks/deployment_affinity_check.py @@ -33,7 +33,6 @@ from litellm.integrations.custom_logger import CustomLogger, Span from litellm.responses.utils import ResponsesAPIRequestUtils from litellm.router_utils.common_utils import ( deployment_ids_failed_in_request, - prefer_web_search_deployments, web_search_alternatives_to_pin, web_search_capable_deployments, ) @@ -539,9 +538,6 @@ class DeploymentAffinityCheck(CustomLogger): at_rate_limit: Final = await self._web_search_deployment_ids_at_limit( healthy_deployments, request_kwargs, parent_otel_span ) - pin_candidates: Final = prefer_web_search_deployments(healthy_deployments, request_kwargs, at_rate_limit) - if self._find_deployment_by_model_id(healthy_deployments=pin_candidates, model_id=model_id) is None: - return (*pin_candidates, pinned) return (*web_search_alternatives_to_pin(pinned, healthy_deployments, request_kwargs, at_rate_limit), pinned) async def _web_search_deployment_ids_at_limit( diff --git a/tests/unit/router_utils/pre_call_checks/test_deployment_affinity_check.py b/tests/unit/router_utils/pre_call_checks/test_deployment_affinity_check.py index 12881f38919..1eae7c2bcd2 100644 --- a/tests/unit/router_utils/pre_call_checks/test_deployment_affinity_check.py +++ b/tests/unit/router_utils/pre_call_checks/test_deployment_affinity_check.py @@ -1479,6 +1479,100 @@ async def test_pin_without_web_search_holds_for_a_web_search_request_once_the_ea assert [response._hidden_params["model_id"] for response in served] == ["search", *["no-search"] * 3] +def _tiered_web_search_router( + optional_pre_call_checks: list[str], # mutable-ok: Router signature + deployments: tuple[tuple[str, bool, int | None, int | None], ...], +) -> litellm.Router: + return litellm.Router( + model_list=[ + { + "model_name": "mixed", + "litellm_params": { + "model": "openai/gpt-x", + "api_key": "fake-key", + "mock_response": "hi", + **({} if order is None else {"order": order}), + **({} if rpm is None else {"rpm": rpm}), + }, + "model_info": {"id": model_id, **({} if supports_web_search else {"supports_web_search": False})}, + } + for model_id, supports_web_search, order, rpm in deployments + ], + optional_pre_call_checks=optional_pre_call_checks, + num_retries=0, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("deployments", "pinned_id", "served_ids"), + [ + ( + (("search-early", True, 1, 1), ("no-search", False, None, None), ("search", True, None, None)), + "no-search", + ["search-early", *["search"] * 3], + ), + ( + (("search-early", True, 1, 1), ("no-search", False, 2, None), ("search", True, 2, None)), + "no-search", + ["search-early", *["search"] * 3], + ), + ( + (("no-search-early", False, 1, None), ("no-search", False, 2, None), ("search", True, 2, None)), + "no-search", + ["search"] * 4, + ), + ( + (("search-pinned", True, None, 1), ("search", True, None, None)), + "search-pinned", + ["search-pinned", *["search"] * 3], + ), + ], + ids=[ + "rate-limited-search-before-unordered-pin", + "rate-limited-search-before-pin-tier", + "flagged-before-pin-tier", + "pinned-search-at-its-rpm", + ], +) +@pytest.mark.parametrize( + ("affinity_check", "pin_key"), + [ + ("deployment_affinity", DeploymentAffinityCheck.get_affinity_cache_key(model_group="mixed", user_key="caller-1")), + ( + "session_affinity", + DeploymentAffinityCheck.get_session_affinity_cache_key( + model_group="mixed", session_id="session-1", user_key="caller-1" + ), + ), + ], + ids=["user-key-pin", "session-pin"], +) +async def test_web_search_request_leaves_a_pin_only_for_an_untried_search_deployment_under_its_limit( + affinity_check: str, + pin_key: str, + deployments: tuple[tuple[str, bool, int | None, int | None], ...], + pinned_id: str, + served_ids: list[str], # mutable-ok: expected id sequence +): + router: Final = _tiered_web_search_router([affinity_check, "enforce_model_rate_limits"], deployments) + await router.cache.async_set_cache(pin_key, {"model_id": pinned_id}, ttl=3600) + messages: Final = [{"role": "user", "content": "hi"}] + + with patch( + "litellm.router_utils.pre_call_checks.model_rate_limit_check.get_utc_datetime", + return_value=datetime(2026, 1, 1, 12, 30, 15, tzinfo=timezone.utc), + ): + served: Final = [ + await router.acompletion( + model="mixed", messages=messages, metadata=dict(_CALLER_METADATA), web_search_options={} + ) + for _ in range(4) + ] + + assert [response._hidden_params["model_id"] for response in served] == served_ids + + @pytest.mark.asyncio async def test_session_pin_on_a_later_order_tier_holds_for_a_web_search_request_when_it_can_search(): router: Final = litellm.Router( From 613ae047f3f661f5bc26e65f5e8afaf0b79287ed Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 30 Oct 2026 22:40:28 -0700 Subject: [PATCH 21/24] test(router): use a fixed clock and a tpm limit in the web search rate-limit tests --- .../test_router_utils_common_utils.py | 46 ++++++++++--------- 1 file changed, 24 insertions(+), 22 deletions(-) diff --git a/tests/unit/router_utils/test_router_utils_common_utils.py b/tests/unit/router_utils/test_router_utils_common_utils.py index ecec1f8b251..27ac6d845cc 100644 --- a/tests/unit/router_utils/test_router_utils_common_utils.py +++ b/tests/unit/router_utils/test_router_utils_common_utils.py @@ -1,6 +1,7 @@ import asyncio import logging from collections.abc import Mapping +from datetime import datetime, timezone from itertools import chain from types import MappingProxyType from typing import Dict, Final, List, Optional, Union @@ -10,6 +11,7 @@ import pytest import litellm from litellm import Router +from litellm.caching.dual_cache import DualCache from litellm.constants import ROUTER_FALLBACK_ERROR_DETAIL_MAX_CHARS from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap from litellm.proxy._types import UserAPIKeyAuth @@ -29,6 +31,7 @@ from litellm.litellm_core_utils.routing_common_utils import ( PROVIDER_SCOPED_CREDENTIAL_PARAMS, warn_on_provider_credential_mismatch, ) +from litellm.router_utils.pre_call_checks.model_rate_limit_check import ModelRateLimitingCheck from litellm.types.router import Deployment, FallbackSkipReason, LiteLLM_Params, ModelInfo, SkippedFallbackTarget from litellm.types.utils import ProviderSpecificModelInfo @@ -933,10 +936,13 @@ _SEARCH_RPM_OF_ONE: Final = ( ) +_ONE_MINUTE: Final = datetime(2026, 10, 28, 12, 0, 30, tzinfo=timezone.utc) + + def _web_search_group_with_model_rate_limits( search_rpm: Mapping[str, Mapping[str, object]] = MappingProxyType({"litellm_params": {"rpm": 1}}), ) -> Router: - return Router( + router: Final = Router( model_list=[ { "model_name": "mixed", @@ -959,8 +965,11 @@ def _web_search_group_with_model_rate_limits( "model_info": {"id": "no-search", "supports_web_search": False}, }, ], - optional_pre_call_checks=["enforce_model_rate_limits"], ) + check: Final = ModelRateLimitingCheck(dual_cache=DualCache(), clock=lambda: _ONE_MINUTE) + router.optional_callbacks = [check] + litellm.logging_callback_manager.add_litellm_callback(check) + return router @pytest.mark.asyncio @@ -1037,7 +1046,7 @@ def _rate_limited_web_search_group(routing_strategy: str) -> Router: model_list=[ { "model_name": "mixed", - "litellm_params": {"model": "openai/gpt-x", "api_key": "fake-key", "mock_response": "hi", "rpm": 1}, + "litellm_params": {"model": "openai/gpt-x", "api_key": "fake-key", "mock_response": "hi", "tpm": 1}, "model_info": {"id": "search"}, }, { @@ -1052,29 +1061,22 @@ def _rate_limited_web_search_group(routing_strategy: str) -> Router: @pytest.mark.asyncio @pytest.mark.parametrize("routing_strategy", ["usage-based-routing-v2", "cost-based-routing"]) -async def test_web_search_request_goes_to_the_other_deployment_once_the_capable_one_is_at_its_rpm_limit( +async def test_web_search_request_goes_to_the_other_deployment_when_the_capable_one_has_no_tpm_room( routing_strategy: str, ) -> None: router: Final = _rate_limited_web_search_group(routing_strategy) - served_by: Final = [] # mutable-ok: collects one id per sequential request - for _ in range(3): - response = await router.acompletion( - model="mixed", messages=[{"role": "user", "content": "hi"}], **_ANTHROPIC_WEB_SEARCH_REQUEST - ) - served_by.append(response._hidden_params["model_id"]) - await asyncio.sleep(0.2) - assert served_by == ["search", "no-search", "no-search"] - - -def test_sync_web_search_request_goes_to_the_other_deployment_once_the_capable_one_is_at_its_rpm_limit() -> None: - router: Final = _rate_limited_web_search_group("usage-based-routing-v2") - served_by: Final = tuple( - router.completion( - model="mixed", messages=[{"role": "user", "content": "hi"}], **_ANTHROPIC_WEB_SEARCH_REQUEST - )._hidden_params["model_id"] - for _ in range(3) + response: Final = await router.acompletion( + model="mixed", messages=[{"role": "user", "content": "hi"}], **_ANTHROPIC_WEB_SEARCH_REQUEST ) - assert served_by == ("search", "no-search", "no-search") + assert response._hidden_params["model_id"] == "no-search" + + +def test_sync_web_search_request_goes_to_the_other_deployment_when_the_capable_one_has_no_tpm_room() -> None: + router: Final = _rate_limited_web_search_group("usage-based-routing-v2") + response: Final = router.completion( + model="mixed", messages=[{"role": "user", "content": "hi"}], **_ANTHROPIC_WEB_SEARCH_REQUEST + ) + assert response._hidden_params["model_id"] == "no-search" def test_sync_pick_falls_back_to_every_deployment_when_the_strategy_finds_none_that_can_search() -> None: From 141c5acbb11fa892e67156a5d8a9e3971af0aa21 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 31 Oct 2026 02:51:33 -0700 Subject: [PATCH 22/24] fix(router): retry an untried deployment after every web search one failed, and leave pins to plain requests in mixed groups --- .../routing_common_utils.py | 27 +++---- litellm/router.py | 5 ++ .../deployment_affinity_check.py | 11 +++ .../test_deployment_affinity_check.py | 70 +++++++++++++++++++ .../test_router_utils_common_utils.py | 40 ++++++++++- 5 files changed, 139 insertions(+), 14 deletions(-) diff --git a/litellm/litellm_core_utils/routing_common_utils.py b/litellm/litellm_core_utils/routing_common_utils.py index 36adc86cec6..1565d12a5cd 100644 --- a/litellm/litellm_core_utils/routing_common_utils.py +++ b/litellm/litellm_core_utils/routing_common_utils.py @@ -36,7 +36,7 @@ from litellm.types.utils import ( _V = TypeVar("_V") _Picked: Final = TypeVar("_Picked") -_DeploymentT: Final = TypeVar("_DeploymentT", bound=dict[str, object]) +_DeploymentT: Final = TypeVar("_DeploymentT", bound=Mapping[str, object]) MID_STREAM_CONTINUATION_KWARG: Final = "_mid_stream_continuation" ROUTER_ONLY_CALL_KWARGS: Final = frozenset({"silent_model", "include_fallback_errors", MID_STREAM_CONTINUATION_KWARG}) @@ -428,7 +428,7 @@ def deployments_the_caller_may_use( return () -def _deployment_supports_web_search(deployment: dict) -> bool: +def _deployment_supports_web_search(deployment: Mapping[str, object]) -> bool: """ Check if a deployment supports web search. @@ -440,12 +440,8 @@ def _deployment_supports_web_search(deployment: dict) -> bool: model_prices_and_context_window.json doesn't have supports_web_search tags on all models yet. TODO: backfill and add fallback. """ - model_info: Final = deployment.get("model_info", {}) - - if "supports_web_search" in model_info: - return model_info["supports_web_search"] - - return True + model_info: Final = as_str_mapping(deployment.get("model_info")) or {} + return bool(model_info.get("supports_web_search", True)) def _is_web_search_tool(tool: object) -> bool: @@ -502,7 +498,7 @@ def filter_web_search_deployments( # Filter out deployments that don't support web search final_deployments: Final = [d for d in healthy_deployments if _deployment_supports_web_search(d)] - if len(healthy_deployments) > 0 and len(final_deployments) == 0: + if healthy_deployments and not final_deployments: verbose_logger.warning("No deployments support web search for request") return final_deployments @@ -575,13 +571,16 @@ def prefer_web_search_deployments( tried: Final = _web_search_tried_deployment_ids(request_kwargs, at_rate_limit) if tried is None: return deployments - untried: Final = (d for d in deployments if _deployment_id(d) not in tried) + untried: Final = [d for d in deployments if _deployment_id(d) not in tried] orders_with_an_untried_capable_deployment: Final = tuple( litellm.utils.get_deployment_order(d) for d in untried if _deployment_supports_web_search(d) ) if not orders_with_an_untried_capable_deployment: - verbose_router_logger.debug("No deployment left for this web search request supports web search, so all stay") - return deployments + if not untried or not any(_deployment_supports_web_search(d) for d in deployments): + verbose_router_logger.debug("No untried deployment for this web search request can search, so all stay") + return deployments + verbose_router_logger.debug("Every deployment that can search was tried, so the untried ones stay") + return untried return [ d for d in deployments @@ -590,6 +589,10 @@ def prefer_web_search_deployments( ] +def web_search_preference_can_narrow(deployments: Sequence[_DeploymentT], request_kwargs: Mapping[str, object]) -> bool: + return 0 < len(web_search_capable_deployments(deployments, request_kwargs)) < len(deployments) + + def web_search_alternatives_to_pin( pinned: _DeploymentT, deployments: Sequence[_DeploymentT], diff --git a/litellm/router.py b/litellm/router.py index 14441331dc1..de468a8ae2d 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -3178,6 +3178,7 @@ class Router(DeprecatedRouterMembers): is_priority_group=self._is_priority_routing_group, session_affinity_repin_on_failover=self.session_affinity_repin_on_failover, deployment_ids_at_limit=self._async_deployment_ids_at_rate_limit, + model_group_deployments=self._model_group_deployments, ) self.optional_callbacks.append(affinity_callback) litellm.logging_callback_manager.add_litellm_callback(affinity_callback) @@ -3226,6 +3227,7 @@ class Router(DeprecatedRouterMembers): is_priority_group=self._is_priority_routing_group, session_affinity_repin_on_failover=self.session_affinity_repin_on_failover, deployment_ids_at_limit=self._async_deployment_ids_at_rate_limit, + model_group_deployments=self._model_group_deployments, ) self.optional_callbacks.append(affinity_callback) litellm.logging_callback_manager.add_litellm_callback(affinity_callback) @@ -10505,6 +10507,9 @@ class Router(DeprecatedRouterMembers): def _model_rate_limit_checks(self) -> tuple[ModelRateLimitingCheck, ...]: return tuple(callback for callback in self._routing_callbacks() if isinstance(callback, ModelRateLimitingCheck)) + def _model_group_deployments(self, model_group: str) -> Sequence[Mapping[str, object]]: + return self._core_get_model_list(model_name=model_group) or () + async def _async_deployment_ids_at_rate_limit( self, deployments: Sequence[dict[str, object]], parent_otel_span: Span | None ) -> frozenset[str]: diff --git a/litellm/router_utils/pre_call_checks/deployment_affinity_check.py b/litellm/router_utils/pre_call_checks/deployment_affinity_check.py index 28416a3109c..791609e0d52 100644 --- a/litellm/router_utils/pre_call_checks/deployment_affinity_check.py +++ b/litellm/router_utils/pre_call_checks/deployment_affinity_check.py @@ -34,6 +34,7 @@ from litellm.litellm_core_utils.routing_common_utils import ( deployment_ids_failed_in_request, web_search_alternatives_to_pin, web_search_capable_deployments, + web_search_preference_can_narrow, ) from litellm.responses.utils import ResponsesAPIRequestUtils from litellm.types.llms.openai import AllMessageValues @@ -118,6 +119,7 @@ class DeploymentAffinityCheck(CustomLogger): session_affinity_repin_on_failover: bool = False, deployment_ids_at_limit: Callable[[Sequence[dict[str, object]], Span | None], Awaitable[frozenset[str]]] | None = None, + model_group_deployments: Callable[[str], Sequence[Mapping[str, object]]] | None = None, ): super().__init__() self.cache = cache @@ -129,6 +131,7 @@ class DeploymentAffinityCheck(CustomLogger): self.is_priority_group = is_priority_group self.session_affinity_repin_on_failover = session_affinity_repin_on_failover self.deployment_ids_at_limit = deployment_ids_at_limit + self.model_group_deployments = model_group_deployments def _get_effective_flags(self, model_group: str) -> tuple[bool, bool, bool]: """ @@ -669,6 +672,14 @@ class DeploymentAffinityCheck(CustomLogger): if not ((enable_user_key and user_key is not None) or session_id is not None): return None + if self.model_group_deployments is not None and web_search_preference_can_narrow( + self.model_group_deployments(deployment_model_name), kwargs + ): + verbose_router_logger.debug( + "DeploymentAffinityCheck: web search request in a group with deployments that cannot search sets no pin" + ) + return None + model_info = kwargs.get("model_info") if not isinstance(model_info, dict): model_info = None diff --git a/tests/unit/router_utils/pre_call_checks/test_deployment_affinity_check.py b/tests/unit/router_utils/pre_call_checks/test_deployment_affinity_check.py index 1eae7c2bcd2..6872c3aa5a2 100644 --- a/tests/unit/router_utils/pre_call_checks/test_deployment_affinity_check.py +++ b/tests/unit/router_utils/pre_call_checks/test_deployment_affinity_check.py @@ -1696,3 +1696,73 @@ async def test_key_pin_holds_for_a_web_search_request_when_tag_filtering_removes } assert served_by == {"flagged-a"} + + + +def _web_search_group_router( + optional_pre_call_checks: list[str], # mutable-ok: Router signature + heavy_can_search: bool, + heavy_weight: int, +) -> litellm.Router: + return litellm.Router( + model_list=[ + { + "model_name": "group", + "litellm_params": { + "model": "openai/gpt-x", + "api_key": "fake-key", + "mock_response": "hi", + "weight": heavy_weight, + }, + "model_info": {"id": "heavy", **({} if heavy_can_search else {"supports_web_search": False})}, + }, + { + "model_name": "group", + "litellm_params": {"model": "openai/gpt-y", "api_key": "fake-key", "mock_response": "hi", "weight": 1}, + "model_info": {"id": "light"}, + }, + ], + optional_pre_call_checks=optional_pre_call_checks, + num_retries=0, + ) + + +_AFFINITY_KINDS: Final = pytest.mark.parametrize( + "optional_pre_call_checks", [["deployment_affinity"], ["session_affinity"]], ids=["user-key-pin", "session-pin"] +) + + +async def _served_by(router: litellm.Router, caller: int, **web_search_request: object) -> object: + response: Final = await router.acompletion( + model="group", + messages=[{"role": "user", "content": "hi"}], + metadata={"user_api_key_hash": f"caller-{caller}", "session_id": f"session-{caller}"}, + **web_search_request, + ) + return response._hidden_params["model_id"] + + +@pytest.mark.asyncio +@_AFFINITY_KINDS +async def test_a_first_web_search_request_does_not_pin_plain_traffic_to_the_deployment_that_can_search( + optional_pre_call_checks: list[str], # mutable-ok: Router signature +): + router: Final = _web_search_group_router(optional_pre_call_checks, heavy_can_search=False, heavy_weight=10**6) + + searched_on: Final = await _served_by(router, 1, web_search_options={}) + plain_on: Final = {await _served_by(router, 1) for _ in range(5)} + + assert searched_on == "light" + assert plain_on == {"heavy"} + + +@pytest.mark.asyncio +@_AFFINITY_KINDS +async def test_a_first_web_search_request_still_pins_a_group_where_every_deployment_can_search( + optional_pre_call_checks: list[str], # mutable-ok: Router signature +): + router: Final = _web_search_group_router(optional_pre_call_checks, heavy_can_search=True, heavy_weight=1) + + for caller in range(10): + searched_on = await _served_by(router, caller, web_search_options={}) + assert {await _served_by(router, caller) for _ in range(3)} == {searched_on} diff --git a/tests/unit/router_utils/test_router_utils_common_utils.py b/tests/unit/router_utils/test_router_utils_common_utils.py index 27ac6d845cc..966e7d96efc 100644 --- a/tests/unit/router_utils/test_router_utils_common_utils.py +++ b/tests/unit/router_utils/test_router_utils_common_utils.py @@ -718,7 +718,7 @@ async def test_web_search_request_prefers_the_capable_deployment_inside_its_orde @pytest.mark.asyncio @pytest.mark.parametrize("metadata_field", ["metadata", "litellm_metadata"]) -async def test_web_search_preference_gives_way_once_the_capable_deployment_failed_this_request( +async def test_web_search_retry_goes_to_an_untried_deployment_once_the_capable_deployment_failed_this_request( metadata_field: str, ) -> None: router: Final = _router_with_a_mixed_web_search_group(False) @@ -735,9 +735,45 @@ async def test_web_search_preference_gives_way_once_the_capable_deployment_faile )["model_info"]["id"] for _ in range(40) } - assert picks == {"search", "no-search"} + assert picks == {"no-search"} + +@pytest.mark.parametrize( + ("deployments", "failed", "expected"), + [ + ( + ( + {"id": "search"}, + {"id": "flag-a", "supports_web_search": False}, + {"id": "flag-b", "supports_web_search": False}, + ), + ("search", "flag-a"), + ("flag-b",), + ), + ( + ({"id": "search"}, {"id": "flag-a", "supports_web_search": False}), + ("search", "flag-a"), + ("search", "flag-a"), + ), + ( + ({"id": "flag-a", "supports_web_search": False}, {"id": "flag-b", "supports_web_search": False}), + ("flag-a",), + ("flag-a", "flag-b"), + ), + ], + ids=["skips-every-tried-one", "keeps-all-once-all-were-tried", "no-deployment-can-search"], +) +def test_web_search_retry_prefers_untried_deployments_only_after_one_that_can_search_was_tried( + deployments: tuple[dict[str, object], ...], failed: tuple[str, ...], expected: tuple[str, ...] +) -> None: + request_kwargs: Final = { + **_ANTHROPIC_WEB_SEARCH_REQUEST, + "metadata": {"previous_models": tuple({"deployment_id": d, "exception_type": "Timeout"} for d in failed)}, + } + kept: Final = prefer_web_search_deployments([{"model_info": info} for info in deployments], request_kwargs) + assert tuple(d["model_info"]["id"] for d in kept) == expected + @pytest.mark.asyncio @_WEB_SEARCH_REQUESTS @pytest.mark.parametrize( From 3d59af33df7b9399421d1d0c068d1a21467bd152 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 31 Oct 2026 04:33:24 -0700 Subject: [PATCH 23/24] fix(router): let weights and fallbacks decide once a web search deployment is only at its limit When every deployment that can search was tried or is at its limit and one of them is only at its limit, keep every deployment so weights, rate limit checks and configured fallbacks decide as on main, instead of sending the request to a flagged deployment that answers without a search. A retry after every capable deployment failed still goes to an untried one. Skip the session repin after a web search request in a group that mixes deployments that can and cannot search, matching the claim hook, and read a pydantic ModelInfo's supports_web_search again. --- .../routing_common_utils.py | 10 ++-- litellm/router.py | 12 ++++- .../deployment_affinity_check.py | 5 ++ .../test_deployment_affinity_check.py | 48 +++++++++++++++++++ .../test_session_id_affinity.py | 2 + .../test_router_utils_common_utils.py | 24 ++++++++++ 6 files changed, 96 insertions(+), 5 deletions(-) diff --git a/litellm/litellm_core_utils/routing_common_utils.py b/litellm/litellm_core_utils/routing_common_utils.py index 1565d12a5cd..b46bdcff54a 100644 --- a/litellm/litellm_core_utils/routing_common_utils.py +++ b/litellm/litellm_core_utils/routing_common_utils.py @@ -5,7 +5,7 @@ from itertools import chain from types import MappingProxyType from typing import TYPE_CHECKING, Final, TypeVar -from pydantic import TypeAdapter, ValidationError +from pydantic import BaseModel, TypeAdapter, ValidationError if TYPE_CHECKING: from litellm.types.llms.openai import OpenAIFileObject @@ -440,8 +440,9 @@ def _deployment_supports_web_search(deployment: Mapping[str, object]) -> bool: model_prices_and_context_window.json doesn't have supports_web_search tags on all models yet. TODO: backfill and add fallback. """ - model_info: Final = as_str_mapping(deployment.get("model_info")) or {} - return bool(model_info.get("supports_web_search", True)) + model_info: Final = deployment.get("model_info") + fields: Final = as_str_mapping(model_info.model_dump() if isinstance(model_info, BaseModel) else model_info) or {} + return bool(fields.get("supports_web_search", True)) def _is_web_search_tool(tool: object) -> bool: @@ -579,6 +580,9 @@ def prefer_web_search_deployments( if not untried or not any(_deployment_supports_web_search(d) for d in deployments): verbose_router_logger.debug("No untried deployment for this web search request can search, so all stay") return deployments + if any(_deployment_supports_web_search(d) and _deployment_id(d) in at_rate_limit for d in deployments): + verbose_router_logger.debug("A deployment that can search is only at its limit, so all stay") + return deployments verbose_router_logger.debug("Every deployment that can search was tried, so the untried ones stay") return untried return [ diff --git a/litellm/router.py b/litellm/router.py index de468a8ae2d..1d3adfe1eaf 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -9855,6 +9855,7 @@ class Router(DeprecatedRouterMembers): await self._repin_session_after_success( _LITELLM_PARAMS_ADAPTER.validate_python(litellm_params), _OPAQUE_VALUE_ADAPTER.validate_python(completion_response), + _OPAQUE_VALUE_ADAPTER.validate_python(kwargs.get("optional_params")), ) # Always track deployment successes for cooldown logic, regardless of TPM/RPM limits @@ -9880,7 +9881,9 @@ class Router(DeprecatedRouterMembers): "litellm.router.Router::deployment_callback_on_success(): Exception occured - %s", e ) - async def _repin_session_after_success(self, litellm_params: Mapping[str, object], response: object) -> None: + async def _repin_session_after_success( + self, litellm_params: Mapping[str, object], response: object, optional_params: object + ) -> None: if not self.session_affinity_repin_on_failover: return try: @@ -9906,7 +9909,12 @@ class Router(DeprecatedRouterMembers): model_ids=list(registered_ids), parent_otel_span=None ) ) - await callback.async_repin_session_after_success(litellm_params, model_group, registered_ids - cooldown_ids) + await callback.async_repin_session_after_success( + litellm_params, + model_group, + registered_ids - cooldown_ids, + _LITELLM_PARAMS_ADAPTER.validate_python(optional_params or {}), + ) except Exception as error: # noqa: BLE001 # Optional pin updates must not interrupt success accounting verbose_router_logger.debug("Session re-pinning failed: %s", error) diff --git a/litellm/router_utils/pre_call_checks/deployment_affinity_check.py b/litellm/router_utils/pre_call_checks/deployment_affinity_check.py index 791609e0d52..5b825b47861 100644 --- a/litellm/router_utils/pre_call_checks/deployment_affinity_check.py +++ b/litellm/router_utils/pre_call_checks/deployment_affinity_check.py @@ -597,6 +597,7 @@ class DeploymentAffinityCheck(CustomLogger): request_kwargs: Mapping[str, object], model_group: str, eligible_deployment_ids: frozenset[str], + request_params: Mapping[str, object], ) -> None: marker_ttl: Final = self._get_marker_session_affinity_ttl(request_kwargs) if not self._get_effective_flags(model_group)[2] and marker_ttl is None: @@ -604,6 +605,10 @@ class DeploymentAffinityCheck(CustomLogger): session_id: Final = self._get_session_id_from_request_kwargs(request_kwargs) if session_id is None: return + if self.model_group_deployments is not None and web_search_preference_can_narrow( + self.model_group_deployments(model_group), request_params + ): + return model_info: Final = TypeAdapter(Mapping[str, object]).validate_python(request_kwargs.get("model_info") or {}) model_id: Final = model_info.get("id") if model_id is None or str(model_id) not in eligible_deployment_ids: diff --git a/tests/unit/router_utils/pre_call_checks/test_deployment_affinity_check.py b/tests/unit/router_utils/pre_call_checks/test_deployment_affinity_check.py index 6872c3aa5a2..0106eb30a04 100644 --- a/tests/unit/router_utils/pre_call_checks/test_deployment_affinity_check.py +++ b/tests/unit/router_utils/pre_call_checks/test_deployment_affinity_check.py @@ -1703,6 +1703,7 @@ def _web_search_group_router( optional_pre_call_checks: list[str], # mutable-ok: Router signature heavy_can_search: bool, heavy_weight: int, + session_affinity_repin_on_failover: bool = False, ) -> litellm.Router: return litellm.Router( model_list=[ @@ -1723,6 +1724,7 @@ def _web_search_group_router( }, ], optional_pre_call_checks=optional_pre_call_checks, + session_affinity_repin_on_failover=session_affinity_repin_on_failover, num_retries=0, ) @@ -1766,3 +1768,49 @@ async def test_a_first_web_search_request_still_pins_a_group_where_every_deploym for caller in range(10): searched_on = await _served_by(router, caller, web_search_options={}) assert {await _served_by(router, caller) for _ in range(3)} == {searched_on} + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("optional_params", "session_pin"), + [ + ({"web_search_options": {}}, None), + ({"tools": [{"type": "web_search_20250305", "name": "web_search"}]}, None), + ({}, {"model_id": "light"}), + ], + ids=["web-search-options", "web-search-tool", "plain"], +) +async def test_a_successful_web_search_request_does_not_repin_a_session_in_a_mixed_group( + optional_params: dict[str, object], # mutable-ok: logging kwargs shape + session_pin: dict[str, str] | None, # mutable-ok: cache value shape +): + router: Final = _web_search_group_router( + ["session_affinity"], heavy_can_search=False, heavy_weight=1, session_affinity_repin_on_failover=True + ) + timestamp: Final = datetime(2026, 1, 1, tzinfo=timezone.utc) + + await router.deployment_callback_on_success( + { + "litellm_params": { + "model_info": {"id": "light"}, + "metadata": { + "model_group": "group", + "deployment_model_name": "group", + "session_id": "session-1", + "user_api_key_hash": "caller-1", + }, + }, + "optional_params": optional_params, + "standard_logging_object": {"total_tokens": 0}, + }, + {"choices": [{"message": {"content": "ok"}}]}, + timestamp, + timestamp, + ) + + assert ( + await router.cache.async_get_cache( + DeploymentAffinityCheck.get_session_affinity_cache_key("group", "session-1", "caller-1") + ) + == session_pin + ) diff --git a/tests/unit/router_utils/pre_call_checks/test_session_id_affinity.py b/tests/unit/router_utils/pre_call_checks/test_session_id_affinity.py index 6b4ea74894a..12978a1f34b 100644 --- a/tests/unit/router_utils/pre_call_checks/test_session_id_affinity.py +++ b/tests/unit/router_utils/pre_call_checks/test_session_id_affinity.py @@ -160,6 +160,7 @@ async def test_repin_session_after_success_reads_split_metadata(route_bucket: st session_bucket: {"session_id": "repin-session", "user_api_key_hash": "caller-1"}, }, {"status": "completed"}, + {}, ) assert await router.cache.async_get_cache(_repin_key()) == {"model_id": "replacement"} finally: @@ -349,6 +350,7 @@ async def test_repin_cooldown_reads_are_limited_to_active_session_route(case: st }, }, {"status": "completed"}, + {}, ) expected: Final = ( { diff --git a/tests/unit/router_utils/test_router_utils_common_utils.py b/tests/unit/router_utils/test_router_utils_common_utils.py index 966e7d96efc..c73e354a8e2 100644 --- a/tests/unit/router_utils/test_router_utils_common_utils.py +++ b/tests/unit/router_utils/test_router_utils_common_utils.py @@ -246,6 +246,11 @@ class TestDeploymentSupportsWebSearch: deployment = {"litellm_params": {"model": "gpt-4"}} assert _deployment_supports_web_search(deployment) is True + @pytest.mark.parametrize("supports_web_search", [True, False]) + def test_model_info_object_setting_is_read(self, supports_web_search: bool) -> None: + deployment: Final = {"model_info": ModelInfo(id="1", supports_web_search=supports_web_search)} + assert _deployment_supports_web_search(deployment) is supports_web_search + class TestFilterWebSearchDeployments: """Test cases for filter_web_search_deployments function""" @@ -774,6 +779,25 @@ def test_web_search_retry_prefers_untried_deployments_only_after_one_that_can_se kept: Final = prefer_web_search_deployments([{"model_info": info} for info in deployments], request_kwargs) assert tuple(d["model_info"]["id"] for d in kept) == expected + +@pytest.mark.parametrize( + ("failed", "at_limit", "expected"), + [((), ("search",), ("search", "flag")), (("search",), (), ("flag",))], + ids=["at-its-limit-keeps-all", "failed-keeps-the-untried-one"], +) +def test_a_web_search_deployment_only_at_its_limit_leaves_every_deployment_to_the_weights_and_fallbacks( + failed: tuple[str, ...], at_limit: tuple[str, ...], expected: tuple[str, ...] +) -> None: + request_kwargs: Final = { + **_ANTHROPIC_WEB_SEARCH_REQUEST, + "metadata": {"previous_models": tuple({"deployment_id": d, "exception_type": "Timeout"} for d in failed)}, + } + infos: Final = ({"id": "search"}, {"id": "flag", "supports_web_search": False}) + kept: Final = prefer_web_search_deployments( + [{"model_info": info} for info in infos], request_kwargs, frozenset(at_limit) + ) + assert tuple(d["model_info"]["id"] for d in kept) == expected + @pytest.mark.asyncio @_WEB_SEARCH_REQUESTS @pytest.mark.parametrize( From 790dd4ec94f3690eeb2d9b894c625a6e4010f4d9 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 31 Oct 2026 06:53:14 -0700 Subject: [PATCH 24/24] fix(router): let a web search request move a session pin off a failed deployment With session_affinity_repin_on_failover, a successful web search request in a group that mixes deployments that can search with flagged ones skipped the repin entirely, so a session pinned on a deployment that had just failed kept that pin and its next plain request went back to it. The repin now only skips creating a pin for such a request. An existing pin is refreshed when its deployment is still eligible and moved to the deployment that answered when it is not, as on main. The new request_params and optional_params parameters also get defaults, so existing callers of the repin methods keep working. --- litellm/router.py | 2 +- .../deployment_affinity_check.py | 16 +++++--- .../test_deployment_affinity_check.py | 39 +++++++++++++++++++ .../test_session_id_affinity.py | 2 - 4 files changed, 50 insertions(+), 9 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index 1d3adfe1eaf..bddb8458376 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -9882,7 +9882,7 @@ class Router(DeprecatedRouterMembers): ) async def _repin_session_after_success( - self, litellm_params: Mapping[str, object], response: object, optional_params: object + self, litellm_params: Mapping[str, object], response: object, optional_params: object = None ) -> None: if not self.session_affinity_repin_on_failover: return diff --git a/litellm/router_utils/pre_call_checks/deployment_affinity_check.py b/litellm/router_utils/pre_call_checks/deployment_affinity_check.py index 5b825b47861..dfb2fe7e5c6 100644 --- a/litellm/router_utils/pre_call_checks/deployment_affinity_check.py +++ b/litellm/router_utils/pre_call_checks/deployment_affinity_check.py @@ -14,6 +14,7 @@ where routing to a consistent deployment is still beneficial. import hashlib from collections.abc import Awaitable, Callable, Mapping, Sequence +from types import MappingProxyType from typing import Final, TypeVar, cast from pydantic import TypeAdapter @@ -25,6 +26,7 @@ from litellm.caching.affinity_cache import ( ROUTER_SESSION_PINS_TARGET, claim_affinity_pin, claim_affinity_pin_in_memory, + read_affinity_pin, set_local_affinity_pin, ) from litellm.caching.dual_cache import DualCache @@ -597,7 +599,7 @@ class DeploymentAffinityCheck(CustomLogger): request_kwargs: Mapping[str, object], model_group: str, eligible_deployment_ids: frozenset[str], - request_params: Mapping[str, object], + request_params: Mapping[str, object] = MappingProxyType({}), ) -> None: marker_ttl: Final = self._get_marker_session_affinity_ttl(request_kwargs) if not self._get_effective_flags(model_group)[2] and marker_ttl is None: @@ -605,10 +607,6 @@ class DeploymentAffinityCheck(CustomLogger): session_id: Final = self._get_session_id_from_request_kwargs(request_kwargs) if session_id is None: return - if self.model_group_deployments is not None and web_search_preference_can_narrow( - self.model_group_deployments(model_group), request_params - ): - return model_info: Final = TypeAdapter(Mapping[str, object]).validate_python(request_kwargs.get("model_info") or {}) model_id: Final = model_info.get("id") if model_id is None or str(model_id) not in eligible_deployment_ids: @@ -616,11 +614,17 @@ class DeploymentAffinityCheck(CustomLogger): cache_key: Final = self.get_session_affinity_cache_key( model_group, session_id, self.get_user_key_from_request_kwargs(request_kwargs) ) + ttl_seconds: Final = marker_ttl if marker_ttl is not None else self.ttl_seconds + is_search_in_mixed_group: Final = self.model_group_deployments is not None and web_search_preference_can_narrow( + self.model_group_deployments(model_group), request_params + ) + if is_search_in_mixed_group and await read_affinity_pin(self.cache, cache_key, ttl_seconds) is None: + return await claim_affinity_pin( self.cache, cache_key, DeploymentAffinityCacheValue(model_id=str(model_id)), - marker_ttl if marker_ttl is not None else self.ttl_seconds, + ttl_seconds, eligible_values=tuple(DeploymentAffinityCacheValue(model_id=value) for value in eligible_deployment_ids), preserve_legacy_model_ids=True, ) diff --git a/tests/unit/router_utils/pre_call_checks/test_deployment_affinity_check.py b/tests/unit/router_utils/pre_call_checks/test_deployment_affinity_check.py index 0106eb30a04..6b3a651361a 100644 --- a/tests/unit/router_utils/pre_call_checks/test_deployment_affinity_check.py +++ b/tests/unit/router_utils/pre_call_checks/test_deployment_affinity_check.py @@ -1814,3 +1814,42 @@ async def test_a_successful_web_search_request_does_not_repin_a_session_in_a_mix ) == session_pin ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("pinned_model_id", "session_pin"), + [("removed", {"model_id": "light"}), ("heavy", {"model_id": "heavy"})], + ids=["ineligible-pin-moves", "eligible-pin-stays"], +) +async def test_a_successful_web_search_request_moves_a_session_only_off_an_ineligible_pin_in_a_mixed_group( + pinned_model_id: str, + session_pin: dict[str, str], # mutable-ok: cache value shape +): + router: Final = _web_search_group_router( + ["session_affinity"], heavy_can_search=False, heavy_weight=1, session_affinity_repin_on_failover=True + ) + cache_key: Final = DeploymentAffinityCheck.get_session_affinity_cache_key("group", "session-1", "caller-1") + await router.cache.async_set_cache(cache_key, {"model_id": pinned_model_id}) + timestamp: Final = datetime(2026, 1, 1, tzinfo=timezone.utc) + + await router.deployment_callback_on_success( + { + "litellm_params": { + "model_info": {"id": "light"}, + "metadata": { + "model_group": "group", + "deployment_model_name": "group", + "session_id": "session-1", + "user_api_key_hash": "caller-1", + }, + }, + "optional_params": {"web_search_options": {}}, + "standard_logging_object": {"total_tokens": 0}, + }, + {"choices": [{"message": {"content": "ok"}}]}, + timestamp, + timestamp, + ) + + assert await router.cache.async_get_cache(cache_key) == session_pin diff --git a/tests/unit/router_utils/pre_call_checks/test_session_id_affinity.py b/tests/unit/router_utils/pre_call_checks/test_session_id_affinity.py index 12978a1f34b..6b4ea74894a 100644 --- a/tests/unit/router_utils/pre_call_checks/test_session_id_affinity.py +++ b/tests/unit/router_utils/pre_call_checks/test_session_id_affinity.py @@ -160,7 +160,6 @@ async def test_repin_session_after_success_reads_split_metadata(route_bucket: st session_bucket: {"session_id": "repin-session", "user_api_key_hash": "caller-1"}, }, {"status": "completed"}, - {}, ) assert await router.cache.async_get_cache(_repin_key()) == {"model_id": "replacement"} finally: @@ -350,7 +349,6 @@ async def test_repin_cooldown_reads_are_limited_to_active_session_route(case: st }, }, {"status": "completed"}, - {}, ) expected: Final = ( {